"""
Virtual Try-On API — diffusion-based, backed by the CatVTON Space.

Exposes two surfaces:
  - a Gradio UI for manual testing
  - api_name="try_on", a base64-in / JSON-out endpoint the web backend calls

Why this Space is a bridge rather than a model host: CatVTON's own Space already
runs the weights on ZeroGPU. Calling it needs gradio_client — its file-upload
handshake is not reachable over plain HTTP (/gradio_api/upload returns 404), so
the Node backend cannot talk to it directly. This Space owns that handshake and
the HF token, and keeps the simple base64 contract the backend was built against.
"""

import base64
import inspect
import io
import json
import os
import tempfile

import gradio as gr
import spaces
from gradio_client import Client, handle_file
from PIL import Image

VTON_SPACE = os.environ.get("VTON_SPACE", "zhengchong/CatVTON")
HF_TOKEN = os.environ.get("HF_TOKEN")

CLOTH_TYPES = ["upper", "lower", "overall"]

_client = None


def get_client() -> Client:
    global _client
    if _client is None:
        # gradio_client renamed this kwarg from hf_token to token in 2.0; the
        # Space and local dev are pinned to different majors.
        params = inspect.signature(Client.__init__).parameters
        kwarg = "token" if "token" in params else "hf_token"
        _client = Client(VTON_SPACE, **{kwarg: HF_TOKEN})
    return _client


def to_b64(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=90)
    return base64.b64encode(buf.getvalue()).decode()


def _write_temp(img: Image.Image) -> str:
    path = tempfile.mktemp(suffix=".png")
    img.convert("RGB").save(path)
    return path


def try_on(
    person_img: Image.Image,
    garment_img: Image.Image,
    cloth_type: str = "upper",
    steps: int = 30,
    guidance: float = 2.5,
    seed: int = 42,
) -> Image.Image:
    person_path = _write_temp(person_img)
    garment_path = _write_temp(garment_img)

    # CatVTON's handler unconditionally reads person_image["layers"][0] as the
    # hand-drawn mask, so the layer has to exist. A fully transparent one is how
    # its UI represents "nothing drawn", which is what selects automasking.
    mask_path = tempfile.mktemp(suffix=".png")
    Image.new("RGBA", person_img.size, (0, 0, 0, 0)).save(mask_path)

    result = get_client().predict(
        person_image={
            "background": handle_file(person_path),
            "layers": [handle_file(mask_path)],
            "composite": handle_file(person_path),
        },
        cloth_image=handle_file(garment_path),
        cloth_type=cloth_type if cloth_type in CLOTH_TYPES else "upper",
        num_inference_steps=steps,
        guidance_scale=guidance,
        seed=seed,
        show_type="result only",
        api_name="/submit_function",
    )

    # predict() hands back a local path, or a dict when the output is a gallery item.
    if isinstance(result, dict):
        result = result.get("path") or result.get("value")
    if isinstance(result, (list, tuple)):
        result = result[0]

    return Image.open(result).convert("RGB")


@spaces.GPU(duration=1)
def _zerogpu_startup_probe():
    """Unused. ZeroGPU refuses to boot a Space with no @spaces.GPU function."""
    return None


def try_on_api(person_b64: str, garment_b64: str, cloth_type: str) -> str:
    """base64 JPEG in, JSON string out: {"output": b64} or {"error": msg}"""
    try:
        person = Image.open(io.BytesIO(base64.b64decode(person_b64)))
        garment = Image.open(io.BytesIO(base64.b64decode(garment_b64)))
        result = try_on(person, garment, cloth_type or "upper")
        return json.dumps({"output": to_b64(result)})
    except Exception as e:
        return json.dumps({"error": f"{type(e).__name__}: {e}"})


with gr.Blocks() as demo:
    gr.Markdown("# Virtual Try-On API")
    gr.Markdown(f"Diffusion try-on via `{VTON_SPACE}`.")

    with gr.Row():
        person_input = gr.Image(label="Person Photo", type="pil")
        garment_input = gr.Image(label="Garment Photo", type="pil")

    cloth_type_input = gr.Radio(
        CLOTH_TYPES, value="upper", label="Garment type"
    )
    output_image = gr.Image(label="Try-On Result")
    submit_btn = gr.Button("Generate Try-On", variant="primary")

    submit_btn.click(
        fn=lambda p, g, t: try_on(p, g, t) if p is not None and g is not None else None,
        inputs=[person_input, garment_input, cloth_type_input],
        outputs=output_image,
        api_name=False,
    )

    # Programmatic endpoint used by the web backend.
    api_person = gr.Textbox(visible=False)
    api_garment = gr.Textbox(visible=False)
    api_cloth_type = gr.Textbox(visible=False, value="upper")
    api_result = gr.Textbox(visible=False)
    api_trigger = gr.Button(visible=False)

    api_trigger.click(
        fn=try_on_api,
        inputs=[api_person, api_garment, api_cloth_type],
        outputs=api_result,
        api_name="try_on",
    )

if __name__ == "__main__":
    demo.launch()