import os
import sys
import subprocess
from pathlib import Path

import gradio as gr
import spaces
import torch
from PIL import Image
from diffusers.image_processor import VaeImageProcessor
from huggingface_hub import snapshot_download

# ============================================================
# CatVTON ZeroGPU — stable Gradio app
# Based on the user's ComfyUI workflow:
#
# LoadCatVTONPipeline_AMD
# Person image       -> target_image
# Garment image      -> refer_image
# Clothing mask      -> mask_image
# ImageInvert        -> represented by the mask polarity note below
# CatVTON_AMD        -> result
#
# IMPORTANT:
# The uploaded UI mask should be:
#   WHITE = clothing / replace area
#   BLACK = background / keep area
# Therefore this version does NOT invert the uploaded mask.
# ============================================================

ROOT = Path("/tmp/catvton_src")
ROOT.mkdir(parents=True, exist_ok=True)

BASE_MODEL = os.getenv(
    "CATVTON_BASE_MODEL",
    "booksforcharlie/stable-diffusion-inpainting",
)
CATVTON_MODEL = os.getenv(
    "CATVTON_MODEL",
    "zhengchong/CatVTON",
)

WIDTH = int(os.getenv("CATVTON_WIDTH", "768"))
HEIGHT = int(os.getenv("CATVTON_HEIGHT", "1024"))


def ensure_catvton_source():
    marker = ROOT / "model" / "pipeline.py"

    if marker.exists():
        return

    if (ROOT / ".git").exists():
        subprocess.run(
            ["git", "-C", str(ROOT), "pull", "--ff-only"],
            check=False,
            stdout=subprocess.DEVNULL,
            stderr=subprocess.DEVNULL,
        )
        return

    subprocess.run(
        [
            "git",
            "clone",
            "--depth",
            "1",
            "--branch",
            "edited",
            "https://github.com/Zheng-Chong/CatVTON.git",
            str(ROOT),
        ],
        check=True,
    )


ensure_catvton_source()
sys.path.insert(0, str(ROOT))

from model.pipeline import CatVTONPipeline
from utils import resize_and_crop, resize_and_padding


# ============================================================
# MODEL
# ============================================================

print("Downloading CatVTON checkpoint...")
catvton_repo = snapshot_download(repo_id=CATVTON_MODEL)

print("Loading CatVTON pipeline...")

pipe = CatVTONPipeline(
    base_ckpt=BASE_MODEL,
    attn_ckpt=catvton_repo,
    attn_ckpt_version="mix",
    weight_dtype=torch.float16,
    use_tf32=True,
    device="cuda",
    # The upstream safety checker can produce false positives on
    # normal fully-clothed fashion images. The user's successful
    # test required this to be disabled.
    skip_safety_check=True,
)

mask_processor = VaeImageProcessor(
    vae_scale_factor=8,
    do_normalize=False,
    do_binarize=True,
    do_convert_grayscale=True,
)


# ============================================================
# IMAGE HELPERS
# ============================================================

def load_rgb(value):
    if value is None:
        return None

    if isinstance(value, Image.Image):
        return value.convert("RGB")

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


def load_mask(value):
    if value is None:
        return None

    if isinstance(value, Image.Image):
        return value.convert("L")

    return Image.open(value).convert("L")


def prepare_inputs(person_file, garment_file, mask_file):
    person = load_rgb(person_file)
    garment = load_rgb(garment_file)
    mask = load_mask(mask_file)

    # Same target/reference preprocessing used by the working app.
    person = resize_and_crop(person, (WIDTH, HEIGHT))
    garment = resize_and_padding(garment, (WIDTH, HEIGHT))

    # IMPORTANT:
    # The UI mask is already:
    #   white = clothing/replacement region
    #   black = preserved region
    #
    # Do not invert it again.
    mask = resize_and_crop(mask, (WIDTH, HEIGHT))

    # Light mask edge smoothing for cleaner garment boundaries.
    mask = mask_processor.blur(mask, blur_factor=9)

    return person, garment, mask


# ============================================================
# GENERATION
# ============================================================

@spaces.GPU(duration=120)
def generate_outfit(person_file, garment_file, mask_file, steps, cfg, seed):
    if person_file is None:
        raise gr.Error("Please upload the person image.")

    if garment_file is None:
        raise gr.Error("Please upload the garment/clothing image.")

    if mask_file is None:
        raise gr.Error(
            "Please upload the clothing mask. "
            "White = clothing area, Black = background."
        )

    person, garment, mask = prepare_inputs(
        person_file,
        garment_file,
        mask_file,
    )

    steps = int(steps)
    cfg = float(cfg)
    seed = int(seed)

    generator = None

    # -1 means random seed.
    if seed >= 0:
        generator = torch.Generator(device="cuda").manual_seed(seed)

    print(
        f"CatVTON inference | steps={steps} "
        f"cfg={cfg} seed={seed}"
    )

    with torch.inference_mode():
        result = pipe(
            image=person,
            condition_image=garment,
            mask=mask,
            num_inference_steps=steps,
            guidance_scale=cfg,
            height=HEIGHT,
            width=WIDTH,
            generator=generator,
        )[0]

    return result


# ============================================================
# UI
# ============================================================

DESCRIPTION = """
# 🐈 CatVTON Outfit Try-On

Upload a **person photo**, a **garment/product photo**, and a **clothing mask**.

**Mask format:** ⚪ White = replace clothing area · ⚫ Black = preserve area

The default settings match the tested workflow: **30 steps / CFG 2.5**.
"""

with gr.Blocks(title="CatVTON Outfit Try-On — ZeroGPU") as demo:

    gr.Markdown(DESCRIPTION)

    with gr.Row():

        # ----------------------------------------------------
        # INPUTS
        # ----------------------------------------------------
        with gr.Column():

            person = gr.Image(
                type="filepath",
                label="👤 Person Image",
            )

            garment = gr.Image(
                type="filepath",
                label="👕 Garment / Clothing Image",
            )

            mask = gr.Image(
                type="filepath",
                label="🎭 Clothing Mask",
            )

            gr.Markdown(
                """
**Mask example**

- White shirt shape = area CatVTON replaces
- Black background = area CatVTON preserves
"""
            )

        # ----------------------------------------------------
        # OUTPUT
        # ----------------------------------------------------
        with gr.Column():

            result = gr.Image(
                type="pil",
                label="✨ Generated Result",
            )

            gr.Markdown(
                "Use the image toolbar to download the generated result."
            )

    # --------------------------------------------------------
    # SETTINGS
    # --------------------------------------------------------

    with gr.Row():

        steps = gr.Slider(
            minimum=10,
            maximum=50,
            step=1,
            value=30,
            label="Inference Steps",
        )

        cfg = gr.Slider(
            minimum=1.0,
            maximum=5.0,
            step=0.1,
            value=2.5,
            label="CFG",
        )

        seed = gr.Number(
            value=-1,
            precision=0,
            label="Seed (-1 = Random)",
        )

    # --------------------------------------------------------
    # BUTTON
    # --------------------------------------------------------

    generate_button = gr.Button(
        "🚀 GENERATE OUTFIT",
        variant="primary",
        size="lg",
    )

    generate_button.click(
        fn=generate_outfit,
        inputs=[
            person,
            garment,
            mask,
            steps,
            cfg,
            seed,
        ],
        outputs=result,
    )

    gr.Markdown(
        """
### Workflow mapping

`Person → target_image`  
`Garment → refer_image`  
`White/black mask → mask_image`  
`CatVTON → Generated Result`

**Model:** `zhengchong/CatVTON`  
**Resolution:** 768 × 1024  
**Default:** 30 steps / CFG 2.5
"""
    )


if __name__ == "__main__":
    demo.queue(max_size=8).launch()