import os
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image
from typing import Iterable

# Pipeline for FLUX.2 Klein
from diffusers import Flux2KleinPipeline
from diffusers.utils import load_image
from huggingface_hub import hf_hub_download

# --- Hardware and Theme Setup ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

colors.orange_red = colors.Color(
    name="orange_red", c50="#FFF0E5", c100="#FFE0CC", c200="#FFC299", c300="#FFA366",
    c400="#FF8533", c500="#FF4500", c600="#E63E00", c700="#CC3700", c800="#B33000",
    c900="#992900", c950="#802200",
)

class OrangeRedTheme(Soft):
    def __init__(
        self, *, primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.orange_red,
        neutral_hue: colors.Color | str = colors.slate, text_size: sizes.Size | str = sizes.text_lg,
        font: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("Outfit"), "Arial", "sans-serif",
        ),
        font_mono: fonts.Font | str | Iterable[fonts.Font | str] = (
            fonts.GoogleFont("IBM Plex Mono"), "ui-monospace", "monospace",
        ),
    ):
        super().__init__(
            primary_hue=primary_hue, secondary_hue=secondary_hue, neutral_hue=neutral_hue,
            text_size=text_size, font=font, font_mono=font_mono,
        )
        super().set(
            background_fill_primary="*primary_50",
            background_fill_primary_dark="*primary_900",
            body_background_fill="linear-gradient(135deg, *primary_200, *primary_100)",
            body_background_fill_dark="linear-gradient(135deg, *primary_900, *primary_800)",
            button_primary_text_color="white",
            button_primary_text_color_hover="white",
            button_primary_background_fill="linear-gradient(90deg, *secondary_500, *secondary_600)",
            button_primary_background_fill_hover="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_dark="linear-gradient(90deg, *secondary_600, *secondary_700)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_600)",
            slider_color="*secondary_500",
            slider_color_dark="*secondary_600",
            block_title_text_weight="600", block_border_width="3px",
            block_shadow="*shadow_drop_lg", button_primary_shadow="*shadow_drop_lg",
            button_large_padding="11px", color_accent_soft="*primary_100",
            block_label_background_fill="*primary_200",
        )

orange_red_theme = OrangeRedTheme()
MAX_SEED = np.iinfo(np.int32).max

# --- Model Loading ---
print("Loading FLUX.2 Klein 9B model...")
pipe = Flux2KleinPipeline.from_pretrained(
    "black-forest-labs/FLUX.2-klein-9B",
    torch_dtype=torch.bfloat16
).to(device)
print("Model loaded successfully.")

# --- LoRA Loading (Updated) ---
print("Loading new LoRA adapters...")
pipe.load_lora_weights(
    "markury/flux2k9b-simpletuner-lora-loona",
    weight_name="pytorch_lora_weights.safetensors",
    adapter_name="simple-tuner"
)
pipe.load_lora_weights(
    "linoyts/Flux2-Klein-Delight-LoRA",
    weight_name="pytorch_lora_weights_v2.safetensors",
    adapter_name="klein-delight"
)
print("All LoRA adapters loaded.")

# Updated map for the new adapters
ADAPTER_MAP = {
    "Simple-Tuner": "simple-tuner",
    "Klein-Delight-Style": "klein-delight",
}

@spaces.GPU
def infer(input_image, prompt, lora_adapter, seed=42, randomize_seed=True, guidance_scale=4.0, steps=4, progress=gr.Progress(track_tqdm=True)):
    # Input image is required for image-to-image tasks
    if not input_image:
        raise gr.Error("Please upload an image to apply a style to.")

    # Dynamically set the adapter based on the dropdown choice
    adapter_name = ADAPTER_MAP.get(lora_adapter)
    if adapter_name:
        print(f"Activating LoRA: {lora_adapter} ({adapter_name})")
        pipe.set_adapters([adapter_name], adapter_weights=[1.0])
    else:
        # If "None" is selected (or an invalid choice), disable LoRAs
        print("No LoRA selected. Disabling adapters.")
        pipe.disable_lora()
        
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    
    original_image = input_image.copy().convert("RGB")
    
    image = pipe(
        image=original_image, 
        prompt=prompt,
        guidance_scale=guidance_scale,
        width=original_image.size[0],
        height=original_image.size[1],
        num_inference_steps=steps,
        generator=torch.Generator(device=device).manual_seed(seed),
    ).images[0]

    return image, seed

@spaces.GPU
def infer_example(input_image, prompt, lora_adapter):
    # Use a fixed seed for reproducible examples
    image, seed = infer(input_image, prompt, lora_adapter, seed=12345, randomize_seed=False)
    return image, seed

# --- UI Layout ---
css="""
#col-container { margin: 0 auto; max-width: 960px; }
#main-title h1 { font-size: 2.2em !important; }
"""

with gr.Blocks() as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **FLUX.2 Klein LoRA Stylizer**", elem_id="main-title")
        gr.Markdown(
            "Apply creative styles to your images using **FLUX.2-klein-9B** and specialized LoRA adapters. "
            "Upload an image, select a style, and write a prompt to guide the transformation."
        )
        
        with gr.Row(equal_height=True):
            with gr.Column():
                input_image = gr.Image(label="Upload Image", type="pil", height=290, sources=["upload", "webcam", "clipboard"])
                prompt = gr.Text(label="Guiding Prompt", show_label=True, placeholder="e.g., a man with a red superhero mask")
                
                lora_adapter = gr.Dropdown(
                    label="Choose a Creative Style",
                    # Updated choices for the new adapters
                    choices=["Simple-Tuner", "Klein-Delight-Style"],
                    value="Klein-Delight-Style"
                )

                run_button = gr.Button("Apply Style", variant="primary")

                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)
                    # Updated defaults suitable for FLUX.2 Klein
                    guidance_scale = gr.Slider(label="Guidance Scale", minimum=0.0, maximum=10.0, step=0.1, value=4.0)       
                    steps = gr.Slider(label="Steps", minimum=1, maximum=50, value=4, step=1)
                    
            with gr.Column():
                output_image = gr.Image(label="Stylized Image", interactive=False, format="png", height=450)
                used_seed = gr.Textbox(label="Used Seed", interactive=False)

        # Updated examples for the new LoRAs
        gr.Examples(
            examples=[
                ["examples/animal.jpg", "a cute red panda, charming and delightful illustration, soft lighting", "Klein-Delight-Style"],
            ],
            inputs=[input_image, prompt, lora_adapter],
            outputs=[output_image, used_seed],
            fn=infer_example,
            cache_examples=False,
        )
            
    run_button.click(
        fn=infer,
        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
        outputs=[output_image, used_seed]
    )

if __name__ == "__main__":
    demo.queue().launch(css=css, theme=orange_red_theme, show_error=True)