import os
import random

import gradio as gr
import numpy as np
import spaces
import torch
from diffusers import Flux2KleinPipeline
from PIL import Image

# ──────────────────────────────────────────────────────────────────────────────
# Smart Character Swap — FLUX.2 [klein]
#
#   Identity source → the face / character to bring in
#   Target scene    → the photo whose pose, lighting, occlusions and color grade
#                     should be kept while the identity is swapped in
#
# Occlusion-aware, lighting-matched character / face swap, trained on FLUX.2 [klein].
# ──────────────────────────────────────────────────────────────────────────────

MAX_SEED = np.iinfo(np.int32).max
dtype = torch.bfloat16
device = "cuda" if torch.cuda.is_available() else "cpu"

BASE_MODEL = "black-forest-labs/FLUX.2-klein-9B"
LORA_REPO = "nhathoangfoto/Flux.2-Klein-9B-SmartCharacterSwap"
LORA_WEIGHTS = "Klein2-9B-SmartCharacterSwap.safetensors"
TRIGGER = "jhuangswap"

DEFAULT_PROMPT = (
    "jhuangswap, masterpiece, high-end photography, realistic portrait, "
    "matching target lighting, highly detailed skin texture, 8k resolution"
)

print("Loading FLUX.2 [klein] 9B...")
pipe = Flux2KleinPipeline.from_pretrained(BASE_MODEL, torch_dtype=dtype)
pipe.to("cuda")
pipe.load_lora_weights(LORA_REPO, weight_name=LORA_WEIGHTS, adapter_name="swap")
print("Pipeline ready.")


def _fit(img, target=1024, mult=16):
    """Resize so the longest side ~= target, snapped to a multiple of `mult`."""
    img = img.convert("RGB")
    w, h = img.size
    scale = target / max(w, h)
    nw = max(mult, int(round(w * scale / mult)) * mult)
    nh = max(mult, int(round(h * scale / mult)) * mult)
    return img.resize((nw, nh), Image.LANCZOS)


@spaces.GPU(duration=120)
def swap(identity, scene, prompt, lora_scale, steps, guidance,
         seed, randomize_seed, progress=gr.Progress(track_tqdm=True)):
    if identity is None:
        raise gr.Error("Please upload an identity source (the face/character to bring in).")
    if scene is None:
        raise gr.Error("Please upload a target scene (the photo to swap the identity into).")

    pipe.set_adapters(["swap"], adapter_weights=[lora_scale])

    scene_f = _fit(scene, target=1024)
    identity_f = _fit(identity, target=1024)

    full_prompt = prompt.strip() if prompt and prompt.strip() else DEFAULT_PROMPT
    if TRIGGER not in full_prompt:
        full_prompt = f"{TRIGGER}, {full_prompt}"

    if randomize_seed:
        seed = random.randint(0, MAX_SEED)
    generator = torch.Generator(device=device).manual_seed(int(seed))

    # Output keeps the target scene's composition; identity is the reference.
    result = pipe(
        image=[scene_f, identity_f],
        prompt=full_prompt,
        height=scene_f.height,
        width=scene_f.width,
        num_inference_steps=int(steps),
        guidance_scale=guidance,
        generator=generator,
    ).images[0]

    return (scene_f, result), seed


css = """
#col { max-width: 1200px; margin: 0 auto; }
.header { text-align:center; padding: 8px 0 4px; }
.header h1 { font-size: 1.7rem; margin: 0; font-weight: 700; }
.header p  { color: var(--body-text-color-subdued); margin: 4px 0 0; font-size: 0.95rem; }
.header a  { color: #6366f1; text-decoration: none; font-weight: 600; }
"""

with gr.Blocks(css=css) as demo:
    with gr.Column(elem_id="col"):
        gr.HTML(
            """
            <div class="header">
              <h1>🎭 Smart Character Swap · FLUX.2 [klein]</h1>
              <p>Swap an <b>identity</b> into a <b>target scene</b> — occlusion-aware, with the scene's
              own lighting and color grade preserved.
              &nbsp;·&nbsp;
              <a href="https://huggingface.co/nhathoangfoto/Flux.2-Klein-9B-SmartCharacterSwap" target="_blank">SmartCharacterSwap LoRA by nhathoangfoto</a></p>
            </div>
            """
        )

        with gr.Row(equal_height=False):
            with gr.Column():
                with gr.Row():
                    identity = gr.Image(label="Identity source (face to bring in)", type="pil", height=300)
                    scene = gr.Image(label="Target scene (pose / lighting to keep)", type="pil", height=300)

                prompt = gr.Textbox(
                    label="Prompt",
                    value=DEFAULT_PROMPT,
                    lines=2,
                    info="Trigger 'jhuangswap' is added automatically if you remove it.",
                )

                with gr.Accordion("Advanced", open=False):
                    lora_scale = gr.Slider(0.5, 1.2, value=0.9, step=0.05, label="LoRA strength")
                    steps = gr.Slider(4, 30, value=8, step=1, label="Steps")
                    guidance = gr.Slider(1.0, 10.0, value=4.0, step=0.1, label="Guidance scale")
                    with gr.Row():
                        seed = gr.Slider(0, MAX_SEED, value=0, step=1, label="Seed")
                        randomize_seed = gr.Checkbox(value=True, label="Randomize")

                run_btn = gr.Button("Swap", variant="primary", size="lg")

            with gr.Column():
                result = gr.ImageSlider(label="Target scene → Swapped result", type="pil", height=480)

        gr.Examples(
            examples=[
                ["examples/identity_B.jpg", "examples/scene_B.jpg"],
                ["examples/identity_A.jpg", "examples/scene_A.jpg"],
            ],
            inputs=[identity, scene],
            outputs=[result, seed],
            fn=lambda i, s: swap(i, s, DEFAULT_PROMPT, 0.9, 8, 4.0, 0, True),
            cache_examples=True,
            cache_mode="lazy",
            label="Identity + scene examples",
        )

    swap_inputs = [identity, scene, prompt, lora_scale, steps, guidance, seed, randomize_seed]
    run_btn.click(fn=swap, inputs=swap_inputs, outputs=[result, seed])
    prompt.submit(fn=swap, inputs=swap_inputs, outputs=[result, seed])

if __name__ == "__main__":
    demo.launch(theme=gr.themes.Citrus(), show_error=True, ssr_mode=False)