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

# Basic settings
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

# Load model (FLUX.1-schnell)
pipe = DiffusionPipeline.from_pretrained(
    "black-forest-labs/FLUX.1-schnell",
    torch_dtype=dtype
).to(device)

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

@spaces.GPU()
def generate_image(prompt, seed, randomize_seed, width, height, steps, guidance_scale):
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device).manual_seed(seed)
    image = pipe(
        prompt=prompt,
        width=width,
        height=height,
        num_inference_steps=steps,
        generator=generator,
        guidance_scale=guidance_scale
    ).images[0]
    return image, seed

def set_prompt(example_text):
    return example_text

# Updated example prompts per tab with richer details:
example_prompts = {
    "Product Design": [
        """A sleek industrial design concept for a coffee machine: 
- Curved metallic body with minimal bezel
- Touchscreen panel for settings
- Modern matte black finish
- Hand-drawn concept sketch style"""
    ],
    "Mindmap": [
        """A handrawn colorful mind map diagram, educational style, vibrant colors, clear hierarchy, golden ratio layout.
KNOWLEDGE
├── ACQUISITION [Brain with Lightning ~60px]
│   ├── READING [Open Book with Glow]
│   ├── PRACTICE [Hands-on Tools]
│   └── OBSERVATION [Eye with Magnifier]
├── PROCESSING [Gear Network ~50px]
│   ├── ANALYSIS [Graph Trending Up]
│   └── SYNTHESIS [Puzzle Pieces]
├── RETENTION [Memory Chip ~45px]
│   ├── SHORT-TERM [Quick Flash]
│   └── LONG-TERM [Solid Archive]
└── APPLICATION
    ├── CREATION [Artist Palette]
    └── INNOVATION [Lightbulb Constellation]"""
    ],
    "Mockup": [
        """A clean hand-drawn style wireframe for a mobile banking app:
- Title screen with logo
- Login screen (username, password, login button)
- Dashboard with 3 main sections (balance, transactions, quick actions)
- Bottom navigation bar (home, transfers, profile)"""
    ],
    "Infographic": [
        """A sophisticated flat-style infographic for a multinational corporation’s annual report:
- Title: "Global Renewable Energy Trends 2025"
- Subtitle: "Market Share and Growth Analysis"
- Visual Elements:
    - Multi-segmented bar charts comparing Solar, Wind, and Hydro energy production across regions
    - Pie chart displaying overall energy distribution: Solar (45%), Wind (30%), Hydro (25%)
    - Trend lines indicating year-over-year growth
- Icons: Sleek, minimalist representations of a sun, wind turbine, and water droplet
- Layout: Clean, grid-based design with ample white space and pastel accents for a modern corporate look
- Annotations: Brief, impactful data callouts highlighting key performance indicators and future forecasts"""
    ],
    "Diagram": [
        """A detailed hand-drawn diagram illustrating an end-to-end business workflow:
- Title: "Integrated Business Process Diagram"
- Components:
    - Market Analysis: Research phase with data charts and competitor mapping
    - Strategy Development: Ideation stage with brainstorming clouds and key focus areas
    - Product Design: Concept sketches with iterative feedback loops
    - Implementation: Process flow with timeline markers and resource allocation icons
    - Post-Launch Review: Feedback, metrics, and continuous improvement cycles
- Visual Elements:
    - Clear, directional arrows connecting each phase
    - Iconography for each component (e.g., magnifying glass, lightbulb, gear, checklist)
- Style: Vibrant, educational yet professional, balancing detailed annotations with visual simplicity
- Layout: Structured with a clear hierarchy and color-coded sections to differentiate process stages"""
    ],
    "Flowchart": [
        """A hand-drawn style flowchart, vibrant colors, minimalistic icons.
BUSINESS WORKFLOW
├── START [Green Button ~40px]
│   ├── COLLECT REQUIREMENTS [Folder Icon]
│   └── ANALYZE DATA [Chart Icon]
├── IMPLEMENTATION [Coding Symbol ~50px]
│   ├── FRONTEND [Browser Icon]
│   └── BACKEND [Server Icon]
├── TEST & INTEGRATION [Gear Icon ~45px]
└── DEPLOY
    └── END [Checkered Flag ~40px]"""
    ]
}

# CSS style: fixed container width and preventing examples from expanding the layout
css = """
* {
    box-sizing: border-box;
}
body {
    background: linear-gradient(135deg, #667eea, #764ba2);
    font-family: 'Helvetica Neue', Helvetica, Arial, sans-serif;
    color: #333;
    margin: 0;
    padding: 0;
}
.gradio-container {
    background: rgba(255, 255, 255, 0.95);
    border-radius: 15px;
    padding: 30px 40px;
    box-shadow: 0 8px 30px rgba(0, 0, 0, 0.3);
    margin: 40px auto;
    width: 1200px;
}
.sidebar {
    background: rgba(255, 255, 255, 0.98);
    border-radius: 10px;
    padding: 20px;
    box-shadow: 0 4px 15px rgba(0, 0, 0, 0.2);
}
button, .btn {
    background: linear-gradient(90deg, #ff8a00, #e52e71);
    border: none;
    color: #fff;
    padding: 12px 24px;
    text-transform: uppercase;
    font-weight: bold;
    letter-spacing: 1px;
    border-radius: 5px;
    cursor: pointer;
    transition: transform 0.2s ease-in-out;
}
button:hover, .btn:hover {
    transform: scale(1.05);
}
.example-accordion {
    width: 100% !important;
    max-width: 100% !important;
}
.example-accordion button {
    width: auto !important;
    white-space: normal !important;
}
"""

with gr.Blocks(css=css, title="Workflow Canvas") as demo:
    gr.Markdown(
        """
        <div style="text-align:center;">
            <h1>Workflow Canvas</h1>
            <p>Generate design concepts and workflow diagrams for your business needs using our multi-tab interface.</p>
            <p><strong>AI Spaces:</strong> Product Design, Mindmap, Mockup, Infographic, Diagram, Flowchart</p>
        </div>
        """
    )
    
    with gr.Row():
        # Left sidebar: Common generation parameters
        with gr.Sidebar(label="Parameters", open=True):
            gr.Markdown("### Generation Parameters")
            seed_slider = gr.Slider(label="Seed", minimum=0, maximum=MAX_SEED, step=1, value=42)
            randomize_seed = gr.Checkbox(label="Randomize Seed", value=True)
            width_slider = gr.Slider(label="Width", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            height_slider = gr.Slider(label="Height", minimum=256, maximum=MAX_IMAGE_SIZE, step=32, value=1024)
            steps_slider = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=20)
            guidance_slider = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=15.0, step=0.5, value=7.5)
        
        # Main area: Tabbed interface with updated tab order:
        # 1: Product Design, 2: Mindmap, 3: Mockup, 4: Infographic, 5: Diagram, 6: Flowchart
        with gr.Column(scale=8):
            with gr.Tabs():
                # Tab 1: Product Design
                with gr.Tab("Product Design"):
                    pd_prompt = gr.Textbox(
                        label="Product Design Prompt",
                        placeholder="Enter a product design concept...",
                        lines=5,
                        value=example_prompts["Product Design"][0]
                    )
                    pd_generate = gr.Button("Generate Product Design")
                    pd_image = gr.Image(label="Generated Product Design", value="w1.webp")
                    with gr.Accordion("Example Prompts", open=True, elem_classes="example-accordion"):
                        for ex in example_prompts["Product Design"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=pd_prompt)
                
                # Tab 2: Mindmap
                with gr.Tab("Mindmap"):
                    mindmap_prompt = gr.Textbox(
                        label="Mindmap Prompt",
                        placeholder="Enter a mind map description...",
                        lines=5,
                        value=example_prompts["Mindmap"][0]
                    )
                    mindmap_generate = gr.Button("Generate Mindmap")
                    mindmap_image = gr.Image(label="Generated Mindmap", value="w6.webp")
                    with gr.Accordion("Example Prompts", open=True, elem_classes="example-accordion"):
                        for ex in example_prompts["Mindmap"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=mindmap_prompt)
                
                # Tab 3: Mockup
                with gr.Tab("Mockup"):
                    mock_prompt = gr.Textbox(
                        label="Mockup Prompt",
                        placeholder="Enter a mockup description...",
                        lines=5,
                        value=example_prompts["Mockup"][0]
                    )
                    mock_generate = gr.Button("Generate Mockup")
                    mock_image = gr.Image(label="Generated Mockup", value="w2.webp")
                    with gr.Accordion("Example Prompts", open=True, elem_classes="example-accordion"):
                        for ex in example_prompts["Mockup"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=mock_prompt)
                
                # Tab 4: Infographic
                with gr.Tab("Infographic"):
                    info_prompt = gr.Textbox(
                        label="Infographic Prompt",
                        placeholder="Enter an infographic description...",
                        lines=5,
                        value=example_prompts["Infographic"][0]
                    )
                    info_generate = gr.Button("Generate Infographic")
                    info_image = gr.Image(label="Generated Infographic", value="w3.webp")
                    with gr.Accordion("Example Prompts", open=True, elem_classes="example-accordion"):
                        for ex in example_prompts["Infographic"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=info_prompt)
                
                # Tab 5: Diagram
                with gr.Tab("Diagram"):
                    diag_prompt = gr.Textbox(
                        label="Diagram Prompt",
                        placeholder="Enter a diagram description...",
                        lines=5,
                        value=example_prompts["Diagram"][0]
                    )
                    diag_generate = gr.Button("Generate Diagram")
                    diag_image = gr.Image(label="Generated Diagram", value="w4.webp")
                    with gr.Accordion("Example Prompts", open=True, elem_classes="example-accordion"):
                        for ex in example_prompts["Diagram"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=diag_prompt)
                
                # Tab 6: Flowchart
                with gr.Tab("Flowchart"):
                    flow_prompt = gr.Textbox(
                        label="Flowchart Prompt",
                        placeholder="Enter a flowchart description...",
                        lines=5,
                        value=example_prompts["Flowchart"][0]
                    )
                    flow_generate = gr.Button("Generate Flowchart")
                    flow_image = gr.Image(label="Generated Flowchart", value="w5.webp")
                    with gr.Accordion("Example Prompts", open=True, elem_classes="example-accordion"):
                        for ex in example_prompts["Flowchart"]:
                            gr.Button(ex, variant="secondary").click(fn=lambda ex=ex: set_prompt(ex), outputs=flow_prompt)
    
    # Bind events for generation buttons
    pd_generate.click(
        fn=generate_image,
        inputs=[pd_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[pd_image, seed_slider]
    )
    
    mindmap_generate.click(
        fn=generate_image,
        inputs=[mindmap_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[mindmap_image, seed_slider]
    )
    
    mock_generate.click(
        fn=generate_image,
        inputs=[mock_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[mock_image, seed_slider]
    )
    
    info_generate.click(
        fn=generate_image,
        inputs=[info_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[info_image, seed_slider]
    )
    
    diag_generate.click(
        fn=generate_image,
        inputs=[diag_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[diag_image, seed_slider]
    )
    
    flow_generate.click(
        fn=generate_image,
        inputs=[flow_prompt, seed_slider, randomize_seed, width_slider, height_slider, steps_slider, guidance_slider],
        outputs=[flow_image, seed_slider]
    )

if __name__ == "__main__":
    demo.queue()
    demo.launch(
        server_name="0.0.0.0",
        server_port=7860,
        share=False,
        show_error=True,
        debug=True
    )