import gradio as gr
import numpy as np
import random
import spaces
import torch
from diffusers import DiffusionPipeline

# Load the model
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"
pipe = DiffusionPipeline.from_pretrained("black-forest-labs/FLUX.1-schnell", torch_dtype=dtype).to(device)

# Constants
MAX_SEED = np.iinfo(np.int32).max
MAX_IMAGE_SIZE = 2048

# Style list for prompt customization
style_list = [
    {"name": "D&D Art", "prompt": "dungeons & dragons style artwork {prompt}. d&d style, key visual, vibrant, studio anime, highly detailed", "negative_prompt": "photo, deformed, black and white, realism, disfigured, low contrast"},
    {"name": "Dark Fantasy", "prompt": "dark and moody dungeons & dragons artwork of {prompt}. gothic ruins, shadowy figures, haunting atmospheres, grim villains, muted colors, intricate textures, sinister undertones", "negative_prompt": "bright, cheerful, cartoonish, lighthearted, futuristic, deformed"},
    {"name": "Epic Battle", "prompt": "dynamic dungeons & dragons artwork of {prompt}. epic battle scene, legendary heroes, fierce monsters, intense action, dramatic lighting, high-detail environment, magical effects, vibrant colors", "negative_prompt": "peaceful, mundane, low energy, modern, sci-fi, simplistic, cartoonish, low contrast"},
    # Add additional styles as needed
{    "name": "Character Portrait",
    "prompt": "dungeons & dragons character portrait of {prompt}. intricate details, expressive faces, heroic poses, rich textures, elaborate costumes, iconic weapons, fantasy aesthetic, studio-quality rendering",
    "negative_prompt": "generic, overly stylized, blurry, simplistic, futuristic, deformed"
},
{
    "name": "Magical Realm",
    "prompt": "mystical dungeons & dragons artwork of {prompt}. enchanted forests, glowing runes, floating islands, otherworldly landscapes, magical auras, vibrant fantasy colors, ultra-detailed",
    "negative_prompt": "urban, dull, realistic, futuristic, lifeless, overly muted, minimalistic"
},
{
    "name": "Legendary Creatures",
    "prompt": "dungeons & dragons artwork of legendary creatures {prompt}. awe-inspiring dragons, ferocious beasts, mythical monsters, intricate scales, detailed textures, dynamic poses, magical auras, cinematic quality",
    "negative_prompt": "simplistic, dull, mundane creatures, low-detail, modern animals, cartoony"
},

{
    "name": "Ancient Ruins",
    "prompt": "dungeons & dragons artwork of {prompt}. ancient temples, crumbling ruins, arcane symbols, epic scale, mystical atmosphere, grand designs, dramatic perspective, ultra-detailed",
    "negative_prompt": "modern structures, low-detail, dull, untextured, non-fantasy elements, minimalistic"
},

{
    "name": "Spellcasting Scene",
    "prompt": "dungeons & dragons artwork of {prompt}. dramatic spellcasting, glowing magic effects, dynamic poses, swirling energy, vibrant light contrasts, powerful sorcerers, rich fantasy aesthetic",
    "negative_prompt": "low energy, non-magical, mundane, modern day, generic designs, dull colors"
},

{
    "name": "Tavern Life",
    "prompt": "dungeons & dragons artwork of {prompt}. lively medieval tavern scene, detailed character interactions, warm firelight, wooden interiors, mugs of ale, bard performances, rich textures, inviting atmosphere",
    "negative_prompt": "cold, empty, futuristic, lifeless, overly simplistic, boring"
},

{
    "name": "D&D Map",
    "prompt": "dungeons & dragons map artwork of {prompt}. In the style of Dungeon Tiles, Intricate fantasy map with a precise 1-inch visible grid layout adhering to D&D rules (5 feet per square), ancient cartography style, clearly labeled locations, key landmarks, navigational markers, topographic textures, detailed compass rose, parchment-like texture, vibrant rich colors, immersive design, optimized for tabletop gameplay and tactical use.",
    "negative_prompt": "modern maps, bland, plain, low detail, inaccurate grids (non-5-foot alignment), futuristic, chaotic, unlabeled locations, minimalist, non-immersive, disproportionate landmarks."
},

{
    "name": "Dungeon Map",
    "prompt": "Simple dungeon grid map for D&D in the style of {prompt}. Dungeon Map Builder. An intricate fantasy map with a precise 1-inch visible grid layout (5 feet per square) following D&D rules, featuring ancient cartography style, labeled dungeon rooms, detailed topographic textures, clear navigational markers, key points of interest (like treasure rooms, traps, and secret passages), vibrant colors for clarity, parchment-like background texture, and optimized for tabletop and tactical gameplay. Designed for immersive and strategic D&D adventures.",
    "negative_prompt": "Modern digital maps, bland designs, low-detail features, inaccurate grids (non-5-foot alignment), futuristic styles, chaotic layouts, unlabeled locations, minimalist aesthetics, disproportionate scales, or non-immersive elements."
},
{
    "name": "Black & White Dungeon Map",
    "prompt": "Detailed top-down black and white dungeon grid map for {prompt}. In the style of a D&D Dungeon Map Builder. An intricate fantasy map with a precise 1-inch visible grid layout (5 feet per square) following D&D rules, featuring ancient cartography style, labeled dungeon rooms with descriptions like: winding corridors, room with barrels and chest, room with beds, room with sarcophagus, room with table and chair, room with bed and figure, room with stairs and plants, room with windows. Detailed topographic textures, ink drawing, sketch, hand-drawn lines, cave-like walls for corridors.",
    "negative_prompt": "Modern digital maps, color maps, bland designs, low-detail features, inaccurate grids (non-5-foot alignment), futuristic styles, chaotic layouts, unlabeled locations, minimalist aesthetics, disproportionate scales, photorealistic elements, 3D rendering, or non-immersive elements."
},
{
    "name": "Colorful Detailed Dungeon Map",
    "prompt": "Detailed top-down dungeon grid map for D&D in the style of {prompt}. Dungeon Map Builder. An intricate fantasy map with a precise 1-inch visible grid layout (5 feet per square) following D&D rules, featuring ancient cartography style, labeled dungeon rooms with descriptions like: winding corridors, room with barrels and chest, room with beds, room with sarcophagus, room with table and chair, room with bed and figure, room with stairs and plants, room with windows. Detailed topographic textures, clear navigational markers, key points of interest, vibrant colors for clarity, parchment-like background texture, cave-like walls for corridors.",
    "negative_prompt": "Modern digital maps, bland designs, low-detail features, inaccurate grids (non-5-foot alignment), futuristic styles, chaotic layouts, unlabeled locations, minimalist aesthetics, disproportionate scales, photorealistic elements, 3D rendering, or non-immersive elements, black and white maps."
},

{
    "name": "D&D Quest Art",
    "prompt": "dungeons & dragons quest artwork of {prompt}. heroic adventurers, perilous journeys, hidden treasures, legendary quests, immersive fantasy landscapes, intricate details, atmospheric lighting, grand scale, narrative-driven composition",
    "negative_prompt": "modern, futuristic, low-detail, overly abstract, mundane settings, simplistic, sci-fi, cartoony, lifeless, uninspired"
},

{
    "name": "D&D Anime Art",
    "prompt": "dungeons & dragons anime-style artwork of {prompt}. stylized characters, vibrant colors, expressive designs, dynamic action poses, magical effects, intricate fantasy backgrounds, studio-quality anime detailing",
    "negative_prompt": "photo-realistic, overly gritty, mundane, futuristic, overly dark, simplistic, deformed"
},

{
    "name": "Arcane Duel",
    "prompt": "dungeons & dragons artwork of {prompt}. two powerful sorcerers battling with swirling magical energy, dramatic lighting, glowing spell effects, mystical atmosphere, ultra-detailed fantasy elements",
    "negative_prompt": "calm, peaceful, dull colors, modern setting, non-fantasy, low-detail"
},
{
    "name": "Treasure Hoard",
    "prompt": "dungeons & dragons artwork of {prompt}. a cavern filled with glittering gold coins, ancient relics, glowing gemstones, treasure chests, detailed textures, vibrant fantasy lighting",
    "negative_prompt": "empty, dull, modern day, futuristic, sci-fi, bland atmosphere"
},
{
    "name": "Underdark",
    "prompt": "dark and moody dungeons & dragons artwork of {prompt}. sprawling underground caverns, glowing fungi, shadowy creatures, bioluminescent details, ominous tones, intricate fantasy detailing",
    "negative_prompt": "bright and sunny, surface world, futuristic, sci-fi, minimal detail, lifeless"
},
{
    "name": "Celestial Realm",
    "prompt": "dungeons & dragons artwork of {prompt}. divine celestial realms, golden light, angelic beings, floating clouds, radiant auras, epic divine architecture, detailed textures",
    "negative_prompt": "dark, demonic, mundane, low-detail, futuristic, simplistic"
},
{
    "name": "D&D Heist Scene",
    "prompt": "dungeons & dragons artwork of {prompt}. a group of adventurers sneaking into a heavily guarded treasure vault, dramatic lighting, intricate lock mechanisms, tense atmosphere, ultra-detailed environment",
    "negative_prompt": "carefree, bland setting, modern day, low-detail, simplistic designs, deformed"
},
{
    "name": "Enchanted Forest",
    "prompt": "dungeons & dragons artwork of {prompt}. magical forest with glowing flora, towering ancient trees, mystical wildlife, detailed textures, serene but vibrant fantasy tones",
    "negative_prompt": "modern city, mundane, lifeless, futuristic, abstract designs, overly simplistic"
},
{
    "name": "Battle at the Gates",
    "prompt": "dungeons & dragons artwork of {prompt}. an epic battle unfolding at massive fortress gates, heavily armored soldiers, dramatic clashes, powerful war machines, intense action, atmospheric lighting",
    "negative_prompt": "peaceful, bland, empty fields, low energy, futuristic, deformed, overly simplistic"
},
{
    "name": "Stormy Seas",
    "prompt": "dungeons & dragons artwork of {prompt}. a ship caught in a raging storm, enormous waves, dark clouds, flashes of lightning, detailed ocean textures, dynamic lighting",
    "negative_prompt": "calm water, lifeless, plain scenery, futuristic, minimalist, abstract"
},
{
    "name": "Desert Wasteland",
    "prompt": "dungeons & dragons artwork of {prompt}. endless desert landscapes, ruins of ancient civilizations, glowing sun on the horizon, swirling sands, intricate details, mystical atmosphere",
    "negative_prompt": "lush forests, futuristic, cityscapes, low-detail, lifeless"
},
{
    "name": "Dwarven Forge",
    "prompt": "dungeons & dragons artwork of {prompt}. a massive underground forge with molten lava streams, glowing furnaces, dwarf blacksmiths, intricate metalworks, rugged and vibrant tones",
    "negative_prompt": "modern machinery, minimalistic, clean, futuristic, deformed"
},
{
    "name": "Ethereal Spires",
    "prompt": "dungeons & dragons artwork of {prompt}. floating crystalline towers above the clouds, glowing lights, otherworldly designs, vibrant celestial colors, ultra-detailed mystical atmosphere",
    "negative_prompt": "mundane, low-detail, dark and grim, simplistic, lifeless"
},
{
    "name": "Feywild Meadow",
    "prompt": "dungeons & dragons artwork of {prompt}. a surreal and magical meadow filled with sparkling lights, dancing fey creatures, vibrant blooming flowers, whimsical details, fantasy colors",
    "negative_prompt": "plain forest, dark and brooding, lifeless, realistic, dull tones"
},
{
    "name": "Dungeon Crawl",
    "prompt": "dungeons & dragons artwork of {prompt}. adventurers exploring a dark dungeon with traps, treasure chests, cobwebbed corners, torch-lit corridors, richly detailed stone textures",
    "negative_prompt": "bright outdoor scenes, futuristic, modern, low detail"
},
{
    "name": "Frosted Peaks",
    "prompt": "dungeons & dragons artwork of {prompt}. icy mountaintops, adventurers climbing through harsh snowstorms, frozen castles, glowing icy runes, cold tones with vivid details",
    "negative_prompt": "warm beaches, futuristic, cartoony, deformed, dull colors"
},
{
    "name": "Cyberpunk Cityscape",
    "prompt": "futuristic cyberpunk artwork of {prompt}. neon-lit streets, towering skyscrapers, flying vehicles, cybernetic details, vibrant colors, cinematic atmosphere, ultra-detailed",
    "negative_prompt": "fantasy, medieval, dull colors, empty city, overly simplistic, bland"
},
{
    "name": "Steampunk World",
    "prompt": "steampunk-style artwork of {prompt}. intricate clockwork mechanisms, steam-powered inventions, Victorian-inspired characters, vibrant but rustic tones, rich textures, detailed environments",
    "negative_prompt": "futuristic, sleek designs, dull, modern cityscapes, simplistic, minimalistic"
},
{
    "name": "Wild West Scene",
    "prompt": "cinematic wild west-style artwork of {prompt}. vast desert landscapes, rugged cowboys, wooden saloons, vintage rifles, atmospheric lighting, detailed textures",
    "negative_prompt": "modern vehicles, futuristic elements, overly vibrant, simplistic, low energy"
},
{
    "name": "Mythological Greece",
    "prompt": "mythological artwork of {prompt}. ancient Greek gods and heroes, marble temples, epic battles, radiant divine light, ultra-detailed settings, rich classical tones",
    "negative_prompt": "sci-fi, modern city, low-detail, lifeless atmosphere, generic textures"
},
{
    "name": "Sci-Fi Space Adventure",
    "prompt": "sci-fi style artwork of {prompt}. massive starships, alien planets, swirling galaxies, futuristic cities, neon highlights, ultra-detailed space themes",
    "negative_prompt": "fantasy elements, medieval setting, dull or muted tones, overly simple designs"
},
{
    "name": "Retro Pixel Art",
    "prompt": "8-bit retro-style pixel art of {prompt}. classic video game aesthetics, vibrant and bold pixel colors, stylized backgrounds, clean and defined sprites",
    "negative_prompt": "realistic textures, photorealistic, overly smooth, detailed textures"
},
{
    "name": "Pop Art Aesthetic",
    "prompt": "vivid pop art illustration of {prompt}. bold colors, high contrast, retro style, comic book-like details, playful and energetic compositions",
    "negative_prompt": "realism, dull tones, blurred effects, abstract"
},
{
    "name": "Noir Detective",
    "prompt": "cinematic noir-style artwork of {prompt}. black and white aesthetics, shadowy figures, fog-filled streets, dramatic lighting, vintage details, atmospheric tension",
    "negative_prompt": "colorful, futuristic, mundane, cheerful, low-detail"
},
{
    "name": "Studio Ghibli-Inspired Fantasy",
    "prompt": "anime-style artwork inspired by Studio Ghibli of {prompt}. magical landscapes, intricate details, vibrant colors, fantastical creatures, peaceful yet surreal atmosphere",
    "negative_prompt": "dark and brooding, sci-fi, plain, minimalist, harsh contrasts"
},
{
    "name": "Art Deco City",
    "prompt": "art deco-inspired cityscape of {prompt}. geometric designs, golden accents, intricate and modern architecture, glowing lights, detailed and luxurious atmosphere",
    "negative_prompt": "medieval settings, fantasy, low-detail, dull textures, minimalism"
},
    {"name": "(No style)", "prompt": "{prompt}", "negative_prompt": ""},
]

styles = {k["name"]: (k["prompt"], k["negative_prompt"]) for k in style_list}
STYLE_NAMES = list(styles.keys())
DEFAULT_STYLE_NAME = "Black & White Dungeon Map"

# Function to apply selected style
def apply_style(style_name: str, positive: str, negative: str = ""):
    p, n = styles.get(style_name, styles[DEFAULT_STYLE_NAME])
    return p.replace("{prompt}", positive), n + (negative or "")

# Inference function
@spaces.GPU()
def infer(
    prompt,
    style,
    seed=42,
    randomize_seed=False,
    width=1024,
    height=1024,
    num_inference_steps=6,
    progress=gr.Progress(track_tqdm=True),
):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    # Apply style to prompt
    styled_prompt, negative_prompt = apply_style(style, prompt)
    generator = torch.Generator().manual_seed(seed)
    image = pipe(
        prompt=styled_prompt,
        width=width,
        height=height,
        num_inference_steps=num_inference_steps,
        generator=generator,
        guidance_scale=0.0,
        negative_prompt=negative_prompt,
    ).images[0]
    return image, seed

# Example prompts
examples = [
    ["A 10x8 battle map located in Ravenloft", "D&D Map"],
    ["A heroic adventurer wielding a flaming sword standing on a cliff", "D&D Art"],
    ["A mystical library with ancient scrolls and glowing runes", "Dark Fantasy"],
    ["A ferocious dragon breathing fire in a dark cavern", "Epic Battle"],
]

css = """
/* General body styling */
body {
    background-color: #1b1b1b;
    font-family: 'Cinzel', serif;
    color: #f5f5f5;
    background-image: url('https://www.transparenttextures.com/patterns/dark-matter.png');
    margin: 0;
    padding: 0;
}
/* Container */
#col-container {
    margin: 0 auto;
    max-width: 550px;
    padding: 20px;
    border: 6px solid #8b4513;
    background: radial-gradient(circle, #302d2b 0%, #3b3634 60%, #1b1b1b 100%);
    border-radius: 20px;
    box-shadow: 0 0 30px rgba(0, 0, 0, 0.9), 0 0 15px rgba(139, 69, 19, 0.7);
    transition: transform 0.3s ease-in-out, box-shadow 0.3s ease-in-out;
}
#col-container:hover {
    transform: scale(1.03);
    box-shadow: 0 0 40px rgba(255, 223, 0, 0.8), 0 0 20px rgba(255, 140, 0, 0.8);
}
/* Fancy title and text effects */
#col-container h1, #col-container h2, #col-container h3 {
    color: #ffebcc;
    text-shadow: 0 0 10px #8b4513, 0 0 20px #ffd700;
}
#col-container p {
    color: #dcdcdc;
    text-shadow: 0 0 6px rgba(139, 69, 19, 0.8);
}
/* Markdown links */
#col-container a {
    color: #e6ac00;
    text-decoration: none;
    font-weight: bold;
    transition: color 0.3s ease-in-out, text-shadow 0.3s ease-in-out;
}
#col-container a:hover {
    color: #fff700;
    text-shadow: 0 0 10px #fff700, 0 0 20px #ffebcc;
}
/* Buttons with D&D glowing effect */
.gr-button {
    background-color: #8b4513;
    color: #f5f5f5;
    font-weight: bold;
    padding: 10px 20px;
    border: 2px solid #b8860b;
    border-radius: 8px;
    cursor: pointer;
    transition: background-color 0.4s ease-in-out, box-shadow 0.4s ease-in-out, transform 0.2s ease-in-out;
}
.gr-button:hover {
    background-color: #b8860b;
    transform: scale(1.05);
    box-shadow: 0 0 12px rgba(255, 223, 0, 0.9), 0 0 24px rgba(255, 140, 0, 0.7);
}
/* Input and Dropdown styling */
input, select, .gr-input, .gr-slider {
    background-color: #302d2b;
    color: #b59e54;
    border: 2px solid #8b4513;
    border-radius: 5px;
    padding: 8px 12px;
    font-family: 'Cinzel', serif;
    transition: border-color 0.3s ease, box-shadow 0.3s ease;
}
input:focus, select:focus, .gr-input:focus, .gr-slider:focus {
    outline: none;
    border-color: #ffd700;
    box-shadow: 0 0 8px #ffd700;
}
/* Slider styling */
.gr-slider > div {
    background-color: #8b4513 !important;
}
/* Checkbox customization */
.gr-checkbox-label {
    font-weight: bold;
    color: #f5f5f5;
    text-shadow: 0 0 5px rgba(139, 69, 19, 0.7);
}
/* Accordion effects */
.gr-accordion {
    background: linear-gradient(145deg, #302d2b, #3b3634);
    color: #f5f5f5;
    border: 2px solid #8b4513;
    border-radius: 10px;
    padding: 8px;
    transition: all 0.3s ease;
}
.gr-accordion:hover {
    box-shadow: 0 0 12px rgba(255, 223, 0, 0.7), 0 0 20px rgba(255, 140, 0, 0.5);
}
/* Image glow effect */
.gr-image {
    border: 4px solid #8b4513;
    border-radius: 10px;
    box-shadow: 0 0 10px #8b4513, 0 0 20px #ffd700;
}
/* Add glow on hover for images */
.gr-image:hover {
    box-shadow: 0 0 20px rgba(255, 140, 0, 0.7), 0 0 40px rgba(255, 223, 0, 0.9);
}
"""

# Interface
with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col-container"):
        # Title and Description
        gr.Markdown(
            """
            # 🛡️ ChatDnD.net Dungeon Generator ⚔️
            **Unleash Your Imagination!** Create heroes, maps, quests, and epic scenes to bring your campaigns to life. 
            Tailored for adventurers seeking inspiration or Dungeon Masters constructing their next grand story. <br>
            [Visit Our Website](https://chatdnd.net) | [Support Us](https://buymeacoffee.com/watchoutformike)
            """
        )

        # Prompt input and style selector
        with gr.Row():
            prompt = gr.Textbox(
                label="🎲 Describe Your Vision:",
                lines=3,
                placeholder="Describe your hero, monster, or legendary landscape..."
            )
            style = gr.Dropdown(
                label="🎨 Select a Style",
                choices=STYLE_NAMES,
                value=DEFAULT_STYLE_NAME,
            )
        
        # Run button and result display
        with gr.Row():
            run_button = gr.Button("Generate Image")
        result = gr.Image(label="🖼️ Your Legendary Vision")
        
        # Advanced settings
        with gr.Accordion("⚙️ Advanced Settings", open=False):
            seed = gr.Slider(
                label="Seed",
                minimum=0,
                maximum=MAX_SEED,
                step=1,
                value=0,
            )
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            
            with gr.Row():
                width = gr.Slider(
                    label="Width",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
                height = gr.Slider(
                    label="Height",
                    minimum=256,
                    maximum=MAX_IMAGE_SIZE,
                    step=32,
                    value=1024,
                )
            
            num_inference_steps = gr.Slider(
                label="Inference Steps",
                minimum=1,
                maximum=50,
                step=1,
                value=6,
            )

        # Examples with styles
        gr.Examples(
            examples=examples,
            inputs=[prompt, style],
            outputs=[result],
            fn=infer,
            cache_examples="lazy",
        )

    # Interactivity
    gr.on(
        triggers=[run_button.click, prompt.submit],
        fn=infer,
        inputs=[prompt, style, seed, randomize_seed, width, height, num_inference_steps],
        outputs=[result, seed],
    )

# Launch the demo
demo.launch()