"""
🎨 Klein LoRA Studio — restyle a photo into a named look, base vs your LoRA.

Pick a named look (Scandinavian, Cottagecore, Watercolor…) and see a before→after
restyle. Optionally load your own trained `.safetensors` LoRA so the look becomes
a consistent *signature* that's truly yours. Image -> Image on FLUX.2 [klein] 4B.
Build Small (Backyard AI). Built on the klein starter's ZeroGPU + LoRA pattern.
"""
from __future__ import annotations

import os
import random
import time

# --- ZeroGPU shim: import `spaces` BEFORE torch -----------------------------
try:
    import spaces  # type: ignore

    GPU = spaces.GPU
except Exception:  # local / non-ZeroGPU fallback

    def GPU(*dargs, **dkwargs):  # noqa: N802
        if len(dargs) == 1 and callable(dargs[0]) and not dkwargs:
            return dargs[0]

        def wrap(fn):
            return fn

        return wrap


import gradio as gr
import torch
from diffusers import Flux2KleinPipeline
from PIL import Image

MODEL_ID = "black-forest-labs/FLUX.2-klein-4B"  # 4B, Apache 2.0, ungated
STEPS = 4
GUIDANCE = 1.0
MAX_SEED = 2**31 - 1

pipe = None
LOAD_ERR = ""
try:
    print(f"Loading {MODEL_ID} on CPU…")
    pipe = Flux2KleinPipeline.from_pretrained(MODEL_ID, torch_dtype=torch.bfloat16)
    print("  loaded.")
except Exception as e:  # noqa: BLE001
    LOAD_ERR = str(e)
    print("Model load failed:", e)


def klein_size(w: int, h: int, target_area: int = 1024 * 1024, divisor: int = 16):
    aspect = w / h
    nh = int((target_area / aspect) ** 0.5)
    nw = int(nh * aspect)
    nw = max(divisor, (nw // divisor) * divisor)
    nh = max(divisor, (nh // divisor) * divisor)
    return nw, nh


LOOKS = {
    "Scandinavian": "Restyle as a clean Scandinavian interior, pale wood, white "
                    "walls, muted neutral palette, soft daylight",
    "Cottagecore": "Restyle as a cozy cottagecore scene, warm florals, vintage "
                   "textiles, soft golden light",
    "Mid-century modern": "Restyle as a mid-century modern interior, walnut tones, "
                          "warm accent colours, retro furniture",
    "Cyberpunk neon": "Restyle with cyberpunk neon lighting, moody atmosphere, "
                      "magenta and cyan glow, rain-slick reflections",
    "Watercolor": "Turn it into a soft watercolor painting, gentle washes, visible "
                  "paper texture",
    "Risograph": "Restyle as a two-colour risograph print, grainy texture, bold "
                 "flat inks",
}


def _apply_lora(p, path: str, scale: float) -> None:
    try:
        p.unload_lora_weights()
    except Exception:
        pass
    p.load_lora_weights(path, adapter_name="user")
    p.set_adapters(["user"], adapter_weights=[float(scale)])


def _unload(p) -> None:
    try:
        p.unload_lora_weights()
    except Exception:
        pass


@GPU(duration=120)
def restyle(input_image: Image.Image | None, look_key: str, lora_file, scale: float):
    if pipe is None:
        raise gr.Error(f"Model isn't loaded (this Space needs a GPU). {LOAD_ERR[:200]}")
    if input_image is None:
        raise gr.Error("Upload a photo first (or pick an example).")
    pipe.to("cuda")
    img = input_image.convert("RGB")
    w, h = klein_size(*img.size)
    if img.size != (w, h):
        img = img.resize((w, h), Image.LANCZOS)
    prompt = LOOKS.get(look_key, next(iter(LOOKS.values())))
    seed = random.randint(0, MAX_SEED)
    t = time.time()
    try:
        _unload(pipe)
        if lora_file is not None:
            path = lora_file if isinstance(lora_file, str) else lora_file.name
            _apply_lora(pipe, path, scale)
            tag = f"+ LoRA {os.path.basename(path)} @ {scale}"
        else:
            tag = "base klein 4B"
        out = pipe(
            prompt=prompt,
            image=img,  # keyword — `image` is positional-first
            width=w,
            height=h,
            num_inference_steps=STEPS,
            guidance_scale=GUIDANCE,
            generator=torch.Generator(device="cuda").manual_seed(seed),
        ).images[0]
    finally:
        _unload(pipe)  # leave the shared pipe clean
    return img, out, f"**{look_key}** · {tag} · seed {seed} · {time.time() - t:.1f}s"


THEME = gr.themes.Soft(
    font=["system-ui", "-apple-system", "Segoe UI", "Roboto", "Helvetica", "Arial", "sans-serif"],
    font_mono=["ui-monospace", "SFMono-Regular", "Consolas", "monospace"],
)
CSS = """
footer {visibility: hidden;}
.gradio-container, .gradio-container .prose, .gradio-container p,
.gradio-container h1, .gradio-container h2, .gradio-container h3 {
  font-family: system-ui, -apple-system, "Segoe UI", Roboto, Helvetica, Arial, sans-serif !important;
}
"""

with gr.Blocks(title="Klein LoRA Studio", theme=THEME, css=CSS) as demo:
    gr.Markdown(
        "# 🎨 Klein LoRA Studio\n"
        "Drop in a photo, pick a **named look**, and see a **before → after** "
        "restyle. Want the look to be *yours* and consistent across a whole set? "
        "Load a `.safetensors` LoRA you trained and it becomes your **signature** "
        "style. Powered by **FLUX.2 [klein] 4B** (4B params, Apache 2.0)."
    )
    with gr.Row():
        with gr.Column():
            in_img = gr.Image(type="pil", label="Your photo", height=300)
            look = gr.Dropdown(list(LOOKS), value="Scandinavian", label="Look")
            lora = gr.File(label="Signature LoRA (optional .safetensors)", file_types=[".safetensors"])
            scale = gr.Slider(0.0, 1.5, value=1.0, step=0.05, label="LoRA strength")
            btn = gr.Button("🎨 Restyle", variant="primary")
        with gr.Column():
            with gr.Row():
                before = gr.Image(label="Before", height=320)
                after = gr.Image(label="After", height=320)
            info = gr.Markdown()
    btn.click(restyle, [in_img, look, lora, scale], [before, after, info])

if __name__ == "__main__":
    # ssr_mode=False: Gradio-5 SSR commonly renders unstyled raw HTML on Spaces.
    demo.queue(max_size=8).launch(
        server_name="0.0.0.0", server_port=7860, show_error=True, ssr_mode=False
    )