# MAOAM-GLaMM demo, adapted for Hugging Face Spaces ZeroGPU.
#
# This is the authors' GLaMM/demo.py (full star-click + text + click+text UI)
# with the minimal changes required to run on ZeroGPU:
#   1. `import spaces` before torch (monkey-patches torch.cuda.*).
#   2. Model built ONCE at module scope; weights resolved from the Hub instead
#      of argparse local paths.
#   3. The only real GPU work (the model forward) runs inside @spaces.GPU.
#   4. Caches anchored to a writable /tmp (ZeroGPU /data and ~/.cache are RO).
#   5. launch() with no server args (Spaces injects them).
#   6. GLaMM meta-tensor load fix kept (low_cpu_mem_usage=False + materialize
#      residual meta params) so the module-scope .to("cuda") cannot crash.
#
# Backend: LLaVA-Llama (GranD pretrained) + SAM ViT-H grounding encoder.

# --- 1. Cache env vars FIRST, before importing torch/transformers ---
import os


def _writable_cache_root():
    for cand in ("/data", "/tmp"):
        try:
            probe = os.path.join(cand, ".hf_write_probe")
            os.makedirs(cand, exist_ok=True)
            with open(probe, "w") as f:
                f.write("ok")
            os.remove(probe)
            return cand
        except Exception:
            continue
    return "/tmp"


_CACHE_ROOT = _writable_cache_root()
os.environ.setdefault("HF_HOME", os.path.join(_CACHE_ROOT, ".cache/huggingface"))
os.environ.setdefault("HF_MODULES_CACHE", "/tmp/hf_modules")
os.environ.setdefault("MPLCONFIGDIR", "/tmp/matplotlib")

import sys

# Make the vendored packages importable exactly the way GLaMM/demo.py expects.
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

# --- 2. import spaces BEFORE torch ---
import spaces

import io
import time
import zipfile
import tempfile
import traceback
from datetime import datetime

import numpy as np
import torch
import torch.nn.functional as F
import torchvision.transforms as transforms
import transformers
import gradio as gr
from PIL import Image
from transformers import CLIPImageProcessor

from model.GLaMM import GLaMMForCausalLM
from utils.hm_utils import add_star_marker
from dataset.datasets import custom_collate_fn_multi
from model.llava import conversation as conversation_lib
from dataset.utils.utils import (
    STAR_QUESTIONS,
    REFERRING_QUESTIONS,
    SEG_QUESTIONS,
    TASK_PROMPT,
)
from tools.glamm_eval_utils import (
    DEFAULT_IM_END_TOKEN,
    DEFAULT_IM_START_TOKEN,
    DEFAULT_IMAGE_TOKEN,
)
from check_load import load_mp_rank_checkpoint

# ---------------------------------------------------------------------------
# Weights configuration
# ---------------------------------------------------------------------------
BASE_MODEL = "MBZUAI/GLaMM-GranD-Pretrained"
VISION_TOWER = "openai/clip-vit-large-patch14-336"
CKPT_REPO = "jpark677/maoam_ckpts"
CKPT_FILE = "glamm/mp_rank_00_model_states.pt"
SAM_URL = "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"


def _ensure_sam_checkpoint() -> str:
    from huggingface_hub import hf_hub_download

    mirror_repo = os.environ.get("SAM_REPO", "")
    mirror_file = os.environ.get("SAM_REPO_FILE", "sam_vit_h_4b8939.pth")
    if mirror_repo:
        try:
            return hf_hub_download(repo_id=mirror_repo, filename=mirror_file)
        except Exception as e:
            print(f"[WARN] SAM mirror download failed ({e}); using fbaipublicfiles URL")

    cache_dir = os.path.join(os.environ["HF_HOME"], "sam")
    os.makedirs(cache_dir, exist_ok=True)
    dest = os.path.join(cache_dir, "sam_vit_h_4b8939.pth")
    if not os.path.exists(dest):
        print(f"[INFO] Downloading SAM ViT-H weights to {dest}")
        torch.hub.download_url_to_file(SAM_URL, dest, progress=False)
    return dest


def _resolve_ckpt() -> str:
    from huggingface_hub import hf_hub_download

    return hf_hub_download(repo_id=CKPT_REPO, filename=CKPT_FILE)


# ---------------------------------------------------------------------------
# GLaMMDemo: the authors' class from GLaMM/demo.py. _setup_model carries the
# ZeroGPU meta-tensor load fix; everything else (star overlays, batch build,
# inference routing) is preserved.
# ---------------------------------------------------------------------------
class GLaMMDemo:
    def __init__(self, model_path, args_dict):
        self.model_path = model_path
        self.args_dict = args_dict
        self.model = None
        self.tokenizer = None
        # On ZeroGPU torch.cuda.is_available() is False at import time even
        # though .to("cuda") is emulated, so force cuda. Real compute happens
        # inside @spaces.GPU.
        self.device = "cuda"
        self._setup_tokenizer()
        self._setup_model()

    @staticmethod
    def _ensure_unit_range(tensor: torch.Tensor) -> torch.Tensor:
        if tensor.numel() == 0:
            return tensor
        tensor = tensor.to(dtype=torch.float32)
        t_min = float(tensor.min().item())
        t_max = float(tensor.max().item())
        if 0.0 <= t_min and t_max <= 1.0:
            return tensor
        if 0.0 <= t_min and t_max <= 255.0:
            tensor = tensor / 255.0
        else:
            denom = max(t_max - t_min, 1e-6)
            tensor = (tensor - t_min) / denom
        return tensor.clamp_(0.0, 1.0)

    def _image_to_tensor(self, image):
        if isinstance(image, torch.Tensor):
            image_tensor = image.detach().clone()
            if image_tensor.ndim == 3 and image_tensor.shape[0] in (1, 3):
                pass
            elif image_tensor.ndim == 3:
                image_tensor = image_tensor.permute(2, 0, 1)
            elif image_tensor.ndim == 2:
                image_tensor = image_tensor.unsqueeze(0)
            image_tensor = image_tensor.to(dtype=torch.float32)
        elif isinstance(image, np.ndarray):
            if image.dtype == np.uint8:
                pil_image = Image.fromarray(image)
                image_tensor = transforms.ToTensor()(pil_image)
            else:
                image_tensor = torch.from_numpy(image).float()
                if image_tensor.ndim == 3 and image_tensor.shape[2] in (1, 3):
                    image_tensor = image_tensor.permute(2, 0, 1)
                elif image_tensor.ndim == 2:
                    image_tensor = image_tensor.unsqueeze(0)
                image_tensor = (
                    image_tensor / 255.0 if image_tensor.max() > 1.0 else image_tensor
                )
        else:
            image_tensor = transforms.ToTensor()(image)
        return self._ensure_unit_range(image_tensor)

    def _setup_tokenizer(self):
        self.tokenizer = transformers.AutoTokenizer.from_pretrained(
            self.model_path,
            model_max_length=self.args_dict.get("model_max_length", 1536),
            padding_side="right",
            use_fast=False,
        )
        self.tokenizer.pad_token = self.tokenizer.unk_token

        if not self.args_dict.get("pretrained", False):
            if self.args_dict.get("use_mm_start_end", True):
                self.tokenizer.add_tokens(
                    [DEFAULT_IM_START_TOKEN, DEFAULT_IM_END_TOKEN], special_tokens=True
                )
            reg_tokens = ["<bbox>", "<point>"]
            segmentation_tokens = ["[SEG]"]
            phrase_tokens = ["<p>", "</p>"]
            special_tokens = reg_tokens + segmentation_tokens + phrase_tokens
            self.tokenizer.add_tokens(special_tokens, special_tokens=True)

        self.tokenizer.add_special_tokens(
            {"additional_special_tokens": ["<bbox>", "[SEG]", "<p>", "</p>"]}
        )

        self.bbox_token_idx = self.tokenizer("<bbox>", add_special_tokens=False).input_ids[0]
        self.seg_token_idx = self.tokenizer("[SEG]", add_special_tokens=False).input_ids[0]
        self.bop_token_idx = self.tokenizer("<p>", add_special_tokens=False).input_ids[0]
        self.eop_token_idx = self.tokenizer("</p>", add_special_tokens=False).input_ids[0]

    def _setup_model(self):
        model_args = {
            "train_mask_decoder": self.args_dict.get("train_mask_decoder", True),
            "out_dim": self.args_dict.get("out_dim", 256),
            "ce_loss_weight": self.args_dict.get("ce_loss_weight", 1.0),
            "dice_loss_weight": self.args_dict.get("dice_loss_weight", 0.5),
            "bce_loss_weight": self.args_dict.get("bce_loss_weight", 2.0),
            "seg_token_idx": self.seg_token_idx,
            "vision_pretrained": self.args_dict.get("vision_pretrained", ""),
            "vision_tower": self.args_dict.get("vision_tower", VISION_TOWER),
            "use_mm_start_end": self.args_dict.get("use_mm_start_end", True),
            "mm_vision_select_layer": self.args_dict.get("mm_vision_select_layer", -2),
            "pretrain_mm_mlp_adapter": self.args_dict.get("pretrain_mm_mlp_adapter", ""),
            "tune_mm_mlp_adapter": self.args_dict.get("tune_mm_mlp_adapter", False),
            "freeze_mm_mlp_adapter": self.args_dict.get("freeze_mm_mlp_adapter", False),
            "mm_use_im_start_end": self.args_dict.get("mm_use_im_start_end", True),
            "with_region": self.args_dict.get("with_region", True),
            "bbox_token_idx": self.bbox_token_idx,
            "eop_token_idx": self.eop_token_idx,
            "bop_token_idx": self.bop_token_idx,
        }
        model_args["num_level_reg_features"] = 4

        # low_cpu_mem_usage=True initializes on the meta device and leaves any
        # param NOT in the base checkpoint (GLaMM/SAM heads, region encoder)
        # unmaterialized, which makes the later .to("cuda") crash with
        # "Cannot copy out of meta tensor". Load fully-materialized on CPU.
        self.model = GLaMMForCausalLM.from_pretrained(
            self.model_path,
            torch_dtype=torch.bfloat16,
            low_cpu_mem_usage=False,
            **model_args,
        )

        resume_ckpt = self.args_dict.get("resume") or ""
        if resume_ckpt:
            if not os.path.exists(resume_ckpt):
                raise FileNotFoundError(f"Checkpoint not found at '{resume_ckpt}'.")
            print(f"[INFO] Loading checkpoint from {resume_ckpt}")
            _, state_dict = load_mp_rank_checkpoint(resume_ckpt)
            cleaned_state = {
                (k[7:] if k.startswith("module.") else k): v
                for k, v in state_dict.items()
            }
            missing, unexpected = self.model.load_state_dict(cleaned_state, strict=False)
            print(f"[INFO] Checkpoint load complete: missing={len(missing)} unexpected={len(unexpected)}")

        self.model.config.eos_token_id = self.tokenizer.eos_token_id
        self.model.config.bos_token_id = self.tokenizer.bos_token_id
        self.model.config.pad_token_id = self.tokenizer.pad_token_id

        self.model.enable_input_require_grads()
        self.model.gradient_checkpointing_enable()

        self.model.get_model().initialize_vision_modules(self.model.get_model().config)
        vision_tower = self.model.get_model().get_vision_tower()
        vision_tower.to(dtype=torch.bfloat16, device=self.device)

        if not self.args_dict.get("pretrained", False):
            self.model.get_model().initialize_glamm_model(self.model.get_model().config)
        else:
            for param in self.model.get_model().grounding_encoder.parameters():
                param.requires_grad = False
            if self.model.get_model().config.train_mask_decoder:
                self.model.get_model().grounding_encoder.mask_decoder.train()
                for p in self.model.get_model().grounding_encoder.mask_decoder.parameters():
                    p.requires_grad = True
            self.model.get_model().text_hidden_fcs.train()
            for p in self.model.get_model().text_hidden_fcs.parameters():
                p.requires_grad = True

        for p in vision_tower.parameters():
            p.requires_grad = False
        for p in self.model.get_model().mm_projector.parameters():
            p.requires_grad = False

        lora_r = self.args_dict.get("lora_r", 0)
        if lora_r == 0:
            for p in self.model.get_model().layers.parameters():
                p.requires_grad = True
            for p in self.model.get_model().mm_projector.parameters():
                p.requires_grad = True

        conversation_lib.default_conversation = conversation_lib.conv_templates[
            self.args_dict.get("conv_type", "llava_v1")
        ]

        self.model.resize_token_embeddings(len(self.tokenizer))

        # Safety net: materialize any param/buffer still on the meta device so
        # the move to cuda cannot raise "Cannot copy out of meta tensor".
        for module in self.model.modules():
            for p_name, p in list(module.named_parameters(recurse=False)):
                if getattr(p, "is_meta", False):
                    module.register_parameter(
                        p_name,
                        torch.nn.Parameter(
                            torch.zeros(p.shape, dtype=p.dtype, device="cpu"),
                            requires_grad=p.requires_grad,
                        ),
                    )
            for b_name, b in list(module.named_buffers(recurse=False)):
                if getattr(b, "is_meta", False):
                    module.register_buffer(
                        b_name, torch.zeros(b.shape, dtype=b.dtype, device="cpu")
                    )

        self.model.to(self.device)
        self.model.eval()

        self.global_enc_processor = CLIPImageProcessor.from_pretrained(
            self.args_dict.get("vision_tower", VISION_TOWER)
        )
        print("[INFO] Model loaded successfully")

    def grounding_enc_processor(self, x: torch.Tensor, image_size: tuple) -> torch.Tensor:
        img_mean = torch.tensor([123.675, 116.28, 103.53], device=x.device).view(1, -1, 1, 1)
        img_std = torch.tensor([58.395, 57.12, 57.375], device=x.device).view(1, -1, 1, 1)
        x = (x - img_mean) / img_std
        h, w = x.shape[-2:]
        target_h, target_w = image_size
        x = F.pad(x, (0, target_w - w, 0, target_h - h))
        return x

    def resize_image_to_square(self, image, target_size=1024):
        image_tensor = self._image_to_tensor(image)
        H, W = image_tensor.shape[-2:]
        if H < W:
            new_h = target_size
            new_w = int(target_size * W / H)
        else:
            new_w = target_size
            new_h = int(target_size * H / W)
        image_tensor = F.interpolate(
            image_tensor.unsqueeze(0), size=(new_h, new_w), mode="bilinear", align_corners=False
        ).squeeze(0)
        image_tensor = transforms.CenterCrop((target_size, target_size))(image_tensor)
        return image_tensor

    def _overlay_all_stars_1024(self, base_tensor_1024_chw, coords_1024, fixed_color):
        img = base_tensor_1024_chw.clone()
        latched = fixed_color
        marker_size = max(8, int(1024 // 32))
        for i, (hh, ww) in enumerate(coords_1024):
            h = int(max(0, min(1023, hh)))
            w = int(max(0, min(1023, ww)))
            try:
                if i == 0 and latched is None:
                    img, c = add_star_marker(img, h, w, size=marker_size)
                    latched = c or "blue"
                else:
                    img, _ = add_star_marker(img, h, w, size=marker_size, color=latched)
            except TypeError:
                img, c = add_star_marker(img, h, w, size=marker_size)
                if i == 0 and latched is None:
                    latched = c or "blue"
        return img, latched

    def create_data_batch(
        self,
        image,
        coords_list_1024,
        text_prompt="Please segment all pixels with the same material as where the <COLOR> star is.",
        fixed_star_color=None,
    ):
        """Create data batch for star/referring using add_star_marker for overlays."""
        if isinstance(image, np.ndarray):
            original_image_pil = Image.fromarray(image)
        else:
            original_image_pil = image

        base_1024_chw = self.resize_image_to_square(original_image_pil, 1024)
        grounding_enc_image = base_1024_chw.clone()
        global_enc_tensor = base_1024_chw.clone()

        if len(coords_list_1024) > 0:
            global_enc_tensor, latched_color = self._overlay_all_stars_1024(
                global_enc_tensor, coords_list_1024, fixed_star_color
            )
            star_color = latched_color or fixed_star_color or "blue"
        else:
            star_color = fixed_star_color

        global_enc_image = transforms.ToPILImage()(global_enc_tensor)

        if star_color is not None:
            processed_prompt = text_prompt.replace("<COLOR>", star_color)
        else:
            processed_prompt = text_prompt.replace("<COLOR>", "").replace("  ", " ").strip()

        conv = conversation_lib.default_conversation.copy()
        conv.messages = []
        begin_str = f"The {DEFAULT_IMAGE_TOKEN} provides an overview of the picture.\n"
        question = begin_str + processed_prompt
        answer = "[SEG]"
        conv.append_message(conv.roles[0], question)
        conv.append_message(conv.roles[1], answer)
        conversation_str = conv.get_prompt()

        fake_batch_item = {
            "filepath": "demo_image_path",
            "image_star": global_enc_tensor,
            "image_without_star": grounding_enc_image.clone(),
            "grounding_image": grounding_enc_image.clone(),
            "masks": torch.zeros(1, 1024, 1024),
            "orig_size": (1024, 1024),
            "sampled_classes": [0],
            "coords": [list(map(int, xy)) for xy in coords_list_1024],
            "global_enc_processor": self.global_enc_processor,
            "star": {
                "conversation": ([conversation_str] if len(coords_list_1024) >= 1 else None),
                "question": question if len(coords_list_1024) >= 1 else None,
            },
            "referring": {
                "conversation": ([conversation_str] if len(coords_list_1024) == 0 else None),
                "question": question if len(coords_list_1024) == 0 else None,
                "desc": processed_prompt if len(coords_list_1024) == 0 else None,
            },
            "vqa": {"conversation": None, "question": None, "answer": None},
        }

        batch_data = custom_collate_fn_multi(
            [fake_batch_item],
            tokenizer=self.tokenizer,
            use_mm_start_end=self.args_dict.get("use_mm_start_end", True),
            inference=True,
        )

        def move_to_device_recursive(obj, device):
            if isinstance(obj, torch.Tensor):
                return obj.to(device)
            elif isinstance(obj, dict):
                return {k: move_to_device_recursive(v, device) for k, v in obj.items()}
            elif isinstance(obj, list):
                return [move_to_device_recursive(item, device) for item in obj]
            else:
                return obj

        data_batch = move_to_device_recursive(batch_data, self.device)

        if "grounding_enc_images" in data_batch:
            grounding_tensors = (
                (data_batch["grounding_enc_images"] * 255).to(torch.uint8).contiguous()
            )
            data_batch["grounding_enc_images"] = self.grounding_enc_processor(
                grounding_tensors.float(), data_batch.get("orig_size", (1024, 1024))
            )

        for key in ["global_enc_images", "images_star", "images_without_star"]:
            if (
                key in data_batch
                and isinstance(data_batch[key], torch.Tensor)
                and data_batch[key].dtype != torch.bfloat16
            ):
                data_batch[key] = data_batch[key].to(dtype=torch.bfloat16)

        if "masks_list" in data_batch and isinstance(data_batch["masks_list"], list):
            data_batch["masks_list"] = [
                mask.to(dtype=torch.bfloat16) if mask.dtype != torch.bfloat16 else mask
                for mask in data_batch["masks_list"]
            ]

        viz_tensor = global_enc_tensor.clone()
        return data_batch, viz_tensor, star_color, global_enc_image

    def inference(self, image, coords_list_1024, text_prompt, fixed_star_color=None):
        """Run inference; routes to star vs referring by number of coords.

        NOTE: this whole method is called inside @spaces.GPU (it builds the
        batch on cuda and runs the forward), so all CUDA work stays there.
        """
        try:
            data_batch, image_tensor, star_color, global_enc_image = self.create_data_batch(
                image, coords_list_1024, text_prompt, fixed_star_color=fixed_star_color
            )
            with torch.no_grad():
                results = self.model(**data_batch)

            task_key = "referring" if len(coords_list_1024) == 0 else "star"
            task_out = results.get(task_key)
            if task_out is None:
                return None, None, None, None, f"No {task_key} task results."

            predictions = task_out.get("pred_masks")
            if predictions is None:
                return (None, None, None, None,
                        f"No predictions for {task_key}. Available keys: {list(task_out.keys())}")

            pred_mask_raw = predictions[0].detach().float().cpu().numpy()
            if pred_mask_raw.ndim == 3 and pred_mask_raw.shape[0] == 1:
                pred_mask_raw = pred_mask_raw[0]

            pred_mask_binary = (pred_mask_raw > 0).astype(np.uint8) * 255
            return pred_mask_binary, pred_mask_raw, image_tensor, star_color, global_enc_image
        except Exception as e:
            error_msg = f"Error during inference: {str(e)}\n{traceback.format_exc()}"
            print(f"[DEBUG] inference: {error_msg}")
            return None, None, None, None, error_msg


# ---------------------------------------------------------------------------
# Build the model ONCE at module scope.
# ---------------------------------------------------------------------------
print("[INFO] Resolving weights...")
_SAM_CKPT = _ensure_sam_checkpoint()
_RESUME_CKPT = _resolve_ckpt()

ARGS = {
    "model_max_length": 1536,
    "use_mm_start_end": True,
    "pretrained": True,
    "train_mask_decoder": True,
    "out_dim": 256,
    "ce_loss_weight": 1.0,
    "dice_loss_weight": 0.5,
    "bce_loss_weight": 2.0,
    "vision_pretrained": _SAM_CKPT,
    "vision_tower": VISION_TOWER,
    "mm_vision_select_layer": -2,
    "pretrain_mm_mlp_adapter": "",
    "tune_mm_mlp_adapter": False,
    "freeze_mm_mlp_adapter": False,
    "mm_use_im_start_end": True,
    "with_region": True,
    "conv_type": "llava_v1",
    "lora_r": 0,
    "resume": _RESUME_CKPT,
}

print("[INFO] Building GLaMM model at module scope...")
DEMO_MODEL = GLaMMDemo(BASE_MODEL, ARGS)
print("[INFO] GLaMM model ready.")


# ---------------------------------------------------------------------------
# The single GPU entry point: the model forward runs here.
# ---------------------------------------------------------------------------
@spaces.GPU(duration=120)
def _gpu_inference(image, coords_1024, text_prompt, fixed_star_color):
    t0 = time.perf_counter()
    out = DEMO_MODEL.inference(image, coords_1024, text_prompt, fixed_star_color=fixed_star_color)
    print(f"[TIMING] inference took {time.perf_counter() - t0:.2f}s")
    return out


# ---------------------------------------------------------------------------
# Overlay helpers (deterministic cyan blend; from demo.py, no matplotlib).
# ---------------------------------------------------------------------------
def _tensor_chw_to_uint8_hwc(t: torch.Tensor) -> np.ndarray:
    t = t.detach().cpu().to(dtype=torch.float32).clamp(0.0, 1.0)
    if t.ndim == 3 and t.shape[0] in (1, 3):
        pass
    elif t.ndim == 3 and t.shape[-1] in (1, 3):
        t = t.permute(2, 0, 1)
    hwc = (t * 255.0).byte().permute(1, 2, 0).numpy()
    if hwc.shape[2] == 1:
        hwc = np.repeat(hwc, 3, axis=2)
    return hwc


def _create_cyan_overlay(base_rgb_uint8, mask_bool, alpha=0.45):
    base = base_rgb_uint8.astype(np.float32)
    overlay = base.copy()
    cyan = np.array([0.0, 255.0, 255.0], dtype=np.float32)
    overlay[mask_bool] = overlay[mask_bool] * (1.0 - alpha) + cyan * alpha
    return np.clip(overlay, 0, 255).astype(np.uint8)


def _png_bytes_from_uint8_hwc(img: np.ndarray) -> bytes:
    pil = Image.fromarray(img.astype(np.uint8))
    buf = io.BytesIO()
    pil.save(buf, format="PNG")
    return buf.getvalue()


def _png_bytes_from_uint8_mask(mask: np.ndarray) -> bytes:
    pil = Image.fromarray(mask.astype(np.uint8), mode="L")
    buf = io.BytesIO()
    pil.save(buf, format="PNG")
    return buf.getvalue()


def _download_with_prompt(orig_pil, clean_tensor_chw, star_tensor_chw, pred_mask_bin, prompt_text):
    if clean_tensor_chw is None or star_tensor_chw is None or pred_mask_bin is None:
        return None

    orig_1024 = _tensor_chw_to_uint8_hwc(clean_tensor_chw)
    star_1024 = _tensor_chw_to_uint8_hwc(star_tensor_chw)

    mask = pred_mask_bin.astype(np.uint8)
    if mask.ndim == 3:
        mask = mask[..., 0]
    mask_bool = (mask > 0)
    overlay_orig = _create_cyan_overlay(orig_1024, mask_bool)
    overlay_star = _create_cyan_overlay(star_1024, mask_bool)

    ts = datetime.now().strftime("%Y%m%d_%H%M%S")
    tmp = tempfile.NamedTemporaryFile(delete=False, suffix=f"_{ts}.zip")
    tmp.close()
    with zipfile.ZipFile(tmp.name, "w", compression=zipfile.ZIP_DEFLATED) as zf:
        zf.writestr("original_1024.png", _png_bytes_from_uint8_hwc(orig_1024))
        zf.writestr("with_star_1024.png", _png_bytes_from_uint8_hwc(star_1024))
        zf.writestr("pred_mask_1024.png", _png_bytes_from_uint8_mask(mask))
        zf.writestr("overlay_on_original_1024.png", _png_bytes_from_uint8_hwc(overlay_orig))
        zf.writestr("overlay_on_star_1024.png", _png_bytes_from_uint8_hwc(overlay_star))
        if orig_pil is not None:
            buf = io.BytesIO()
            orig_pil.save(buf, format="PNG")
            zf.writestr("original_fullres.png", buf.getvalue())
        if prompt_text is not None:
            zf.writestr("prompt.txt", str(prompt_text))
    return tmp.name


# ---------------------------------------------------------------------------
# Gradio UI (the authors' full star-click interface).
# ---------------------------------------------------------------------------
MAX_STARS = 5

_SELECTION_MODES = {
    "Material: click": (
        STAR_QUESTIONS[0],
        "Place one or more stars on the image, then click **Submit**. `<COLOR>` is filled automatically.",
    ),
    "Material: text": (
        REFERRING_QUESTIONS[0],
        "Replace **`<DESCRIPTION>`** with your material description (e.g. *shiny chrome metal*). No stars needed.",
    ),
    "Material: click + text": (
        f"Please segment all pixels made of the material described below, where the <COLOR> star is.\nDescription: <DESCRIPTION>\n\n{TASK_PROMPT}",
        "Place a star, then replace **`<DESCRIPTION>`** with your material description.",
    ),
    "Object: text": (
        SEG_QUESTIONS[0].replace("{class_name}", "<OBJECT>"),
        "Replace **`<OBJECT>`** with your object expression (e.g. *the man in a red shirt*). No stars needed.",
    ),
}
_DEFAULT_MODE = "Material: click"


def on_selection_change(mode):
    prompt, hint = _SELECTION_MODES.get(mode, _SELECTION_MODES[_DEFAULT_MODE])
    return prompt, hint


def _coerce_to_pil(image_obj):
    """Accept the many shapes an image can arrive in (numpy, PIL, filepath str,
    or a Gradio file-data dict) and return an RGB PIL image, or None."""
    if image_obj is None:
        return None
    if isinstance(image_obj, np.ndarray):
        return Image.fromarray(image_obj.astype(np.uint8)).convert("RGB")
    if isinstance(image_obj, Image.Image):
        return image_obj.convert("RGB")
    if isinstance(image_obj, dict):
        image_obj = image_obj.get("path") or image_obj.get("url") or image_obj.get("name")
    if isinstance(image_obj, str):
        return Image.open(image_obj).convert("RGB")
    raise TypeError(f"Unsupported image input type: {type(image_obj)}")


def on_image_upload(image_np_or_pil):
    orig_pil = _coerce_to_pil(image_np_or_pil)
    if orig_pil is None:
        return None, [], None, None, None, None, "Upload an image to start."
    clean_chw = DEMO_MODEL.resize_image_to_square(orig_pil, 1024)
    disp_chw = clean_chw.clone()
    disp_np = np.array(transforms.ToPILImage()(disp_chw))
    return (
        disp_np,
        [],
        None,
        orig_pil,
        clean_chw,
        disp_chw,
        "Image loaded. Click up to 5 points, then Submit.",
    )


def _add_star_to_tensor(disp_tensor_chw, coords_1024, fixed_color, h, w):
    marker_size = 32
    disp = disp_tensor_chw
    try:
        if fixed_color is None:
            disp, c = add_star_marker(disp, h, w, size=marker_size)
            fixed_color = c or "blue"
        else:
            disp, _ = add_star_marker(disp, h, w, size=marker_size, color=fixed_color)
    except TypeError:
        disp, c = add_star_marker(disp, h, w, size=marker_size)
        if fixed_color is None:
            fixed_color = c or "blue"
    coords_1024 = coords_1024 + [(h, w)]
    return disp, coords_1024, fixed_color


def on_click_add_star(image_disp_np, coords_1024, fixed_color, disp_tensor_chw, evt: gr.SelectData):
    if image_disp_np is None or disp_tensor_chw is None:
        return image_disp_np, coords_1024, fixed_color, disp_tensor_chw, "Please upload an image first."
    if evt is None:
        return image_disp_np, coords_1024, fixed_color, disp_tensor_chw, "Click anywhere on the image to add a star."
    if coords_1024 is None:
        coords_1024 = []
    if len(coords_1024) >= MAX_STARS:
        return image_disp_np, coords_1024, fixed_color, disp_tensor_chw, f"Max {MAX_STARS} stars reached."
    h, w = int(evt.index[1]), int(evt.index[0])
    disp, coords_1024, fixed_color = _add_star_to_tensor(disp_tensor_chw, coords_1024, fixed_color, h, w)
    disp_np = np.array(transforms.ToPILImage()(disp))
    return disp_np, coords_1024, fixed_color, disp, f"Star #{len(coords_1024)} @ (h={h}, w={w})."


def on_undo_last(coords_1024, fixed_color, clean_tensor_chw):
    if clean_tensor_chw is None:
        return None, coords_1024, fixed_color, None, "Nothing to undo."
    if not coords_1024:
        return (
            np.array(transforms.ToPILImage()(clean_tensor_chw.clone())),
            [], None, clean_tensor_chw.clone(), "Nothing to undo.",
        )
    new_coords = coords_1024[:-1]
    disp = clean_tensor_chw.clone()
    latched = fixed_color
    marker_size = max(8, int(1024 // 32))
    for i, (h, w) in enumerate(new_coords):
        try:
            if i == 0 and latched is None:
                disp, c = add_star_marker(disp, int(h), int(w), size=marker_size)
                latched = c or "blue"
            else:
                disp, _ = add_star_marker(disp, int(h), int(w), size=marker_size, color=latched)
        except TypeError:
            disp, c = add_star_marker(disp, int(h), int(w), size=marker_size)
            if i == 0 and latched is None:
                latched = c or "blue"
    disp_np = np.array(transforms.ToPILImage()(disp))
    return disp_np, new_coords, latched, disp, f"Removed last star. {len(new_coords)} remaining."


def on_clear_stars(clean_tensor_chw):
    if clean_tensor_chw is None:
        return None, [], None, None, "Nothing to clear."
    disp = clean_tensor_chw.clone()
    disp_np = np.array(transforms.ToPILImage()(disp))
    return disp_np, [], None, disp, "Cleared all stars."


def on_submit(orig_pil, coords_1024, text_prompt, fixed_color, clean_tensor_chw):
    if orig_pil is None:
        return None, None, "Please upload an image first.", None, None, None
    # Defensive: state may arrive as a filepath/dict (e.g. if examples were
    # cached), so normalize back to a PIL image before inference.
    if not isinstance(orig_pil, Image.Image):
        orig_pil = _coerce_to_pil(orig_pil)

    coords_1024 = coords_1024 or []
    pred_mask, raw_mask, image_1024_with_star, latched_color, global_enc_image = _gpu_inference(
        orig_pil, coords_1024, text_prompt, fixed_color
    )
    if pred_mask is None:
        return None, None, f"Inference failed: {raw_mask}", None, None, None

    if len(coords_1024) == 0:
        final_prompt = text_prompt.replace("<COLOR>", "").replace("  ", " ").strip()
        status = f"REFERRING task (0 stars).\nPrompt: {final_prompt}"
    else:
        used_color = fixed_color or latched_color or "blue"
        final_prompt = text_prompt.replace("<COLOR>", used_color)
        status = f"STAR task with {len(coords_1024)} point(s). Color={used_color}.\nPrompt: {final_prompt}"

    mask_2d = pred_mask
    if mask_2d.ndim == 3:
        mask_2d = mask_2d[0] if mask_2d.shape[0] == 1 else mask_2d[..., 0]
    mask_bool = mask_2d > 0

    if clean_tensor_chw is not None:
        orig_1024_np = _tensor_chw_to_uint8_hwc(clean_tensor_chw)
        overlay_np = _create_cyan_overlay(orig_1024_np, mask_bool)
    else:
        overlay_np = None

    binary_np = (mask_bool.astype(np.uint8) * 255)
    coverage = 100.0 * float(mask_bool.mean())
    status = f"{status}\nSelected pixels: {coverage:.1f}% of the (1024x1024) image."
    return overlay_np, binary_np, status, pred_mask, image_1024_with_star, final_prompt


with gr.Blocks(title="MAOAM-GLaMM: Object & Material Selection") as demo:
    gr.Markdown(
        "# MAOAM: Unified Object & Material Selection with VLMs\n"
        "**Backend: GLaMM** = LLaVA-Llama (GranD pretrained) + SAM ViT-H.\n\n"
        "1) Upload an image. &nbsp; 2) Choose a **Selection type**. &nbsp; "
        "3) Follow the hint under the prompt (click star points and/or edit the text). &nbsp; "
        "4) Press **Submit**."
    )

    coords_state = gr.State([])
    fixed_color_state = gr.State(None)
    orig_pil_state = gr.State(None)
    clean_tensor_state = gr.State(None)
    disp_tensor_state = gr.State(None)
    last_pred_mask_state = gr.State(None)
    last_star_tensor_state = gr.State(None)
    last_prompt_state = gr.State(None)

    with gr.Row():
        with gr.Column(scale=1):
            input_image = gr.Image(label="Input / Click to add star(s)", type="numpy", height=400)
            selection_dropdown = gr.Dropdown(
                choices=list(_SELECTION_MODES.keys()),
                value=_DEFAULT_MODE,
                label="Selection type",
            )
            text_prompt = gr.Textbox(
                label="Text prompt",
                value=_SELECTION_MODES[_DEFAULT_MODE][0],
                lines=3,
            )
            hint_md = gr.Markdown(_SELECTION_MODES[_DEFAULT_MODE][1])
            with gr.Row():
                undo_btn = gr.Button("Undo last star", variant="secondary")
                clear_stars_btn = gr.Button("Clear stars", variant="secondary")
            submit_btn = gr.Button("Submit", variant="primary")
            download_btn = gr.Button("Download (original+mask+overlays)", variant="secondary")
            download_file = gr.File(label="Download zip")

        with gr.Column(scale=1):
            overlay_image = gr.Image(label="Overlaid Image", height=400)
            binary_mask_image = gr.Image(label="Binary Mask", height=400)
            status_text = gr.Textbox(
                label="Status",
                value="Upload an image, click up to 5 star points, then Submit.",
                interactive=False,
                lines=4,
            )
            coords_table = gr.Dataframe(
                headers=["h", "w"],
                datatype=["number", "number"],
                row_count=5,
                col_count=(2, "fixed"),
                interactive=False,
                label="Star coordinates (1024 space)",
            )

    def _load_example(image_np, _mode, _prompt):
        # gr.Examples click sets the input components but does NOT fire
        # input_image.upload, so the gr.State values (orig/clean/disp tensors)
        # would stay None and Submit would say "upload an image first".
        # Replicate the upload handler here to populate all state.
        return on_image_upload(image_np)

    gr.Examples(
        examples=[
            ["examples/living_room.jpg", "Object: text", "the two cats"],
            ["examples/living_room.jpg", "Material: text", "soft pink fabric"],
            ["examples/street.jpg", "Object: text", "the television screen"],
        ],
        inputs=[input_image, selection_dropdown, text_prompt],
        fn=_load_example,
        outputs=[input_image, coords_state, fixed_color_state, orig_pil_state,
                 clean_tensor_state, disp_tensor_state, status_text],
        run_on_click=True,
        cache_examples=False,
    )

    input_image.upload(
        on_image_upload,
        inputs=[input_image],
        outputs=[input_image, coords_state, fixed_color_state, orig_pil_state,
                 clean_tensor_state, disp_tensor_state, status_text],
    ).then(
        lambda coords: [[h, w] for (h, w) in (coords or [])],
        inputs=[coords_state], outputs=[coords_table],
    )

    input_image.select(
        on_click_add_star,
        inputs=[input_image, coords_state, fixed_color_state, disp_tensor_state],
        outputs=[input_image, coords_state, fixed_color_state, disp_tensor_state, status_text],
    ).then(
        lambda coords: [[h, w] for (h, w) in (coords or [])],
        inputs=[coords_state], outputs=[coords_table],
    )

    undo_btn.click(
        on_undo_last,
        inputs=[coords_state, fixed_color_state, clean_tensor_state],
        outputs=[input_image, coords_state, fixed_color_state, disp_tensor_state, status_text],
    ).then(
        lambda coords: [[h, w] for (h, w) in (coords or [])],
        inputs=[coords_state], outputs=[coords_table],
    )

    clear_stars_btn.click(
        on_clear_stars,
        inputs=[clean_tensor_state],
        outputs=[input_image, coords_state, fixed_color_state, disp_tensor_state, status_text],
    ).then(
        lambda coords: [[h, w] for (h, w) in (coords or [])],
        inputs=[coords_state], outputs=[coords_table],
    )

    submit_btn.click(
        on_submit,
        inputs=[orig_pil_state, coords_state, text_prompt, fixed_color_state, clean_tensor_state],
        outputs=[overlay_image, binary_mask_image, status_text,
                 last_pred_mask_state, last_star_tensor_state, last_prompt_state],
    )

    download_btn.click(
        _download_with_prompt,
        inputs=[orig_pil_state, clean_tensor_state, last_star_tensor_state,
                last_pred_mask_state, last_prompt_state],
        outputs=[download_file],
    )

    selection_dropdown.change(
        on_selection_change,
        inputs=[selection_dropdown],
        outputs=[text_prompt, hint_md],
    )

    gr.Markdown(
        "---\n"
        "**Paper:** MAOAM: Unified Object & Material Selection with Vision-Language Models "
        "(SIGGRAPH 2026). Project page: https://jadenpark0.github.io/project_pages/maoam/\n\n"
        "License: Adobe Research."
    )


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