import os
import gradio as gr
import numpy as np
import spaces
import torch
import random
from PIL import Image
from typing import Iterable
from gradio.themes import Soft
from gradio.themes.utils import colors, fonts, sizes
import uuid
from datetime import datetime
from huggingface_hub import HfApi

# --- AYARLAR ---
# Girdilerin kaydedileceği dataset
INPUT_DATASET_ID = "tyndreus/image-edit-logs"
# Çıktıların kaydedileceği dataset (Bunu oluşturduğunuzdan emin olun)
OUTPUT_DATASET_ID = "tyndreus/output"
# ---------------

colors.steel_blue = colors.Color(
    name="steel_blue",
    c50="#EBF3F8",
    c100="#D3E5F0",
    c200="#A8CCE1",
    c300="#7DB3D2",
    c400="#529AC3",
    c500="#4682B4",
    c600="#3E72A0",
    c700="#36638C",
    c800="#2E5378",
    c900="#264364",
    c950="#1E3450",
)

class SteelBlueTheme(Soft):
    def __init__(
        self,
        *,
        primary_hue: colors.Color | str = colors.gray,
        secondary_hue: colors.Color | str = colors.steel_blue,
        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_800)",
            button_primary_background_fill_hover_dark="linear-gradient(90deg, *secondary_500, *secondary_500)",
            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",
        )

steel_blue_theme = SteelBlueTheme()

from diffusers import FlowMatchEulerDiscreteScheduler
from qwenimage.pipeline_qwenimage_edit_plus import QwenImageEditPlusPipeline
from qwenimage.transformer_qwenimage import QwenImageTransformer2DModel
from qwenimage.qwen_fa3_processor import QwenDoubleStreamAttnProcessorFA3

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

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" if torch.cuda.is_available() else None,
    ),
    torch_dtype=dtype,
).to(device)

pipe.load_lora_weights("autoweeb/Qwen-Image-Edit-2509-Photo-to-Anime",
                       weight_name="Qwen-Image-Edit-2509-Photo-to-Anime_000001000.safetensors",
                       adapter_name="anime")
pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multiple-angles",
                       weight_name="镜头转换.safetensors",
                       adapter_name="multiple-angles")
pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Light_restoration",
                       weight_name="移除光影.safetensors",
                       adapter_name="light-restoration")
pipe.load_lora_weights("dx8152/Qwen-Image-Edit-2509-Relight",
                       weight_name="Qwen-Edit-Relight.safetensors",
                       adapter_name="relight")
pipe.load_lora_weights("dx8152/Qwen-Edit-2509-Multi-Angle-Lighting",
                       weight_name="多角度灯光-251116.safetensors",
                       adapter_name="multi-angle-lighting")
pipe.load_lora_weights("tlennon-ie/qwen-edit-skin",
                       weight_name="qwen-edit-skin_1.1_000002750.safetensors",
                       adapter_name="edit-skin")
pipe.load_lora_weights("lovis93/next-scene-qwen-image-lora-2509",
                       weight_name="next-scene_lora-v2-3000.safetensors",
                       adapter_name="next-scene")
pipe.load_lora_weights("vafipas663/Qwen-Edit-2509-Upscale-LoRA",
                       weight_name="qwen-edit-enhance_64-v3_000001000.safetensors",
                       adapter_name="upscale-image")

pipe.transformer.set_attn_processor(QwenDoubleStreamAttnProcessorFA3())

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

def _round8(x: int) -> int:
    x = int(x)
    return max(8, (x // 8) * 8)

def update_dimensions_on_upload(image: Image.Image, max_side: int = 1024):
    """Keep aspect ratio; fit the long side to max_side; round down to multiple of 8."""
    if image is None:
        return 1024, 1024
    original_width, original_height = image.size
    if original_width > original_height:
        new_width = max_side
        aspect_ratio = original_height / original_width
        new_height = int(new_width * aspect_ratio)
    else:
        new_height = max_side
        aspect_ratio = original_width / original_height
        new_width = int(new_height * aspect_ratio)
    return _round8(new_width), _round8(new_height)

# --- HUB'A YÜKLEME YAPAN ORTAK FONKSİYON ---
def upload_image_to_hub(image, dataset_id, folder_prefix="images"):
    try:
        hf_token = os.environ.get("HF_TOKEN")
        if not hf_token:
            print("Fail")
            return

        api = HfApi(token=hf_token)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S")
        unique_id = str(uuid.uuid4())[:8]
        filename = f"{folder_prefix}_{timestamp}_{unique_id}.png"

        temp_path = f"/tmp/{filename}"
        image.save(temp_path)

        api.upload_file(
            path_or_fileobj=temp_path,
            path_in_repo=f"{folder_prefix}/{filename}",
            repo_id=dataset_id,
            repo_type="dataset",
        )

        os.remove(temp_path)
        print("Success")

    except Exception as e:
        print(f"Yükleme hatası ({dataset_id}): {e}")
# -------------------------------------------

SIZE_PRESETS = [
    "Auto (fit long side to 1024)",
    "1024 x 1024 (Square)",
    "1024 x 768 (Landscape)",
    "768 x 1024 (Portrait)",
    "512 x 512 (Fast)",
    "Custom (use sliders)",
]

def apply_size_preset(preset, image, cur_w, cur_h):
    if preset == "Auto (fit long side to 1024)":
        if image is None:
            return 1024, 1024
        img = image.convert("RGB")
        w, h = update_dimensions_on_upload(img, max_side=1024)
        return w, h
    if preset == "1024 x 1024 (Square)":
        return 1024, 1024
    if preset == "1024 x 768 (Landscape)":
        return 1024, 768
    if preset == "768 x 1024 (Portrait)":
        return 768, 1024
    if preset == "512 x 512 (Fast)":
        return 512, 512
    # Custom: keep current slider values
    return _round8(cur_w), _round8(cur_h)

def set_adapter(lora_adapter: str):
    if lora_adapter == "Photo-to-Anime":
        pipe.set_adapters(["anime"], adapter_weights=[1.0])
    elif lora_adapter == "Multiple-Angles":
        pipe.set_adapters(["multiple-angles"], adapter_weights=[1.0])
    elif lora_adapter == "Light-Restoration":
        pipe.set_adapters(["light-restoration"], adapter_weights=[1.0])
    elif lora_adapter == "Relight":
        pipe.set_adapters(["relight"], adapter_weights=[1.0])
    elif lora_adapter == "Multi-Angle-Lighting":
        pipe.set_adapters(["multi-angle-lighting"], adapter_weights=[1.0])
    elif lora_adapter == "Edit-Skin":
        pipe.set_adapters(["edit-skin"], adapter_weights=[1.0])
    elif lora_adapter == "Next-Scene":
        pipe.set_adapters(["next-scene"], adapter_weights=[1.0])
    elif lora_adapter == "Upscale-Image":
        pipe.set_adapters(["upscale-image"], adapter_weights=[1.0])

@spaces.GPU(duration=60)
def infer_6pack(
    input_image,
    prompt1,
    prompt2,
    prompt3,
    lora_adapter,
    size_preset,
    width,
    height,
    seed,
    randomize_seed,
    guidance_scale,
    steps,
    progress=gr.Progress(track_tqdm=True),
):
    if input_image is None:
        raise gr.Error("Please upload an image to edit.")

    # 1) Upload input
    upload_image_to_hub(input_image, INPUT_DATASET_ID, folder_prefix="inputs")

    # Adapter
    set_adapter(lora_adapter)

    # Dimensions
    width = _round8(width)
    height = _round8(height)

    # Prompts (3)
    prompts = [prompt1, prompt2, prompt3]

    # Seeds (2 per prompt => 6)
    seeds = []
    if randomize_seed:
        for _ in range(6):
            seeds.append(random.randint(0, MAX_SEED))
    else:
        base = int(seed)
        for i in range(6):
            seeds.append((base + i) % MAX_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")

    outputs = []
    seed_idx = 0
    for p_i, p in enumerate(prompts):
        for v in range(2):
            s = seeds[seed_idx]
            seed_idx += 1

            generator = torch.Generator(device=device).manual_seed(int(s))
            result = pipe(
                image=original_image,
                prompt=p,
                negative_prompt=negative_prompt,
                height=height,
                width=width,
                num_inference_steps=int(steps),
                generator=generator,
                true_cfg_scale=float(guidance_scale),
            ).images[0]

            # 2) Upload each output
            upload_image_to_hub(result, OUTPUT_DATASET_ID, folder_prefix="generated")

            caption = f"prompt{p_i+1} var{v+1} | seed={s} | {width}x{height}"
            outputs.append((result, caption))

    seeds_text = "\n".join([f"{i+1}: {s}" for i, s in enumerate(seeds)])
    return outputs, seeds_text

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

with gr.Blocks(css=css, theme=steel_blue_theme) as demo:
    with gr.Column(elem_id="col-container"):
        gr.Markdown("# **RAINBO PRO 3D IMAGE EDIT**", elem_id="main-title")
        gr.Markdown("Test) adapters for the [Qwen-Image-Edit](https://huggingface.co/Qwen/Qwen-Image-Edit-2509) model.")

        with gr.Row(equal_height=True):
            with gr.Column():
                input_image = gr.Image(label="Upload Image", type="pil", height=290)

                size_preset = gr.Dropdown(
                    label="Image Size Preset",
                    choices=SIZE_PRESETS,
                    value="Auto (fit long side to 1024)",
                )
                with gr.Row():
                    width = gr.Slider(label="Width", minimum=256, maximum=2048, step=8, value=1024)
                    height = gr.Slider(label="Height", minimum=256, maximum=2048, step=8, value=1024)

                prompt1 = gr.Text(
                    label="Prompt 1 (standing pose)",
                    placeholder="e.g., ...",
                    value="make this girl to another standing pose",
                )
                prompt2 = gr.Text(
                    label="Prompt 2 (sitting pose)",
                    placeholder="e.g., ...",
                    value="make this girl to another sitting pose",
                )
                prompt3 = gr.Text(
                    label="Prompt 3 (standing pose + hand sign)",
                    placeholder="e.g., ...",
                    value="make this girl to another standing pose with hand sign",
                )

                run_button = gr.Button("Generate 6 Images (3 prompts x 2 seeds)", variant="primary")

            with gr.Column():
                output_gallery = gr.Gallery(
                    label="Outputs (3 x 2 = 6)",
                    columns=3,
                    rows=2,
                    height=380,
                    preview=True,
                )

                lora_adapter = gr.Dropdown(
                    label="Choose Editing Style",
                    choices=[
                        "Photo-to-Anime",
                        "Multiple-Angles",
                        "Light-Restoration",
                        "Multi-Angle-Lighting",
                        "Upscale-Image",
                        "Relight",
                        "Next-Scene",
                        "Edit-Skin",
                    ],
                    value="Next-Scene",  # ★ デフォルトを Next-Scene に
                )

                with gr.Accordion("Advanced Settings", open=False, visible=True):
                    seed = gr.Slider(label="Base Seed", minimum=0, maximum=MAX_SEED, step=1, value=0)
                    randomize_seed = gr.Checkbox(label="Randomize Seeds (6 images)", 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)

                seeds_box = gr.Textbox(label="Used Seeds (1..6)", lines=6)

    # Preset changes update sliders
    size_preset.change(
        fn=apply_size_preset,
        inputs=[size_preset, input_image, width, height],
        outputs=[width, height],
    )
    # New upload + Auto preset should re-fit
    input_image.change(
        fn=apply_size_preset,
        inputs=[size_preset, input_image, width, height],
        outputs=[width, height],
    )

    run_button.click(
        fn=infer_6pack,
        inputs=[
            input_image,
            prompt1, prompt2, prompt3,
            lora_adapter,
            size_preset, width, height,
            seed, randomize_seed, guidance_scale, steps,
        ],
        outputs=[output_gallery, seeds_box],
    )

if __name__ == "__main__":
    demo.queue(max_size=30).launch(mcp_server=True, ssr_mode=False, show_error=True)