"""
FLUX.2 Turbo LoRA Explorer - HuggingFace Space
A unique space for exploring LoRA styles with FLUX's fast inference

Features:
- 8+ Style LoRA Presets (Anime, Realistic, Sketch, Cyberpunk, etc.)
- AI-Powered Random Prompt Generator
- Image-to-Image generation
- Side-by-side comparison mode
- LoRA strength control
- Full parameter control
"""

# Monkey-patch to fix Gradio bug: TypeError: argument of type 'bool' is not iterable
# and APIInfoParseError: Cannot parse schema True
# See: https://github.com/gradio-app/gradio/issues/11084
def _patch_gradio_client():
    try:
        import gradio_client.utils as client_utils

        # Patch get_type to handle boolean schemas
        original_get_type = getattr(client_utils, 'get_type', None)
        if original_get_type:
            def patched_get_type(schema):
                if isinstance(schema, bool):
                    return "Any"
                return original_get_type(schema)
            client_utils.get_type = patched_get_type

        # Patch _json_schema_to_python_type to handle boolean schemas
        original_json_schema = getattr(client_utils, '_json_schema_to_python_type', None)
        if original_json_schema:
            def patched_json_schema(schema, defs):
                if isinstance(schema, bool):
                    return "Any"
                return original_json_schema(schema, defs)
            client_utils._json_schema_to_python_type = patched_json_schema
    except Exception:
        pass  # If patch fails, continue anyway

_patch_gradio_client()

import gradio as gr
import modal
import os
import random
from PIL import Image
from io import BytesIO
from dotenv import load_dotenv

# Load environment variables from .env file (for local development)
load_dotenv()

# Modal configuration
APP_NAME = "flux2-turbo-lora-explorer"

# Style LoRA presets with descriptions
STYLE_LORAS = {
    "None (Base Model)": {
        "id": "none",
        "description": "Pure FLUX output without any style modification",
    },
    "Anime": {
        "id": "anime",
        "description": "Japanese anime style with vibrant colors and clean lines",
    },
    "Realistic": {
        "id": "realistic",
        "description": "Photorealistic enhancement for lifelike images",
    },
    "Sketch": {
        "id": "sketch",
        "description": "Hand-drawn sketch style with pencil textures",
    },
    "Vintage": {
        "id": "vintage",
        "description": "Retro vintage poster aesthetic from the 50s-60s",
    },
    "Cyberpunk": {
        "id": "cyberpunk",
        "description": "80s cyberpunk with neon lights and futuristic vibes",
    },
    "Watercolor": {
        "id": "watercolor",
        "description": "Soft watercolor painting with flowing colors",
    },
    "Pixel Art": {
        "id": "pixel",
        "description": "Retro pixel art style for gaming aesthetics",
    },
}

ASPECT_RATIOS = {
    "1:1 (1024x1024)": (1024, 1024),
    "16:9 (1344x768)": (1344, 768),
    "9:16 (768x1344)": (768, 1344),
    "4:3 (1152x896)": (1152, 896),
    "3:2 (1216x832)": (1216, 832),
}


def generate_ai_prompt(style: str):
    """Generate a creative prompt - using local fallbacks for now"""
    # Local prompts (Modal integration disabled temporarily)
    fallback_prompts = [
        "a mystical forest with glowing mushrooms and ethereal mist, magical atmosphere",
        "an astronaut riding a horse on Mars, cinematic lighting, epic scale",
        "a cozy coffee shop interior with warm lighting and rain outside the window",
        "a dragon sleeping on a pile of gold coins in a crystal cave",
        "a futuristic city skyline at sunset with flying cars and neon signs",
        "a portrait of a wise old wizard with a long white beard, detailed face",
        "an underwater palace with mermaids and bioluminescent creatures",
        "a steampunk airship flying through clouds at golden hour",
    ]
    return random.choice(fallback_prompts)


def generate_image(
    prompt: str,
    style: str,
    lora_scale: float,
    aspect_ratio: str,
    num_steps: int,
    guidance_scale: float,
    seed: int,
):
    """Generate image using Modal backend"""

    if not prompt.strip():
        raise gr.Error("Please enter a prompt")

    width, height = ASPECT_RATIOS.get(aspect_ratio, (1024, 1024))
    lora_id = STYLE_LORAS.get(style, {}).get("id", "none")

    if seed == -1:
        seed = random.randint(0, 2**32 - 1)

    try:
        Flux2TurboLoRA = modal.Cls.from_name(APP_NAME, "Flux2TurboLoRA")

        image_bytes = Flux2TurboLoRA().generate.remote(
            prompt=prompt,
            lora_id=lora_id,
            lora_scale=lora_scale,
            num_steps=num_steps,
            guidance_scale=guidance_scale,
            width=width,
            height=height,
            seed=seed,
        )

        image = Image.open(BytesIO(image_bytes))

        info = f"""
**Generation Info:**
- Style: {style}
- LoRA Scale: {lora_scale}
- Steps: {num_steps}
- Guidance: {guidance_scale}
- Seed: {seed}
- Size: {width}x{height}
"""
        return image, info, seed

    except Exception as e:
        raise gr.Error(f"Generation failed: {str(e)}")




def compare_styles(
    prompt: str,
    styles: list[str],
    aspect_ratio: str,
    num_steps: int,
    guidance_scale: float,
    seed: int,
):
    """Generate comparison images with different styles"""

    if not prompt.strip():
        raise gr.Error("Please enter a prompt")

    if len(styles) < 2:
        raise gr.Error("Please select at least 2 styles to compare")

    width, height = ASPECT_RATIOS.get(aspect_ratio, (1024, 1024))

    if seed == -1:
        seed = random.randint(0, 2**32 - 1)

    Flux2TurboLoRA = modal.Cls.from_name(APP_NAME, "Flux2TurboLoRA")
    model = Flux2TurboLoRA()

    images = []
    for i, style in enumerate(styles):

        lora_id = STYLE_LORAS.get(style, {}).get("id", "none")

        try:
            image_bytes = model.generate.remote(
                prompt=prompt,
                lora_id=lora_id,
                lora_scale=1.0,
                num_steps=num_steps,
                guidance_scale=guidance_scale,
                width=width,
                height=height,
                seed=seed,
            )

            img = Image.open(BytesIO(image_bytes))
            images.append((img, style))

        except Exception as e:
            images.append((None, f"{style} - Error"))

    return images


def random_seed():
    """Generate random seed"""
    return random.randint(0, 2**32 - 1)


# Build Gradio Interface
with gr.Blocks(
    title="FLUX LoRA Explorer",
    theme=gr.themes.Soft(
        primary_hue="violet",
        secondary_hue="slate",
    ),
    css="""
    .header { text-align: center; margin-bottom: 1.5rem; }
    .header h1 { font-size: 2.5rem; margin-bottom: 0.5rem; background: linear-gradient(135deg, #667eea 0%, #764ba2 100%); -webkit-background-clip: text; -webkit-text-fill-color: transparent; }
    .style-card { border: 2px solid #e0e0e0; border-radius: 12px; padding: 1rem; margin: 0.5rem; transition: all 0.2s; }
    .style-card:hover { border-color: #667eea; transform: translateY(-2px); }
    footer { display: none !important; }
    """
) as demo:

    gr.HTML("""
    <div class="header">
        <h1>🎨 FLUX.2-dev-Turbo LoRA Explorer</h1>
        <p>6x faster image generation with 8-step Turbo inference</p>
        <p><strong>⚡ 8-Step Turbo</strong> • <strong>8+ Style LoRAs</strong> • <strong>🤖 AI Prompts</strong> • <strong>🖼️ Image-to-Image</strong></p>
    </div>
    """)

    with gr.Tabs():
        # ============ TAB 1: GENERATE ============
        with gr.TabItem("🎨 Generate", id="generate"):
            with gr.Row():
                with gr.Column(scale=1):
                    prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Describe the image you want to create...",
                        lines=3,
                    )

                    style = gr.Dropdown(
                        choices=list(STYLE_LORAS.keys()),
                        value="None (Base Model)",
                        label="🎭 Style LoRA",
                    )

                    style_description = gr.Markdown(
                        value="*Pure FLUX output without any style modification*"
                    )

                    with gr.Row():
                        ai_prompt_btn = gr.Button("🤖 AI Prompt", variant="secondary", size="sm")
                        random_btn = gr.Button("🎲 Random Seed", size="sm")

                    with gr.Accordion("⚙️ Advanced Settings", open=False):
                        lora_scale = gr.Slider(
                            minimum=0.0,
                            maximum=1.5,
                            value=1.0,
                            step=0.05,
                            label="LoRA Strength",
                            info="How strongly the style affects the output"
                        )

                        aspect_ratio = gr.Dropdown(
                            choices=list(ASPECT_RATIOS.keys()),
                            value="1:1 (1024x1024)",
                            label="Aspect Ratio",
                        )

                        with gr.Row():
                            num_steps = gr.Slider(
                                minimum=4,
                                maximum=50,
                                value=8,  # Turbo: 8 steps (6x faster!)
                                step=1,
                                label="Steps (Turbo: 8)",
                            )
                            guidance_scale = gr.Slider(
                                minimum=1.0,
                                maximum=10.0,
                                value=2.5,  # Turbo optimal
                                step=0.1,
                                label="Guidance Scale",
                            )

                        seed = gr.Number(
                            value=-1,
                            label="Seed (-1 for random)",
                            precision=0,
                        )

                    generate_btn = gr.Button("🚀 Generate", variant="primary", size="lg")

                with gr.Column(scale=1):
                    output_image = gr.Image(
                        label="Generated Image",
                        type="pil",
                        height=512,
                    )
                    output_info = gr.Markdown()
                    output_seed = gr.Number(label="Used Seed", interactive=False)

            # Update style description
            def update_description(style):
                desc = STYLE_LORAS.get(style, {}).get("description", "")
                return f"*{desc}*"

            style.change(update_description, inputs=style, outputs=style_description)
            ai_prompt_btn.click(generate_ai_prompt, inputs=style, outputs=prompt)
            random_btn.click(random_seed, outputs=seed)

            generate_btn.click(
                generate_image,
                inputs=[prompt, style, lora_scale, aspect_ratio, num_steps, guidance_scale, seed],
                outputs=[output_image, output_info, output_seed],
            )

        # ============ TAB 2: IMAGE TO IMAGE ============
        with gr.TabItem("🖼️ Image to Image", id="img2img"):
            gr.Markdown("""
            ### 🚧 Coming Soon

            **Image-to-Image is not yet supported for FLUX.2-dev-Turbo.**

            The diffusers library hasn't released `FluxImg2ImgPipeline` for FLUX.2 yet.
            This feature will be enabled once official support is available.

            In the meantime, try the **Generate** tab for text-to-image generation!

            ---

            *Expected features when available:*
            - Transform existing images with text descriptions
            - Control transformation strength (subtle to dramatic changes)
            - Apply style LoRAs to your images
            """)

        # ============ TAB 3: COMPARE STYLES ============
        with gr.TabItem("⚖️ Compare Styles", id="compare"):
            gr.Markdown("""
            ### Compare Multiple Styles
            Generate the same prompt with different style LoRAs to find your favorite look.
            All images use the same seed for fair comparison.
            """)

            with gr.Row():
                with gr.Column(scale=1):
                    compare_prompt = gr.Textbox(
                        label="Prompt",
                        placeholder="Enter a prompt to compare across styles...",
                        lines=2,
                        value="a majestic dragon flying over a medieval castle",
                    )

                    compare_styles_select = gr.CheckboxGroup(
                        choices=list(STYLE_LORAS.keys()),
                        value=["None (Base Model)", "Anime", "Realistic", "Cyberpunk"],
                        label="Select Styles to Compare (2-4 recommended)",
                    )

                    with gr.Row():
                        compare_aspect = gr.Dropdown(
                            choices=list(ASPECT_RATIOS.keys()),
                            value="1:1 (1024x1024)",
                            label="Aspect Ratio",
                        )
                        compare_seed = gr.Number(
                            value=-1,
                            label="Seed (-1 for random)",
                            precision=0,
                        )

                    with gr.Row():
                        compare_steps = gr.Slider(4, 50, value=8, step=1, label="Steps (Turbo: 8)")
                        compare_guidance = gr.Slider(1.0, 10.0, value=2.5, step=0.1, label="Guidance")

                    compare_btn = gr.Button("⚖️ Compare Styles", variant="primary")

                with gr.Column(scale=2):
                    compare_gallery = gr.Gallery(
                        label="Style Comparison",
                        columns=2,
                        rows=2,
                        height=600,
                        object_fit="contain",
                    )

            compare_btn.click(
                compare_styles,
                inputs=[compare_prompt, compare_styles_select, compare_aspect, compare_steps, compare_guidance, compare_seed],
                outputs=compare_gallery,
            )

        # ============ TAB 4: STYLE GALLERY ============
        with gr.TabItem("🖼️ Style Gallery", id="gallery"):
            gr.Markdown("### Available Style LoRAs")
            gr.Markdown("Browse all available style presets and their descriptions.")

            for name, info in STYLE_LORAS.items():
                with gr.Row():
                    gr.Markdown(f"""
                    **{name}**
                    - ID: `{info['id']}`
                    - {info['description']}
                    """)

        # ============ TAB 5: ABOUT ============
        with gr.TabItem("ℹ️ About", id="about"):
            gr.Markdown("""
            ## About FLUX LoRA Explorer

            This space lets you explore different artistic styles using LoRA adapters
            with FLUX's powerful image generation.

            ### Features

            | Feature | Description |
            |---------|-------------|
            | **8+ Style Presets** | Anime, Realistic, Sketch, Cyberpunk, Watercolor, and more |
            | **AI Prompt Generator** | Let AI create creative prompts for you using Qwen2.5 |
            | **Image-to-Image** | Transform existing images with text descriptions |
            | **Style Comparison** | Compare multiple styles side-by-side with the same seed |
            | **LoRA Strength Control** | Fine-tune how strongly the style affects your image |
            | **Full Parameter Control** | Steps, guidance, seed, aspect ratio |

            ### How LoRAs Work

            LoRA (Low-Rank Adaptation) is a technique that allows fine-tuning large models
            efficiently. Each style LoRA was trained on specific artistic styles and can be
            applied on top of the base FLUX model to transform your images.

            ### Tips

            - **LoRA Strength**: Start with 1.0 and adjust. Lower values (0.5-0.8) give subtle effects.
            - **Steps**: More steps = better quality but slower. 28 is a good balance.
            - **Guidance**: Higher values follow the prompt more closely. 3.5 is recommended.
            - **Seed**: Use the same seed to compare styles fairly.
            - **Image-to-Image Strength**: 0.3-0.5 for subtle edits, 0.7-0.9 for dramatic changes.

            ### Credits

            - Base Model: [FLUX.1-dev](https://huggingface.co/black-forest-labs/FLUX.1-dev) by Black Forest Labs
            - AI Prompts: [Qwen2.5-1.5B-Instruct](https://huggingface.co/Qwen/Qwen2.5-1.5B-Instruct)
            - GPU Backend: [Modal.com](https://modal.com)
            - Community LoRAs from various HuggingFace creators
            """)

if __name__ == "__main__":
    demo.launch(ssr_mode=False)