try:
    import spaces
    SPACES_AVAILABLE = True
    print("✅ Spaces available - ZeroGPU mode")
except ImportError:
    SPACES_AVAILABLE = False
    print("⚠️ Spaces not available - running in regular mode")

import gradio as gr
import torch
from diffusers import StableDiffusionPipeline
from PIL import Image
import datetime
import io
import json
import os
from typing import Optional

# ======================
# 配置区（你只需修改这里即可扩展）
# ======================

# 1. 基础模型
BASE_MODEL = "SG161222/RealisticVisionV6.0"

# 2. 固定LoRA（不可选，自动加载）
FIXED_LORAS = [
    ("Lykon/epiCRealism_LoRA", 0.8),      # 质量增强
    ("latent-consistency/lora-dreamshaper", 0.7),  # 姿势控制
]

# 3. 风格模板（自动拼接到用户提示词前）
STYLE_PROMPTS = {
    "None": "",
    "Realistic": "photorealistic, ultra-detailed skin, natural lighting, 8k, professional photography, f/1.8, shallow depth of field, Canon EOS R5, ",
    "Anime": "anime style, cel shading, vibrant colors, detailed eyes, studio ghibli, trending on pixiv, ",
    "Comic": "comic book style, bold outlines, dynamic angles, comic panel, Marvel style, inked lines, ",
    "Watercolor": "watercolor painting, soft brush strokes, translucent layers, artistic, painterly, paper texture, ",
}

# 4. 可选LoRA下拉菜单（用户可选1个，None表示清除）
OPTIONAL_LORAS = [
    "None",
    "Add Detail: https://huggingface.co/latent-consistency/lora-add-detail",
    "Vintage Photo: https://huggingface.co/ckpt/LoRA-vintage-photo",
    "Cinematic: https://huggingface.co/latent-consistency/lora-cinematic",
    "Portrait Enhancer: https://huggingface.co/deforum/Portrait-Enhancer-LoRA",
    "Soft Focus: https://huggingface.co/latent-consistency/lora-soft-focus",
]

# 解析可选LoRA的名称和ID
OPTIONAL_LORA_MAP = {}
for item in OPTIONAL_LORAS:
    if item != "None":
        name, url = item.split(": ", 1)
        OPTIONAL_LORA_MAP[name] = url
    else:
        OPTIONAL_LORA_MAP["None"] = None

# 默认参数
DEFAULT_SEED = -1
DEFAULT_WIDTH = 1024
DEFAULT_HEIGHT = 1024
DEFAULT_LORA_SCALE = 0.8
DEFAULT_STEPS = 30
DEFAULT_CFG = 7.5

# ======================
# 全局变量：延迟加载模型
# ======================
pipe = None
device = "cuda" if torch.cuda.is_available() else "cpu"

def load_pipeline():
    global pipe
    if pipe is None:
        print("🚀 Loading base model...")
        pipe = StableDiffusionPipeline.from_pretrained(
            BASE_MODEL,
            torch_dtype=torch.float16,
            safety_checker=None,
            requires_safety_checker=False,
        ).to(device)
        pipe.enable_attention_slicing()
        pipe.enable_vae_slicing()
        pipe.enable_model_cpu_offload()  # 适配ZeroGPU
        print("✅ Base model loaded.")
    return pipe

def unload_pipeline():
    global pipe
    if pipe is not None:
        del pipe
        torch.cuda.empty_cache()
        pipe = None
        print("🗑️ Pipeline unloaded.")

# ======================
# 主生成函数
# ======================
def generate_image(
    prompt, negative_prompt, style, seed, width, height, optional_lora_name, lora_scale,
    steps, cfg_scale
):
    global pipe

    # 加载模型（懒加载）
    pipe = load_pipeline()

    # 处理种子
    if seed == -1:
        seed = torch.randint(0, 2**32, (1,)).item()
    generator = torch.Generator(device=device).manual_seed(seed)

    # 拼接风格提示词
    full_prompt = STYLE_PROMPTS[style] + prompt
    full_negative_prompt = negative_prompt

    # 加载固定LoRA（每次生成前都加载，确保状态正确）
    for lora_id, scale in FIXED_LORAS:
        pipe.load_lora_weights(lora_id, adapter_name=lora_id)
        pipe.set_adapters([lora_id], adapter_weights=[scale])

    # 加载可选LoRA（如果非None）
    if optional_lora_name != "None":
        lora_url = OPTIONAL_LORA_MAP[optional_lora_name]
        pipe.load_lora_weights(lora_url, adapter_name=optional_lora_name)
        pipe.set_adapters([lora_id for lora_id, _ in FIXED_LORAS] + [optional_lora_name], 
                          adapter_weights=[scale for _, scale in FIXED_LORAS] + [lora_scale])
    else:
        # 清除所有可选LoRA，只保留固定
        pipe.set_adapters([lora_id for lora_id, _ in FIXED_LORAS], 
                          adapter_weights=[scale for _, scale in FIXED_LORAS])

    # 生成图像
    image = pipe(
        prompt=full_prompt,
        negative_prompt=full_negative_prompt,
        num_inference_steps=steps,
        guidance_scale=cfg_scale,
        width=width,
        height=height,
        generator=generator,
    ).images[0]

    # 生成元数据
    metadata = {
        "prompt": full_prompt,
        "negative_prompt": full_negative_prompt,
        "base_model": BASE_MODEL,
        "fixed_loras": [lora_id for lora_id, _ in FIXED_LORAS],
        "optional_lora": optional_lora_name if optional_lora_name != "None" else None,
        "lora_scale": lora_scale,
        "seed": seed,
        "steps": steps,
        "cfg_scale": cfg_scale,
        "style": style,
        "width": width,
        "height": height,
        "timestamp": datetime.datetime.now().isoformat()
    }

    # 生成文件名
    timestamp = datetime.datetime.now().strftime("%y%m%d%H%M")
    filename_base = f"{seed}-{timestamp}"

    # 保存为WebP（高质量）
    img_buffer = io.BytesIO()
    image.save(img_buffer, format="WEBP", quality=95, method=6)
    img_buffer.seek(0)

    # 保存元数据为TXT
    metadata_buffer = io.StringIO()
    json.dump(metadata, metadata_buffer, indent=2, ensure_ascii=False)
    metadata_buffer.seek(0)

    # 返回：图像、元数据、文件名
    return (
        image,
        json.dumps(metadata, indent=2, ensure_ascii=False),
        f"{filename_base}.webp",
        f"{filename_base}.txt",
        img_buffer.getvalue(),
        metadata_buffer.getvalue().encode('utf-8')
    )

# ======================
# Gradio UI
# ======================
with gr.Blocks(
    theme=gr.themes.Soft(
        primary_hue="blue",
        secondary_hue="green",
        neutral_hue="slate",
    ).set(
        body_background_fill="linear-gradient(135deg, #1e40af, #059669)",
        button_primary_background_fill="white",
        button_primary_text_color="#1e40af",
        input_background_fill="rgba(255,255,255,0.9)",
    ),
    css="""
    body { font-family: 'Helvetica Neue', 'Segoe UI', 'Arial', sans-serif; }
    .gr-button { font-family: 'Helvetica Neue', 'Arial', sans-serif; font-weight: 500; }
    .gr-textarea { font-family: 'Consolas', 'Monaco', 'Courier New', monospace; }
    """,
) as demo:
    gr.Markdown(
        """
        # 🎨 AI Photo Generator (RealisticVision + LoRA)
        **PRO + ZeroGPU Optimized | Multi-LoRA | Style Templates | Metadata Export**
        """
    )

    with gr.Row():
        with gr.Column(scale=3):
            # a. 提示词输入框
            prompt_input = gr.Textbox(
                label="Prompt (Positive)",
                placeholder="A beautiful woman, golden hour, soft sunlight...",
                lines=5,
                max_lines=20,
                elem_classes=["gr-textarea"]
            )

            # b. 负提示词输入框
            negative_prompt_input = gr.Textbox(
                label="Negative Prompt",
                placeholder="blurry, low quality, deformed, cartoon, anime, text, watermark...",
                lines=5,
                max_lines=20,
                elem_classes=["gr-textarea"]
            )

            # c. 风格选择（单选）
            style_radio = gr.Radio(
                choices=list(STYLE_PROMPTS.keys()),
                label="Style",
                value="Realistic",
                elem_classes=["gr-radio"]
            )

            # d. 种子选择
            with gr.Row():
                seed_input = gr.Slider(
                    minimum=-1,
                    maximum=99999999,
                    step=1,
                    value=DEFAULT_SEED,
                    label="Seed (-1 = Random)"
                )
                seed_reset = gr.Button("Reset Seed")

            # e. 宽度选择
            with gr.Row():
                width_input = gr.Slider(
                    minimum=512,
                    maximum=1536,
                    step=64,
                    value=DEFAULT_WIDTH,
                    label="Width"
                )
                width_reset = gr.Button("Reset Width")

            # f. 高度选择
            with gr.Row():
                height_input = gr.Slider(
                    minimum=512,
                    maximum=1536,
                    step=64,
                    value=DEFAULT_HEIGHT,
                    label="Height"
                )
                height_reset = gr.Button("Reset Height")

            # g. LoRA选择（下拉）
            optional_lora_dropdown = gr.Dropdown(
                choices=list(OPTIONAL_LORA_MAP.keys()),
                label="Optional LoRA",
                value="None",
                elem_classes=["gr-dropdown"]
            )

            # h. LoRA控制
            with gr.Row():
                lora_scale_slider = gr.Slider(
                    minimum=0.0,
                    maximum=1.5,
                    step=0.05,
                    value=DEFAULT_LORA_SCALE,
                    label="LoRA Scale"
                )
                lora_reset = gr.Button("Reset LoRA Scale")

            # i. 功能控制（Steps & CFG）
            with gr.Row():
                steps_slider = gr.Slider(
                    minimum=10,
                    maximum=100,
                    step=1,
                    value=DEFAULT_STEPS,
                    label="Steps"
                )
                cfg_slider = gr.Slider(
                    minimum=1.0,
                    maximum=20.0,
                    step=0.5,
                    value=DEFAULT_CFG,
                    label="CFG Scale"
                )
                gen_reset = gr.Button("Reset Generation")

            # m. 生成按钮
            generate_btn = gr.Button("✨ Generate Image", variant="primary", size="lg")

        with gr.Column(scale=2):
            # j. 图片显示区
            image_output = gr.Image(label="Generated Image", height=768, format="webp")

            # k. 元数据显示区
            metadata_output = gr.Textbox(
                label="Metadata (JSON)",
                lines=12,
                max_lines=20,
                elem_classes=["gr-textarea"]
            )

            # l. 下载按钮（并列）
            with gr.Row():
                download_img_btn = gr.Button("⬇️ Download Image (WebP)")
                download_meta_btn = gr.Button("⬇️ Download Metadata (TXT)")

            # 隐藏文件输出（用于下载）
            hidden_img_file = gr.File(visible=False)
            hidden_meta_file = gr.File(visible=False)

    # ======================
    # 事件绑定
    # ======================

    # 重置种子
    seed_reset.click(fn=lambda: -1, outputs=seed_input)
    # 重置宽度
    width_reset.click(fn=lambda: DEFAULT_WIDTH, outputs=width_input)
    # 重置高度
    height_reset.click(fn=lambda: DEFAULT_HEIGHT, outputs=height_input)
    # 重置LoRA缩放
    lora_reset.click(fn=lambda: DEFAULT_LORA_SCALE, outputs=lora_scale_slider)
    # 重置生成参数
    gen_reset.click(
        fn=lambda: (DEFAULT_STEPS, DEFAULT_CFG),
        outputs=[steps_slider, cfg_slider]
    )

    # 生成
    generate_btn.click(
        fn=generate_image,
        inputs=[
            prompt_input, negative_prompt_input, style_radio,
            seed_input, width_input, height_input,
            optional_lora_dropdown, lora_scale_slider,
            steps_slider, cfg_slider
        ],
        outputs=[
            image_output, metadata_output,
            hidden_img_file, hidden_meta_file,
            hidden_img_file, hidden_meta_file
        ]
    )

    # 下载图片
    download_img_btn.click(
        fn=None,
        inputs=[hidden_img_file],
        outputs=None,
        js="(f) => { const a = document.createElement('a'); a.href = f; a.download = f.split('/').pop(); document.body.appendChild(a); a.click(); document.body.removeChild(a); }"
    )

    # 下载元数据
    download_meta_btn.click(
        fn=None,
        inputs=[hidden_meta_file],
        outputs=None,
        js="(f) => { const a = document.createElement('a'); a.href = f; a.download = f.split('/').pop(); document.body.appendChild(a); a.click(); document.body.removeChild(a); }"
    )

    # 设置文件下载（通过返回值触发）
    generate_button.click(
        fn=lambda img_bytes, meta_bytes, img_name, meta_name: (
            gr.File(value=io.BytesIO(img_bytes), label=img_name, visible=True),
            gr.File(value=io.BytesIO(meta_bytes), label=meta_name, visible=True)
        ),
        inputs=[hidden_img_file, hidden_meta_file, hidden_img_file, hidden_meta_file],
        outputs=[hidden_img_file, hidden_meta_file]
    )

# ======================
# 启动
# ======================
if __name__ == "__main__":
    demo.launch()