"""
🛍️ Shopfront — Product Photo Studio.

Upload a plain phone snap of a product; klein restages it on a clean, well-lit
scene while keeping the product itself intact, and returns a grid of variations
to pick from. Image -> Image on FLUX.2 [klein] 4B. Build Small (Backyard AI).
Built on the klein starter's verified ZeroGPU + pipeline 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
VARIANTS = 4

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):
    """Snap (w, h) to multiples of 16 under klein's 4096-patch ceiling."""
    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


# Each scene is an edit instruction: describe the *change* (background, surface,
# light), not the product — klein keeps the subject and restages around it.
SCENES = {
    "White marble": "a clean white marble surface, soft bright daylight, minimal "
                    "studio background, gentle reflection, professional product photo",
    "Linen flat-lay": "a top-down flat-lay on natural linen fabric, soft diffused "
                      "light, a few tasteful props, professional product photography",
    "Sunlit windowsill": "a sunlit wooden windowsill, warm morning light, soft "
                         "natural shadows, cozy lifestyle product photo",
    "Studio grey": "a seamless soft grey studio backdrop, even softbox lighting, "
                   "subtle reflection, clean e-commerce style",
    "Botanical": "soft green foliage and fresh natural light, botanical setting, "
                 "professional product photo",
}
# Prepended to every scene so the product stays recognizable.
GUARD = ("Keep the product itself unchanged and recognizable; change only the "
         "background, surface and lighting. Restage it on ")

_EX = os.path.join(os.path.dirname(__file__), "examples")
EXAMPLES = [os.path.join(_EX, f) for f in ("latte.jpg", "room.jpg", "street.jpg")
            if os.path.exists(os.path.join(_EX, f))]


@GPU(duration=120)
def stage(input_image: Image.Image | None, scene_key: str):
    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 product 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 = GUARD + SCENES.get(scene_key, next(iter(SCENES.values())))
    out, t = [], time.time()
    for _ in range(VARIANTS):
        seed = random.randint(0, MAX_SEED)
        # Footgun: pass prompt and image by keyword (`image` is positional-first).
        res = pipe(
            prompt=prompt,
            image=img,
            width=w,
            height=h,
            num_inference_steps=STEPS,
            guidance_scale=GUIDANCE,
            generator=torch.Generator(device="cuda").manual_seed(seed),
        ).images[0]
        out.append(res)
    return out, f"{VARIANTS} variations · {scene_key} · klein 4B · {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="Shopfront — Product Photo Studio", theme=THEME, css=CSS) as demo:
    gr.Markdown(
        "# 🛍️ Shopfront — Product Photo Studio\n"
        "Selling something handmade? Upload a plain phone photo of your product and "
        "**Shopfront** restages it on a clean, well-lit scene — keeping the product "
        "itself intact — then hands you **four variations** to choose from. No prompt "
        "writing, no studio. 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 product photo", height=320)
            scene = gr.Dropdown(list(SCENES), value="White marble", label="Scene")
            btn = gr.Button("📸 Restage it", variant="primary")
            if EXAMPLES:
                gr.Examples(EXAMPLES, in_img, label="No product handy? Try a photo")
        with gr.Column():
            out = gr.Gallery(label="Pick your favourite", columns=2, height=420, object_fit="contain")
            info = gr.Markdown()
    btn.click(stage, [in_img, scene], [out, 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
    )