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

# --- Custom Local Imports ---
# Note: Ensure these files (pipeline_qwenimage_edit_plus.py, etc.) 
# are present in the same directory or installed in the environment.
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

# --- Theme Imports ---
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes

# --- Custom Theme Definition ---
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)",
            button_secondary_text_color="black",
            button_secondary_text_color_hover="white",
            button_secondary_background_fill="linear-gradient(90deg, *primary_300, *primary_300)",
            button_secondary_background_fill_hover="linear-gradient(90deg, *primary_400, *primary_400)",
            button_secondary_background_fill_dark="linear-gradient(90deg, *primary_500, *primary_600)",
            button_secondary_background_fill_hover_dark="linear-gradient(90deg, *primary_500, *primary_500)",
            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()

# --- Device Setup ---
device = torch.device("cuda" if torch.cuda.is_available() else "cpu")

print("CUDA_VISIBLE_DEVICES=", os.environ.get("CUDA_VISIBLE_DEVICES"))
print("torch.__version__ =", torch.__version__)
print("torch.version.cuda =", torch.version.cuda)
print("cuda available:", torch.cuda.is_available())
print("cuda device count:", torch.cuda.device_count())
if torch.cuda.is_available():
    print("current device:", torch.cuda.current_device())
    print("device name:", torch.cuda.get_device_name(torch.cuda.current_device()))

print("Using device:", device)

# --- Model Loading ---
dtype = torch.bfloat16

pipe = QwenImageEditPlusPipeline.from_pretrained(
    "Qwen/Qwen-Image-Edit-2509",
    transformer=QwenImageTransformer2DModel.from_pretrained(
        "linoyts/Qwen-Image-Edit-Rapid-AIO", 
        subfolder='transformer',
        torch_dtype=dtype,
        device_map='cuda'
    ),
    torch_dtype=dtype
).to(device)

# Apply FA3 Optimization
try:
    pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())
    print("Flash Attention 3 Processor set successfully.")
except Exception as e:
    print(f"Warning: Could not set FA3 processor: {e}")

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

# --- Dynamic LoRA Configuration ---
# These are architectural placeholders. To make the styles work, update 'repo' and 'weights' 
# to point to actual HuggingFace repositories containing valid LoRA weights.
ADAPTER_SPECS = {
    "Cinematic-DSLR": {
        "repo": "prithivMLmods/Qwen-Image-Edit-2509-LoRAs-Fast",
        "weights": "placeholder_weights.safetensors", 
        "adapter_name": "cinematic-dslr",
        "description": "High-end cinema look with professional color grading."
    },
    "Portrait-Pro": {
        "repo": "prithivMLmods/Qwen-Image-Edit-2509-LoRAs-Fast",
        "weights": "placeholder_weights.safetensors",
        "adapter_name": "portrait-pro",
        "description": "Optimized for studio portrait lighting and skin detail."
    },
    "High-Key-Lighting": {
        "repo": "prithivMLmods/Qwen-Image-Edit-2509-LoRAs-Fast",
        "weights": "placeholder_weights.safetensors",
        "adapter_name": "high-key",
        "description": "Bright, even lighting typical of commercial photography."
    },
    "Editorial-Style": {
        "repo": "prithivMLmods/Qwen-Image-Edit-2509-LoRAs-Fast",
        "weights": "placeholder_weights.safetensors",
        "adapter_name": "editorial",
        "description": "Magazine-style composition and contrast."
    }
}

# Track what is currently loaded in memory for hot-swapping
LOADED_ADAPTERS = set()

def update_dimensions_on_upload(image):
    """Calculates optimal dimensions based on image aspect ratio."""
    if image is None:
        return 1024, 1024
    
    original_width, original_height = image.size
    
    if original_width > original_height:
        new_width = 1024
        aspect_ratio = original_height / original_width
        new_height = int(new_width * aspect_ratio)
    else:
        new_height = 1024
        aspect_ratio = original_width / original_height
        new_width = int(new_height * aspect_ratio)
        
    # Ensure dimensions are multiples of 8 (standard for diffusion models)
    new_width = (new_width // 8) * 8
    new_height = (new_height // 8) * 8
    
    return new_width, new_height

@spaces.GPU
def infer(
    input_image,
    prompt,
    lora_adapter,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True)
):
    """
    Main inference function with dynamic LoRA hot-loading.
    """
    # Cleanup memory before starting
    gc.collect()
    torch.cuda.empty_cache()

    if input_image is None:
        raise gr.Error("Please upload an image to edit.")

    # 1. Get Config for Selected Adapter
    spec = ADAPTER_SPECS.get(lora_adapter)
    if not spec:
        # Fallback to base model if config missing
        print(f"Configuration not found for: {lora_adapter}. Using base model.")
        adapter_name = "base"
    else:
        adapter_name = spec["adapter_name"]

    # 2. Lazy Loading Logic (Hot Swapping)
    # Only loads if not currently in memory to save bandwidth/startup time
    if spec and adapter_name not in LOADED_ADAPTERS:
        print(f"--- Hot Loading Adapter: {lora_adapter} ---")
        try:
            pipe.load_lora_weights(
                spec["repo"], 
                weight_name=spec["weights"], 
                adapter_name=adapter_name
            )
            LOADED_ADAPTERS.add(adapter_name)
        except Exception as e:
            # Fallback for demonstration if placeholder weights don't exist
            print(f"Info: Could not load weights for {lora_adapter}: {e}")
            gr.Warning(f"Could not load specific style weights for '{lora_adapter}'. Using base model instead.")
            # Ensure we don't try to set this adapter if it failed to load
            adapter_name = "base" 
    else:
        print(f"--- Adapter {lora_adapter} already active in memory or using base model. ---")

    # 3. Activate the specific adapter
    # If 'base' (or fallback), we disable adapters. Otherwise, set the specific one.
    if adapter_name == "base":
        pipe.disable_lora()
    else:
        pipe.set_adapters([adapter_name], adapter_weights=[1.0])

    # 4. Standard Inference Setup
    if randomize_seed:
        seed = random.randint(0, MAX_SEED)

    generator = torch.Generator(device=device).manual_seed(seed)
    negative_prompt = "worst quality, low quality, bad anatomy, bad hands, text, error, missing fingers, extra digit, fewer digits, cropped, jpeg artifacts, signature, watermark, username, blurry"

    original_image = input_image.convert("RGB")
    width, height = update_dimensions_on_upload(original_image)

    try:
        result = pipe(
            image=original_image,
            prompt=prompt,
            negative_prompt=negative_prompt,
            height=height,
            width=width,
            num_inference_steps=steps,
            generator=generator,
            true_cfg_scale=guidance_scale,
        ).images[0]
        
        return result, seed

    except Exception as e:
        raise gr.Error(f"Error during inference: {e}")
    finally:
        # Cleanup
        gc.collect()
        torch.cuda.empty_cache()

@spaces.GPU
def infer_example(input_image, prompt, lora_adapter):
    """Helper function for Gradio Examples."""
    if input_image is None:
        return None, 0
    
    input_pil = input_image.convert("RGB")
    guidance_scale = 1.0
    steps = 4
    result, seed = infer(input_pil, prompt, lora_adapter, 0, True, guidance_scale, steps)
    return result, seed

# --- Gradio 6 Application ---
# Gradio 6 Syntax: gr.Blocks() takes NO parameters. All config goes in demo.launch()

with gr.Blocks() as demo:
    gr.HTML("""
    <div style="display: flex; justify-content: space-between; align-items: center; margin-bottom: 10px;">
        <h1 style="margin: 0;">Qwen-Image-Edit-2509-LoRAs-Fast</h1>
    </div>
    """)
    
    gr.Markdown(
        "Perform diverse image edits using specialized [LoRA](https://huggingface.co/models?other=base_model:adapter:Qwen/Qwen-Image-Edit-2509) "
        "adapters for the [Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) model. "
        "This demo features **dynamic hot-loading**, downloading LoRA weights only when you select them."
    )

    with gr.Row(equal_height=True):
        with gr.Column():
            input_image = gr.Image(label="Upload Image", type="pil", height=290)
            
            prompt = gr.Text(
                label="Edit Prompt",
                show_label=True,
                placeholder="e.g., apply cinematic lighting...",
                lines=2
            )

            run_button = gr.Button("Edit Image", variant="primary", size="lg")

        with gr.Column():
            output_image = gr.Image(label="Output Image", interactive=False, format="png", height=353)
            
            with gr.Row():
                # Dynamic keys based on the config dict
                lora_adapter = gr.Dropdown(
                    label="Choose Editing Style",
                    choices=list(ADAPTER_SPECS.keys()),
                    value="Cinematic-DSLR",
                    info="Select a style to hot-load"
                )
            
            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)
                guidance_scale = gr.Slider(label="Guidance Scale", minimum=1.0, maximum=10.0, step=0.1, value=1.0)
                steps = gr.Slider(label="Inference Steps", minimum=1, maximum=50, step=1, value=4)
    
    gr.Examples(
        examples=[
            ["examples/1.jpg", "Apply cinematic dslr style.", "Cinematic-DSLR"],
            ["examples/5.jpg", "Enhance portrait lighting.", "Portrait-Pro"],
            ["examples/4.jpg", "Switch to high key lighting.", "High-Key-Lighting"],
        ],
        inputs=[input_image, prompt, lora_adapter],
        outputs=[output_image, seed],
        fn=infer_example,
        cache_examples=False,
        label="Examples"
    )

    # Gradio 6 Event Listeners
    run_button.click(
        fn=infer,
        inputs=[input_image, prompt, lora_adapter, seed, randomize_seed, guidance_scale, steps],
        outputs=[output_image, seed],
        api_visibility="public"
    )

css="""
#col-container {
    margin: 0 auto;
    max-width: 1000px;
}
.gradio-container {
    font-family: 'Outfit', sans-serif !important;
}
"""

if __name__ == "__main__":
    # Gradio 6 Launch Syntax
    # All app-level parameters (theme, css, footer_links) go here.
    demo.queue(max_size=30).launch(
        css=css, 
        theme=orange_red_theme, 
        mcp_server=True, 
        ssr_mode=False, 
        show_error=True,
        footer_links=[{"label": "Built with anycoder", "url": "https://huggingface.co/spaces/akhaliq/anycoder"}]
    )