"""Gradio demo for PerceptionDLM parallel region captioning.

This app runs the PerceptionDLM model on ZeroGPU. Users upload an image and
one or more binary masks, and the model generates captions for all masked
regions in parallel via a single denoising process. The decoding animation
replays each diffusion step so you can watch captions emerge token by token.
"""

import os

os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import html as html_lib
import time
from typing import Dict, List, Tuple

import spaces  # MUST come before torch / any CUDA-touching import
import torch
import numpy as np
from PIL import Image

import gradio as gr
from transformers import AutoModel, AutoProcessor
from huggingface_hub import snapshot_download
import json

# ---------------------------------------------------------------------------
# Model loading at module scope (ZeroGPU intercepts .to("cuda"))
# ---------------------------------------------------------------------------
MODEL_ID = "MSALab/PerceptionDLM"
DTYPE = torch.bfloat16

# The model's config.json references "bitersun/LLaDA-8B-Instruct-HF" as the
# _name_or_path for language_model_config, but that repo no longer exists on
# the Hub.  All the needed remote-code files (configuration_llada.py,
# modeling_llada.py, etc.) live in MSALab/PerceptionDLM itself.  We patch the
# auto_map in language_model_config to prefix the repo ID using the "--"
# syntax, so get_class_from_dynamic_module downloads code from the correct
# repo while keeping _name_or_path intact (the "llada" check in
# modeling_pdmllm.py relies on it).
print(f"Downloading model from {MODEL_ID} ...")
_local_model_dir = snapshot_download(
    repo_id=MODEL_ID,
    repo_type="model",
)

# Patch config.json: prefix auto_map values with the correct repo ID
_config_path = os.path.join(_local_model_dir, "config.json")
with open(_config_path) as f:
    _config_dict = json.load(f)

_lm_cfg = _config_dict.get("language_model_config", {})
_lm_auto_map = _lm_cfg.get("auto_map", {})
_patched = False
for key, val in _lm_auto_map.items():
    if not val.startswith(MODEL_ID + "--"):
        _lm_auto_map[key] = f"{MODEL_ID}--{val}"
        _patched = True
if _patched:
    with open(_config_path, "w") as f:
        json.dump(_config_dict, f, indent=2)
    print(f"Patched language_model_config.auto_map -> prefix {MODEL_ID}--")

print(f"Loading processor from {_local_model_dir} ...")
PROCESSOR = AutoProcessor.from_pretrained(_local_model_dir, trust_remote_code=True)
TOKENIZER = PROCESSOR.tokenizer

print(f"Loading model from {_local_model_dir} ...")
MODEL = AutoModel.from_pretrained(
    _local_model_dir,
    torch_dtype=DTYPE,
    trust_remote_code=True,
    attn_implementation="sdpa",
)
MODEL.processor = PROCESSOR
MODEL.to("cuda")
MODEL.eval()
print("Model loaded.")

# ---------------------------------------------------------------------------
# SAM 3 (mask-generation tool) — interactive click-to-mask.
# This is an ADDITIONAL tool for producing the binary masks PerceptionDLM
# consumes; it does NOT replace the perception model. We use Meta's standalone
# `sam3` package (independent of transformers, so it coexists with the pinned
# transformers version PerceptionDLM requires). The interactive image
# predictor does single-image promptable segmentation: click point(s) / draw a
# box -> one binary mask. The facebook/sam3 weights are gated, so an HF_TOKEN
# with access must be available in the environment.
#
# The model is built lazily on first use inside the GPU worker: `build_tracker`
# imports triton at construction time, which is only meaningful on the GPU
# worker under ZeroGPU.
# ---------------------------------------------------------------------------
print("Preparing SAM 3 (lazy, built on first GPU call) ...")
_SAM3_PREDICTOR = None


def _get_sam3_predictor():
    """Build (once) and return the SAM 3 interactive image predictor.

    We build the SAM 3 tracker with its own vision backbone and load weights
    directly from the gated facebook/sam3 checkpoint: the tracker's own
    parameters come from the ``tracker.*`` keys, and the shared SAM 3 vision
    backbone comes from the ``detector.backbone.*`` keys. (The package's
    ``build_sam3_image_model`` builds the tracker without a backbone, which
    works for its concept/video paths but leaves the single-image interactive
    predictor without image features.)
    """
    global _SAM3_PREDICTOR
    if _SAM3_PREDICTOR is None:
        from sam3 import model_builder as _mb
        from sam3.model.sam1_task_predictor import SAM3InteractiveImagePredictor
        from huggingface_hub import hf_hub_download

        tracker = _mb.build_tracker(
            apply_temporal_disambiguation=False, with_backbone=True
        )
        ckpt = torch.load(
            hf_hub_download(repo_id="facebook/sam3", filename="sam3.pt"),
            map_location="cpu",
            weights_only=True,
        )
        if "model" in ckpt and isinstance(ckpt["model"], dict):
            ckpt = ckpt["model"]
        state = {
            k[len("tracker."):]: v
            for k, v in ckpt.items()
            if k.startswith("tracker.")
        }
        for k, v in ckpt.items():
            if k.startswith("detector.backbone."):
                state["backbone." + k[len("detector.backbone."):]] = v
        tracker.load_state_dict(state, strict=False)
        _SAM3_PREDICTOR = SAM3InteractiveImagePredictor(tracker.cuda().eval())
    return _SAM3_PREDICTOR

# ---------------------------------------------------------------------------
# Constants
# ---------------------------------------------------------------------------
MASK_ID = 126336  # token id for the LLaDA diffusion backbone
MASK_PLACEHOLDER = "\ue000"  # sentinel for not-yet-revealed tokens
DEFAULT_PROMPT = "Describe each masked region in detail."

OVERLAY_COLORS = [
    (239, 68, 68),    # red
    (16, 185, 129),   # green
    (59, 130, 246),   # blue
    (245, 158, 11),   # amber
    (236, 72, 153),   # pink
    (139, 92, 246),   # violet
    (6, 182, 212),    # cyan
    (132, 204, 22),   # lime
]

# ---------------------------------------------------------------------------
# Preprocessing helpers (adapted from demo/infer_pdmllm.py)
# ---------------------------------------------------------------------------

def find_closest_aspect_ratio(aspect_ratio, target_ratios, width, height, image_size):
    best_ratio_diff = float('inf')
    best_ratio = (1, 1)
    area = width * height
    for ratio in target_ratios:
        target_aspect_ratio = ratio[0] / ratio[1]
        ratio_diff = abs(aspect_ratio - target_aspect_ratio)
        if ratio_diff < best_ratio_diff:
            best_ratio_diff = ratio_diff
            best_ratio = ratio
        elif ratio_diff == best_ratio_diff:
            if area > 0.5 * image_size * image_size * ratio[0] * ratio[1]:
                best_ratio = ratio
    return best_ratio


def dynamic_preprocess(image, min_num=1, max_num=6, image_size=512, use_thumbnail=True):
    orig_width, orig_height = image.size
    aspect_ratio = orig_width / orig_height
    target_ratios = set(
        (i, j) for n in range(min_num, max_num + 1)
        for i in range(1, n + 1) for j in range(1, n + 1)
        if i * j <= max_num and i * j >= min_num
    )
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
    target_aspect_ratio = find_closest_aspect_ratio(
        aspect_ratio, target_ratios, orig_width, orig_height, image_size
    )
    target_width = image_size * target_aspect_ratio[0]
    target_height = image_size * target_aspect_ratio[1]
    blocks = target_aspect_ratio[0] * target_aspect_ratio[1]
    resized_img = image.resize((target_width, target_height))
    processed_images = []
    for i in range(blocks):
        box = (
            (i % (target_width // image_size)) * image_size,
            (i // (target_width // image_size)) * image_size,
            ((i % (target_width // image_size)) + 1) * image_size,
            ((i // (target_width // image_size)) + 1) * image_size,
        )
        split_img = resized_img.crop(box)
        processed_images.append(split_img)
    assert len(processed_images) == blocks
    if use_thumbnail and len(processed_images) != 1:
        thumbnail_img = image.resize((image_size, image_size))
        processed_images.append(thumbnail_img)
    return processed_images


def sort_masks_by_area(masks: List[np.ndarray]):
    areas = [np.sum(m) for m in masks]
    return np.argsort(np.array(areas))[::-1]


def build_visual_prompt_matrices(
    masks: List[np.ndarray],
    prompt_numbers: int,
) -> tuple:
    if len(masks) > prompt_numbers:
        raise ValueError(
            f"Number of masks ({len(masks)}) exceeds prompt_numbers ({prompt_numbers})."
        )
    height, width = masks[0].shape
    prompt_indexes = list(range(prompt_numbers))
    selected_prompt_indexes = prompt_indexes[:len(masks)]
    selected_prompt_tokens = [f"<Prompt{i}>" for i in selected_prompt_indexes]

    filled_matrices = []
    for prompt_id, mask in zip(selected_prompt_indexes, masks):
        filled_matrix = np.full((height, width), 255, dtype=np.uint8)
        fill_area = (filled_matrix == 255) & mask.astype(bool)
        filled_matrix[fill_area] = prompt_id
        filled_matrices.append(filled_matrix)

    visual_prompt_images = [Image.fromarray(m) for m in filled_matrices]
    return visual_prompt_images, selected_prompt_tokens, selected_prompt_indexes


def build_bboxes(masks: List[np.ndarray], tokenizer) -> Dict[str, tuple]:
    height, width = masks[0].shape
    bboxes: Dict[str, tuple] = {}
    for idx, mask in enumerate(masks):
        coords = np.argwhere(mask > 0)
        if coords.size == 0:
            continue
        y_min, x_min = coords.min(axis=0)
        y_max, x_max = coords.max(axis=0)
        token_id = tokenizer.convert_tokens_to_ids(f"<|reserved_token_{idx}|>")
        bboxes[str(token_id)] = (
            x_min / width,
            y_min / height,
            x_max / width,
            y_max / height,
        )
    return bboxes


def compute_aspect_ratio(image: Image.Image, processor, num_tiles: int) -> torch.Tensor:
    min_tiles = getattr(processor, "min_sub_img", 1)
    max_tiles = getattr(processor, "max_sub_img", 6)
    if hasattr(processor, "image_size"):
        image_size = processor.image_size[0] if isinstance(processor.image_size, tuple) else processor.image_size
    else:
        size = getattr(processor, "size", 512)
        if isinstance(size, dict):
            image_size = size.get("height", size.get("shortest_edge", 512))
        else:
            image_size = size
    aspect_ratio = image.width / image.height
    target_ratios = {
        (i, j)
        for n in range(min_tiles, max_tiles + 1)
        for i in range(1, n + 1)
        for j in range(1, n + 1)
        if min_tiles <= i * j <= max_tiles
    }
    target_ratios = sorted(target_ratios, key=lambda x: x[0] * x[1])
    grid_w, grid_h = find_closest_aspect_ratio(aspect_ratio, target_ratios, image.width, image.height, image_size)
    return torch.tensor([[grid_w, grid_h]], dtype=torch.int64)


def build_prompt_text(tokenizer, num_image_token: int, num_tiles: int, questions: List[str], gen_len: int, num_masks: int) -> str:
    img_ctx = "".join(["<IMG_CONTEXT>"] * (num_image_token * num_tiles))
    parts = ["system\nYou are a helpful assistant.\n"]
    parts.append("user\n")
    parts.append(
        f"<img>{img_ctx}</img>"
        + "\n".join([f"<|reserved_token_{i}|>" for i in range(num_masks)])
        + f"\n{questions[0]}\n"
    )
    parts.append("assistant\n")
    mask_seq = "<|mdm_mask|>" * gen_len
    parts.append("\n".join([f"<|Mask_Cap_{i}|>{mask_seq}" for i in range(num_masks)]))
    return "".join(parts) + ""


def split_assistant_blocks(text: str, num_masks: int) -> List[str]:
    blocks = text.split("assistant\n")
    assistant_text = blocks[-1].split("")[0] if len(blocks) > 1 else text
    captions = []
    for i in range(num_masks):
        start_tag = f"<|Mask_Cap_{i}|>"
        next_tag = f"<|Mask_Cap_{i + 1}|>"
        start_pos = assistant_text.find(start_tag)
        if start_pos == -1:
            captions.append("")
            continue
        content_start = start_pos + len(start_tag)
        end_pos = assistant_text.find(next_tag, content_start) if i < num_masks - 1 else len(assistant_text)
        if end_pos == -1:
            end_pos = len(assistant_text)
        captions.append(assistant_text[content_start:end_pos].strip())
    return captions


# ---------------------------------------------------------------------------
# Helper utilities for the UI
# ---------------------------------------------------------------------------

def _to_binary_mask(mask_img: Image.Image, target_size: Tuple[int, int]) -> np.ndarray:
    arr = np.array(mask_img.convert("L").resize(target_size, Image.NEAREST))
    return (arr > 0).astype(np.uint8)


def make_overlay(pil_image: Image.Image, masks: List[np.ndarray], max_side: int = 768):
    base = pil_image.convert("RGB")
    w, h = base.size
    scale = min(1.0, max_side / max(w, h))
    if scale < 1.0:
        new_size = (max(1, int(w * scale)), max(1, int(h * scale)))
        base = base.resize(new_size, Image.BILINEAR)
    annotations = []
    for idx, mask in enumerate(masks):
        m = mask.astype(np.uint8)
        if scale < 1.0:
            m = np.array(
                Image.fromarray(m * 255).resize(base.size, Image.NEAREST)
            ) > 0
            m = m.astype(np.uint8)
        annotations.append((m, f"Region {idx}"))
    return (base, annotations)


def make_preset_thumbnail(image_path: str, mask_paths: List[str]) -> Image.Image:
    img = Image.open(image_path).convert("RGB")
    base = np.array(img).astype(np.float32)
    for idx, mp in enumerate(mask_paths):
        m = _to_binary_mask(Image.open(mp), img.size).astype(bool)
        color = np.array(OVERLAY_COLORS[idx % len(OVERLAY_COLORS)], dtype=np.float32)
        base[m] = 0.45 * base[m] + 0.55 * color
    out = Image.fromarray(base.astype(np.uint8))
    out.thumbnail((320, 320))
    return out


def _save_mask_png(mask_arr: np.ndarray) -> str:
    """Persist a binary (0/255) mask array to a temp PNG and return its path."""
    import tempfile
    tmp = tempfile.NamedTemporaryFile(suffix=".png", delete=False)
    Image.fromarray(mask_arr.astype(np.uint8), mode="L").save(tmp.name)
    return tmp.name


def _draw_click_markers(pil_image: Image.Image, points: List[List[float]], labels: List[int]) -> Image.Image:
    """Return a copy of the image with click points drawn (green=include, red=exclude)."""
    from PIL import ImageDraw
    img = pil_image.convert("RGB").copy()
    draw = ImageDraw.Draw(img)
    r = max(4, int(min(img.size) * 0.012))
    for (x, y), lab in zip(points, labels):
        color = (16, 185, 129) if lab == 1 else (239, 68, 68)
        draw.ellipse([x - r, y - r, x + r, y + r], fill=color, outline=(255, 255, 255), width=2)
    return img


# ---------------------------------------------------------------------------
# Decoding animation helpers
# ---------------------------------------------------------------------------

def decode_step_captions(step_tokens: torch.Tensor, num_masks: int) -> List[str]:
    """Decode a single denoising step's token state into per-mask captions."""
    ids = step_tokens[0].tolist()
    pieces = []
    for tid in ids:
        if tid == MASK_ID:
            pieces.append(MASK_PLACEHOLDER)
        else:
            pieces.append(TOKENIZER.decode([tid], skip_special_tokens=False))
    raw = "".join(pieces)
    captions = []
    for i in range(num_masks):
        start_tag = f"<|Mask_Cap_{i}|>"
        next_tag = f"<|Mask_Cap_{i + 1}|>"
        start_pos = raw.find(start_tag)
        if start_pos == -1:
            captions.append("")
            continue
        content_start = start_pos + len(start_tag)
        end_pos = raw.find(next_tag, content_start) if i < num_masks - 1 else len(raw)
        if end_pos == -1:
            end_pos = len(raw)
        text = raw[content_start:end_pos]
        for tok in ("", "<|mdm_mask|>"):
            text = text.replace(tok, "")
        captions.append(text.strip())
    return captions


def _render_caption_body(cap: str, prev_cap: str, color: tuple, highlight: bool = True) -> str:
    rgb = f"rgb{color}"
    out = []
    prev_revealed = prev_cap.replace(MASK_PLACEHOLDER, "") if prev_cap else ""
    seen_real = 0
    for ch in cap:
        if ch == MASK_PLACEHOLDER:
            out.append(
                f'<span class="tok-pending" style="background:{rgb};"></span>'
            )
        else:
            seen_real += 1
            is_new = highlight and seen_real > len(prev_revealed)
            esc = html_lib.escape(ch)
            if is_new:
                out.append(f'<span class="tok-new" style="background:{rgb};">{esc}</span>')
            else:
                out.append(esc)
    if not cap:
        return '<span class="tok-empty">…</span>'
    return "".join(out)


def render_caption_html(
    captions: List[str],
    prev_captions: List[str],
    step_idx: int,
    total_steps: int,
) -> str:
    last_step = max(total_steps - 1, 1)
    is_final = step_idx >= total_steps - 1
    pct = int(round(step_idx / last_step * 100))
    css = """
<style>
.dec-wrap { font-family: -apple-system, BlinkMacSystemFont, 'Segoe UI', Roboto, sans-serif; }
.dec-head { display:flex; align-items:center; gap:12px; margin-bottom:14px; }
.dec-step { font-size:0.95em; font-weight:600; color:#475569; white-space:nowrap; }
.dec-progress { flex:1; height:6px; background:#e2e8f0; border-radius:3px; overflow:hidden; }
.dec-progress-fill { height:100%; background:linear-gradient(90deg,#6366f1,#a855f7); border-radius:3px; transition:width 0.25s ease; }
.cap-card { border:1px solid #e2e8f0; border-radius:12px; padding:14px 16px; margin-bottom:12px; background:#fff; box-shadow:0 1px 3px rgba(0,0,0,0.04); }
.cap-title { display:flex; align-items:center; gap:8px; font-weight:600; font-size:0.9em; margin-bottom:8px; color:#1e293b; }
.cap-dot { width:13px; height:13px; border-radius:50%; flex-shrink:0; }
.cap-body { font-size:0.95em; line-height:1.75; color:#0f172a; word-break:break-word; }
.tok-pending { display:inline-block; width:0.55em; height:0.55em; border-radius:50%; margin:0 1px; opacity:0.35; vertical-align:middle; animation:tokpulse 1.1s ease-in-out infinite; }
@keyframes tokpulse { 0%,100%{opacity:0.18;transform:scale(0.8);} 50%{opacity:0.6;transform:scale(1.05);} }
.tok-new { color:#fff; border-radius:4px; padding:0 2px; animation:tokreveal 0.45s ease-out; }
@keyframes tokreveal { from{opacity:0;transform:translateY(-3px) scale(0.9);} to{opacity:1;transform:none;} }
.tok-empty { color:#94a3b8; font-style:italic; }
</style>
"""
    parts = [css, '<div class="dec-wrap">']
    parts.append(
        f'<div class="dec-head"><span class="dec-step">Step {step_idx} / {last_step}</span>'
        f'<div class="dec-progress"><div class="dec-progress-fill" style="width:{pct}%;"></div></div></div>'
    )
    for i, cap in enumerate(captions):
        color = OVERLAY_COLORS[i % len(OVERLAY_COLORS)]
        prev = prev_captions[i] if prev_captions and i < len(prev_captions) else ""
        body = _render_caption_body(cap, prev, color, highlight=not is_final)
        parts.append(
            f'<div class="cap-card"><div class="cap-title">'
            f'<span class="cap-dot" style="background:rgb{color};"></span>Region {i}</div>'
            f'<div class="cap-body">{body}</div></div>'
        )
    parts.append("</div>")
    return "".join(parts)


# ---------------------------------------------------------------------------
# GPU inference function (decorated with @spaces.GPU for ZeroGPU)
# ---------------------------------------------------------------------------

@spaces.GPU(duration=180)
def run_inference_gpu(
    pil_image: Image.Image,
    mask_images: List[Image.Image],
    prompt: str,
    gen_length: int,
    steps: int,
    temperature: float,
    top_p: float,
) -> List[List[str]]:
    """Run the full PerceptionDLM pipeline and return per-step decoding history.

    Each element is a list of per-mask caption strings for that denoising step.
    All CUDA tensors are decoded to text inside the GPU worker so only plain
    Python data crosses the pickle boundary.
    """
    prompt = prompt or DEFAULT_PROMPT
    target_size = pil_image.size

    masks_list = [_to_binary_mask(m, target_size) for m in mask_images]

    sub_images = dynamic_preprocess(
        pil_image,
        min_num=PROCESSOR.min_sub_img,
        max_num=PROCESSOR.max_sub_img,
        image_size=PROCESSOR.image_size[0],
        use_thumbnail=True,
    )
    pixel_values = PROCESSOR.image_processor.preprocess(
        images=sub_images, return_tensors="pt"
    )["pixel_values"].to("cuda").to(DTYPE)
    aspect_ratio = compute_aspect_ratio(
        pil_image, PROCESSOR, num_tiles=pixel_values.shape[0]
    ).to("cuda")

    sort_idx = sort_masks_by_area(masks_list)
    masks_list = [masks_list[i] for i in sort_idx]

    bboxes = build_bboxes(masks_list, TOKENIZER)
    visual_prompt_images, prompt_tokens, _ = build_visual_prompt_matrices(
        masks_list, prompt_numbers=MODEL.config.prompt_numbers
    )

    mask_values_list = []
    for vp_img in visual_prompt_images:
        vp_rgb = vp_img.convert("RGB")
        sub_masks = dynamic_preprocess(
            vp_rgb,
            min_num=PROCESSOR.min_sub_img,
            max_num=PROCESSOR.max_sub_img,
            image_size=PROCESSOR.image_size[0],
            use_thumbnail=True,
        )
        mv = PROCESSOR.image_processor.preprocess(
            images=sub_masks, return_tensors="pt"
        )["pixel_values"].to("cuda").to(DTYPE)
        mask_values_list.append(mv)

    questions = [prompt for _ in masks_list]
    prompt_text = build_prompt_text(
        tokenizer=TOKENIZER,
        num_image_token=MODEL.config.num_image_token,
        num_tiles=pixel_values.shape[0],
        questions=questions,
        gen_len=gen_length,
        num_masks=len(masks_list),
    )

    model_inputs = TOKENIZER(prompt_text, return_tensors="pt")
    input_ids = model_inputs["input_ids"].to("cuda")

    _, all_steps = MODEL.generate_replace_noise(
        pixel_values=pixel_values,
        global_mask_values_list=mask_values_list,
        aspect_ratios=aspect_ratio,
        bboxes=[bboxes],
        input_ids=input_ids,
        steps=steps,
        temperature=temperature,
        top_p=top_p,
        tokenizer=TOKENIZER,
        prompt_tokens=prompt_tokens,
    )

    num_masks = len(masks_list)
    history = [decode_step_captions(step_tok, num_masks) for step_tok in all_steps]
    return history


# ---------------------------------------------------------------------------
# SAM 3 mask generation (decorated with @spaces.GPU for ZeroGPU)
# ---------------------------------------------------------------------------

@spaces.GPU(duration=60)
def sam3_generate_mask_gpu(
    pil_image: Image.Image,
    points: List[List[float]],
    labels: List[int],
) -> np.ndarray:
    """Generate a binary segmentation mask from click points using SAM 3.

    ``points`` is a list of ``[x, y]`` pixel coordinates the user clicked on
    the image; ``labels`` are matching 1 (include) / 0 (exclude) flags. Returns
    a uint8 HxW array (0/255) at the original image resolution.
    """
    image = pil_image.convert("RGB")
    predictor = _get_sam3_predictor()

    point_coords = np.array([[float(x), float(y)] for x, y in points], dtype=np.float32)
    point_labels = np.array([int(l) for l in labels], dtype=np.int32)

    with torch.inference_mode():
        predictor.set_image(image)
        masks, scores, _ = predictor.predict(
            point_coords=point_coords,
            point_labels=point_labels,
            multimask_output=False,
        )
    # masks: (C, H, W) numpy at original resolution; C=1 with multimask_output=False.
    mask = np.asarray(masks[0]) > 0
    return (mask.astype(np.uint8) * 255)


# ---------------------------------------------------------------------------
# Preset examples
# ---------------------------------------------------------------------------
PRESET_DIR = os.path.dirname(os.path.abspath(__file__))
PRESETS: Dict[str, dict] = {}
demo_img = os.path.join(PRESET_DIR, "demo.jpg")
if os.path.exists(demo_img):
    masks = sorted(
        os.path.join(PRESET_DIR, f)
        for f in os.listdir(PRESET_DIR)
        if f.startswith("demo_mask_") and f.endswith(".jpg")
    )
    if masks:
        PRESETS["demo.jpg with 3 masks"] = {"image": demo_img, "masks": masks}

PRESET_KEYS = list(PRESETS.keys())

# ---------------------------------------------------------------------------
# Build the Gradio interface
# ---------------------------------------------------------------------------

CUSTOM_CSS = """
#col-container { max-width: 1200px; margin: 0 auto; }
.dark .gradio-container { color: var(--body-text-color); }
.region-anno { overflow:hidden; }
.region-anno img, .region-anno canvas { max-width:100%; height:auto; object-fit:contain; }
"""

color_map = {
    f"Region {i}": "#%02x%02x%02x" % OVERLAY_COLORS[i % len(OVERLAY_COLORS)]
    for i in range(len(OVERLAY_COLORS))
}

with gr.Blocks(
    title="PerceptionDLM Region Captioning",
    theme=gr.themes.Citrus(),
    css=CUSTOM_CSS,
) as demo:
    gr.Markdown(
        "# 🎯 PerceptionDLM Region Captioning\n"
        "A diffusion multimodal LLM that captions any region of an image **in parallel**. "
        "Upload an image and one or more binary masks, then run inference — "
        "hover over a region to highlight it, and replay the diffusion decoding to watch each "
        "caption emerge token by token.\n\n"
        "Model: [MSALab/PerceptionDLM](https://huggingface.co/MSALab/PerceptionDLM) · "
        "Paper: [arXiv:2606.19534](https://arxiv.org/abs/2606.19534) · "
        "Code: [GitHub](https://github.com/MSALab-PKU/PerceptionDLM)"
    )

    with gr.Row(elem_id="col-container"):
        with gr.Column(scale=1):
            # gr.Markdown("### Input")
            if PRESETS:
                preset_gallery = gr.Gallery(
                    value=[
                        (make_preset_thumbnail(c["image"], c["masks"]), name)
                        for name, c in PRESETS.items()
                    ],
                    columns=3,
                    height="auto",
                    object_fit="cover",
                    allow_preview=False,
                    label=None,
                    show_label=False,
                    visible=False
                )
                #gr.Markdown("*Click a thumbnail to load preset*")
            image_in = gr.Image(type="pil", label="Image", image_mode="RGB")

            gr.Markdown("### Regions to caption")
            mask_mode = gr.Radio(
                choices=[
                    "Generate mask via SAM 3 (click on image)",
                    "Upload a mask file",
                ],
                value="Generate mask via SAM 3 (click on image)",
                label="How to provide masks",
            )

            # ---- Default mode: SAM 3 interactive click-to-mask ----
            with gr.Group(visible=True) as sam3_group:
                gr.Markdown(
                    "Click on the image below to place points, then press "
                    "**Generate mask via SAM 3**. Add the resulting region and "
                    "repeat to caption several regions."
                )
                point_type = gr.Radio(
                    choices=["include", "exclude"],
                    value="include",
                    label="Click type",
                )
                sam3_image = gr.Image(
                    type="pil", label="Click to place points", interactive=True,
                    image_mode="RGB",
                )
                with gr.Row():
                    clear_pts_btn = gr.Button("Clear points", variant="secondary")
                    gen_mask_btn = gr.Button("Generate mask via SAM 3", variant="primary")
                sam3_mask_preview = gr.Image(label="Generated mask (preview)", interactive=False)
                with gr.Row():
                    add_region_btn = gr.Button("➕ Add this region", variant="primary")
                    clear_regions_btn = gr.Button("Clear all regions", variant="secondary")
                sam3_regions_status = gr.Markdown("*No regions added yet.*")

            # ---- Secondary mode: upload pre-existing mask file(s) ----
            with gr.Group(visible=False) as upload_group:
                mask_in = gr.File(
                    file_count="multiple",
                    file_types=["image"],
                    label="Mask images (binary, ≥1)",
                )

            prompt_in = gr.Textbox(value=DEFAULT_PROMPT, label="Prompt")
            with gr.Accordion("Advanced settings", open=False):
                with gr.Row():
                    gen_len_in = gr.Slider(8, 128, value=64, step=8, label="Gen length")
                    steps_in = gr.Slider(8, 128, value=32, step=8, label="Steps")
            run_btn = gr.Button("Run inference", variant="primary")

        with gr.Column(scale=1):
            gr.Markdown("### Output")
            overlay_out = gr.AnnotatedImage(
                label="Regions (hover to highlight)",
                color_map=color_map,
                elem_classes=["region-anno"],
            )
            with gr.Row():
                step_slider = gr.Slider(
                    0, 1, value=0, step=1, label="Decoding step",
                    interactive=True, scale=4,
                )
                play_btn = gr.Button("▶ Play", variant="secondary", scale=1)
            captions_out = gr.HTML()

    # State
    history_state = gr.State([])
    num_masks_state = gr.State(0)
    # SAM 3 interactive state
    sam3_points_state = gr.State([])       # list of [x, y]
    sam3_labels_state = gr.State([])       # list of 1/0
    sam3_regions_state = gr.State([])      # list of generated mask file paths
    sam3_last_mask_state = gr.State(None)  # last generated mask file path

    # ---- Preset loading via gallery click ----
    def load_preset(evt: gr.SelectData):
        name = PRESET_KEYS[evt.index]
        case = PRESETS[name]
        img = Image.open(case["image"]).convert("RGB")
        return img, case["masks"]

    if PRESETS:
        preset_gallery.select(
            load_preset, inputs=None, outputs=[image_in, mask_in]
        )

    # ---- Mode toggle: SAM 3 (default) vs. upload ----
    def _toggle_mode(mode):
        use_sam3 = mode.startswith("Generate mask via SAM 3")
        return gr.update(visible=use_sam3), gr.update(visible=not use_sam3)

    mask_mode.change(
        _toggle_mode, inputs=[mask_mode], outputs=[sam3_group, upload_group]
    )

    # ---- Mirror the main image into the SAM 3 click canvas ----
    def _sync_sam3_image(image):
        # New image resets any pending clicks / preview for that image.
        return image, [], [], None, gr.update(value=None)

    image_in.change(
        _sync_sam3_image,
        inputs=[image_in],
        outputs=[sam3_image, sam3_points_state, sam3_labels_state,
                 sam3_last_mask_state, sam3_mask_preview],
    )

    # ---- Record a click on the SAM 3 canvas ----
    def _on_sam3_click(image, ptype, points, labels, evt: gr.SelectData):
        if image is None:
            return image, points, labels
        x, y = evt.index
        points = points + [[float(x), float(y)]]
        labels = labels + [1 if ptype == "include" else 0]
        marked = _draw_click_markers(image, points, labels)
        return marked, points, labels

    sam3_image.select(
        _on_sam3_click,
        inputs=[image_in, point_type, sam3_points_state, sam3_labels_state],
        outputs=[sam3_image, sam3_points_state, sam3_labels_state],
    )

    # ---- Clear pending points ----
    def _clear_points(image):
        return image, [], [], None, gr.update(value=None)

    clear_pts_btn.click(
        _clear_points,
        inputs=[image_in],
        outputs=[sam3_image, sam3_points_state, sam3_labels_state,
                 sam3_last_mask_state, sam3_mask_preview],
    )

    # ---- Generate a mask from the current points via SAM 3 ----
    def _generate_mask(image, points, labels):
        if image is None:
            raise gr.Error("Please provide an image first.")
        if not points:
            raise gr.Error("Click on the image to place at least one point.")
        mask_arr = sam3_generate_mask_gpu(image, points, labels)
        path = _save_mask_png(mask_arr)
        return path, mask_arr

    gen_mask_btn.click(
        _generate_mask,
        inputs=[image_in, sam3_points_state, sam3_labels_state],
        outputs=[sam3_last_mask_state, sam3_mask_preview],
    )

    # ---- Add the generated mask as a region ----
    def _add_region(image, last_mask, regions):
        if not last_mask:
            raise gr.Error("Generate a mask via SAM 3 before adding a region.")
        regions = regions + [last_mask]
        status = f"**{len(regions)} region(s) added.** Ready to run inference."
        # Reset pending points/preview so the next region starts fresh.
        return regions, [], [], None, gr.update(value=None), image, status

    add_region_btn.click(
        _add_region,
        inputs=[image_in, sam3_last_mask_state, sam3_regions_state],
        outputs=[sam3_regions_state, sam3_points_state, sam3_labels_state,
                 sam3_last_mask_state, sam3_mask_preview, sam3_image,
                 sam3_regions_status],
    )

    # ---- Clear all added regions ----
    def _clear_regions(image):
        return [], [], [], None, gr.update(value=None), image, "*No regions added yet.*"

    clear_regions_btn.click(
        _clear_regions,
        inputs=[image_in],
        outputs=[sam3_regions_state, sam3_points_state, sam3_labels_state,
                 sam3_last_mask_state, sam3_mask_preview, sam3_image,
                 sam3_regions_status],
    )

    # ---- Run inference ----
    def _on_run(image, mode, sam3_regions, mask_files, prompt, gen_len, steps):
        """Run PerceptionDLM inference and return overlay + decoding animation."""
        if image is None:
            raise gr.Error("Please provide an image.")
        use_sam3 = mode.startswith("Generate mask via SAM 3")
        if use_sam3:
            if not sam3_regions:
                raise gr.Error(
                    "Generate at least one region with SAM 3 (click the image, "
                    "generate a mask, then 'Add this region')."
                )
            mask_paths = list(sam3_regions)
        else:
            if not mask_files:
                raise gr.Error("Please provide at least one mask image.")
            mask_paths = [f if isinstance(f, str) else f.name for f in mask_files]
        mask_images = [Image.open(p) for p in mask_paths]

        history = run_inference_gpu(
            image, mask_images, prompt, int(gen_len), int(steps),
            temperature=0.0, top_p=1.0,
        )

        # Rebuild masks_list for overlay (same sorting as inside GPU fn)
        target_size = image.size
        masks_list = [_to_binary_mask(m, target_size) for m in mask_images]
        sort_idx = sort_masks_by_area(masks_list)
        masks_list = [masks_list[i] for i in sort_idx]
        overlay = make_overlay(image, masks_list)

        total = len(history)
        last = total - 1
        html = render_caption_html(
            history[last], history[last - 1] if last > 0 else [], last, total
        )
        slider_update = gr.update(minimum=0, maximum=last, value=last, step=1)
        return history, len(masks_list), overlay, slider_update, html

    run_btn.click(
        _on_run,
        inputs=[image_in, mask_mode, sam3_regions_state, mask_in,
                prompt_in, gen_len_in, steps_in],
        outputs=[history_state, num_masks_state, overlay_out, step_slider, captions_out],
        api_name="run_inference",
    )

    # ---- Step slider scrubbing ----
    def _on_step(step_idx, history):
        if not history:
            return gr.update()
        total = len(history)
        i = int(step_idx)
        i = max(0, min(i, total - 1))
        prev = history[i - 1] if i > 0 else []
        return render_caption_html(history[i], prev, i, total)

    step_slider.change(_on_step, inputs=[step_slider, history_state], outputs=[captions_out])

    # ---- Play animation ----
    def _on_play(history):
        if not history:
            yield gr.update(), gr.update()
            return
        total = len(history)
        for i in range(total):
            prev = history[i - 1] if i > 0 else []
            html = render_caption_html(history[i], prev, i, total)
            yield gr.update(value=i), html
            if i < total - 1:
                time.sleep(0.25)

    play_btn.click(_on_play, inputs=[history_state], outputs=[step_slider, captions_out])

    # ---- Examples ----
    # Each row maps 1:1 to inputs=[image_in, mask_in, prompt_in]. Because
    # mask_in is a file_count="multiple" component, its example value must be a
    # SINGLE element that is itself the list of mask paths (not the mask paths
    # spread across the row), otherwise the values shift into the wrong fields.
    UPLOAD_MODE = "Upload a mask file"
    example_entries = []
    for name, case in PRESETS.items():
        # Columns: image, mask-mode, list-of-masks, prompt -> aligned 1:1 with
        # inputs=[image_in, mask_mode, mask_in, prompt_in].
        example_entries.append([case["image"], UPLOAD_MODE, case["masks"], DEFAULT_PROMPT])

    def _run_example(image, mode, mask_files, prompt):
        """Run an example (examples supply masks via the upload mode)."""
        return _on_run(
            image, mode, [], mask_files, prompt,
            gen_len_in.value, steps_in.value,
        )

    if example_entries:
        gr.Examples(
            examples=example_entries,
            inputs=[image_in, mask_mode, mask_in, prompt_in],
            fn=_run_example,
            outputs=[history_state, num_masks_state, overlay_out, step_slider, captions_out],
            cache_examples=False,
        )


demo.queue()
if __name__ == "__main__":
    demo.launch(mcp_server=True)