"""
SAM 3.1 — Promptable Concept Segmentation demo
================================================
A live, language-driven segmentation demo built on Meta's Segment Anything Model 3.1.

Type a short noun phrase (e.g. "horse", "saddle"); the model finds and segments
*every* matching instance in the image. No boxes, no clicks, no retraining.

Model:   facebook/sam3.1  (image Promptable Concept Segmentation path)
Runtime: Hugging Face Spaces — works on a standard GPU Space or on ZeroGPU.

Deploy notes are in README.md (gated-model access + HF_TOKEN + hardware).
"""

import os
import time
import colorsys
from contextlib import nullcontext

import numpy as np
from PIL import Image, ImageDraw, ImageFont
import torch
import gradio as gr

# --------------------------------------------------------------------------------------
# ZeroGPU support (optional). On a standard GPU Space this becomes a transparent no-op.
# --------------------------------------------------------------------------------------
try:
    import spaces  # provided by the ZeroGPU runtime

    GPU = spaces.GPU
except Exception:  # not on Spaces / package missing → identity decorator

    def GPU(*args, **kwargs):
        # Supports both `@GPU` and `@GPU(duration=...)`
        if len(args) == 1 and callable(args[0]) and not kwargs:
            return args[0]

        def _deco(fn):
            return fn

        return _deco


# --------------------------------------------------------------------------------------
# Configuration
# --------------------------------------------------------------------------------------
MODEL_ID = os.environ.get("MODEL_ID", "facebook/sam3.1")
FALLBACK_MODEL_ID = os.environ.get("FALLBACK_MODEL_ID", "facebook/sam3")
HF_TOKEN = (
    os.environ.get("HF_TOKEN")
    or os.environ.get("HUGGING_FACE_HUB_TOKEN")
    or os.environ.get("HUGGINGFACE_TOKEN")
)

# ZeroGPU sets SPACES_ZERO_GPU; in that case CUDA is attached only inside @GPU calls,
# but we can still target "cuda" because the `spaces` runtime patches device placement.
_ZERO_GPU = bool(os.environ.get("SPACES_ZERO_GPU"))
DEVICE = "cuda" if (torch.cuda.is_available() or _ZERO_GPU) else "cpu"

EXAMPLE_PROMPTS = [
    "horse",
    "saddle",
    "person",
    "object used for riding control",
]

KEY_MESSAGE = "Segmentation is fully driven by language prompts — no retraining required."

# Lazily-loaded singletons
_MODEL = None
_PROCESSOR = None
_LOADED_ID = None


# --------------------------------------------------------------------------------------
# Model loading
# --------------------------------------------------------------------------------------
def load_model():
    """Load SAM 3.1 (image PCS) once. Falls back to SAM 3 if 3.1 is unavailable."""
    global _MODEL, _PROCESSOR, _LOADED_ID
    if _MODEL is not None:
        return

    from transformers import Sam3Model, Sam3Processor

    candidates = [MODEL_ID]
    if FALLBACK_MODEL_ID and FALLBACK_MODEL_ID != MODEL_ID:
        candidates.append(FALLBACK_MODEL_ID)

    last_err = None
    for mid in candidates:
        try:
            processor = Sam3Processor.from_pretrained(mid, token=HF_TOKEN)
            model = Sam3Model.from_pretrained(mid, token=HF_TOKEN)
            model.eval()
            try:
                model.to(DEVICE)
            except Exception:
                # On ZeroGPU the move is handled when the GPU is attached; ignore here.
                pass
            _MODEL, _PROCESSOR, _LOADED_ID = model, processor, mid
            if mid != MODEL_ID:
                print(f"[sam3.1-demo] '{MODEL_ID}' unavailable; loaded fallback '{mid}'.")
            else:
                print(f"[sam3.1-demo] Loaded '{mid}' on {DEVICE}.")
            return
        except Exception as e:  # try next candidate
            last_err = e
            print(f"[sam3.1-demo] Could not load '{mid}': {e}")

    raise RuntimeError(
        f"Failed to load any SAM 3 model from {candidates}. Last error: {last_err}"
    )


def _friendly_error(err: Exception) -> str:
    """Turn a load/inference exception into actionable guidance."""
    text = str(err).lower()
    gated = any(k in text for k in ["401", "403", "gated", "access", "token", "authorized"])
    if gated:
        return (
            "Couldn't access the model weights. SAM 3 / 3.1 are gated: request access on the "
            "Hugging Face model page, then add your token as a Space secret named "
            "<b>HF_TOKEN</b> (Settings → Variables and secrets), and restart the Space."
        )
    return f"Something went wrong while running the model: {err}"


# --------------------------------------------------------------------------------------
# Inference (GPU-scoped). Everything returned here is CPU/NumPy so it stays valid
# after the GPU is released (important for ZeroGPU).
# --------------------------------------------------------------------------------------
def _amp_ctx():
    """bfloat16 autocast on CUDA for speed; pass-through elsewhere."""
    if DEVICE == "cuda":
        return torch.autocast("cuda", dtype=torch.bfloat16)
    return nullcontext()


def _to_np(x):
    if x is None:
        return None
    if hasattr(x, "detach"):
        return x.detach().to("cpu").float().numpy()
    if isinstance(x, np.ndarray):
        return x
    if isinstance(x, (list, tuple)):
        if len(x) == 0:
            return np.zeros((0,))
        if hasattr(x[0], "detach"):
            return np.stack([t.detach().to("cpu").float().numpy() for t in x])
        return np.asarray(x)
    return np.asarray(x)


def _postprocess(outputs, target_sizes, threshold):
    res = _PROCESSOR.post_process_instance_segmentation(
        outputs,
        threshold=float(threshold),
        mask_threshold=0.5,
        target_sizes=target_sizes,
    )[0]
    masks = _to_np(res.get("masks"))
    boxes = _to_np(res.get("boxes"))
    scores = _to_np(res.get("scores"))
    return masks, boxes, scores


@GPU(duration=120)
def _infer_single(image: Image.Image, prompt: str, threshold: float):
    """Segment one text prompt on one image. Returns CPU arrays + metadata."""
    load_model()
    inputs = _PROCESSOR(images=image, text=prompt, return_tensors="pt").to(_MODEL.device)
    target_sizes = inputs["original_sizes"].tolist()

    t0 = time.perf_counter()
    with torch.no_grad():
        try:
            with _amp_ctx():
                outputs = _MODEL(**inputs)
        except RuntimeError:
            outputs = _MODEL(**inputs)  # rare: fall back to full precision
    if DEVICE == "cuda":
        torch.cuda.synchronize()
    ms = (time.perf_counter() - t0) * 1000.0

    masks, boxes, scores = _postprocess(outputs, target_sizes, threshold)
    return masks, boxes, scores, ms, _LOADED_ID, _MODEL.device.type


@GPU(duration=180)
def _infer_many(image: Image.Image, prompts, threshold: float):
    """Segment several prompts on one image, reusing vision features for speed.

    The whole batch runs inside a single GPU call, so the cached vision embeddings
    stay valid (safe on ZeroGPU).
    """
    load_model()
    img_inputs = _PROCESSOR(images=image, return_tensors="pt").to(_MODEL.device)
    target_sizes = img_inputs["original_sizes"].tolist()

    results = []
    t0 = time.perf_counter()
    with torch.no_grad():
        with _amp_ctx():
            vision_embeds = _MODEL.get_vision_features(pixel_values=img_inputs.pixel_values)
        for prompt in prompts:
            text_inputs = _PROCESSOR(text=prompt, return_tensors="pt").to(_MODEL.device)
            with _amp_ctx():
                outputs = _MODEL(vision_embeds=vision_embeds, **text_inputs)
            masks, boxes, scores = _postprocess(outputs, target_sizes, threshold)
            results.append((prompt, masks, boxes, scores))
    if DEVICE == "cuda":
        torch.cuda.synchronize()
    ms = (time.perf_counter() - t0) * 1000.0
    return results, ms, _LOADED_ID, _MODEL.device.type


# --------------------------------------------------------------------------------------
# Rendering (CPU). Builds the semi-transparent overlay and the mask-only view.
# --------------------------------------------------------------------------------------
def _palette(n: int):
    """Evenly-spaced, vivid colors (golden-ratio hue spacing) — one per instance."""
    cols = []
    for i in range(max(n, 1)):
        h = (i * 0.61803398875) % 1.0
        r, g, b = colorsys.hsv_to_rgb(h, 0.72, 1.0)
        cols.append((int(r * 255), int(g * 255), int(b * 255)))
    return cols


def _mask_list(masks, h, w):
    """Normalize whatever the model returned into a list of HxW bool arrays."""
    out = []
    if masks is None:
        return out
    arr = masks
    if arr.ndim == 2:
        arr = arr[None, ...]
    for i in range(arr.shape[0]):
        m = arr[i]
        if m.ndim == 3:
            m = m[0]
        m = m > 0.5
        if m.shape[:2] != (h, w):
            m = (
                np.asarray(
                    Image.fromarray((m.astype(np.uint8) * 255)).resize(
                        (w, h), Image.NEAREST
                    )
                )
                > 127
            )
        out.append(m)
    return out


def _boundary(mask: np.ndarray) -> np.ndarray:
    """1px boundary via 4-neighbour erosion (no SciPy/OpenCV dependency)."""
    e = mask.copy()
    e[1:, :] &= mask[:-1, :]
    e[:-1, :] &= mask[1:, :]
    e[:, 1:] &= mask[:, :-1]
    e[:, :-1] &= mask[:, 1:]
    return mask & ~e


def _font(size: int):
    for path in (
        "DejaVuSans-Bold.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans-Bold.ttf",
        "DejaVuSans.ttf",
        "/usr/share/fonts/truetype/dejavu/DejaVuSans.ttf",
    ):
        try:
            return ImageFont.truetype(path, size)
        except Exception:
            continue
    return ImageFont.load_default()


def render(image: Image.Image, masks, boxes, scores, prompt: str,
           alpha: float = 0.5, show_boxes: bool = False):
    """Return (overlay_image, mask_only_image)."""
    base = np.asarray(image.convert("RGB")).astype(np.float32)
    h, w = base.shape[:2]

    mlist = _mask_list(masks, h, w)
    cols = _palette(len(mlist))

    overlay = base.copy()
    mask_only = np.zeros_like(base)

    for i, m in enumerate(mlist):
        c = np.array(cols[i], dtype=np.float32)
        overlay[m] = overlay[m] * (1.0 - alpha) + c * alpha
        edge = _boundary(m)
        overlay[edge] = c                       # crisp instance outline
        mask_only[m] = c
        mask_only[edge] = np.minimum(c + 70, 255)

    overlay_img = Image.fromarray(overlay.clip(0, 255).astype(np.uint8))
    mask_only_img = Image.fromarray(mask_only.astype(np.uint8))

    if show_boxes and boxes is not None and len(boxes) and scores is not None:
        draw = ImageDraw.Draw(overlay_img, "RGBA")
        fsize = max(13, int(w / 55))
        font = _font(fsize)
        line_w = max(2, int(w / 480))
        for i in range(min(len(boxes), len(mlist) or len(boxes))):
            x1, y1, x2, y2 = [float(v) for v in boxes[i][:4]]
            c = cols[i % len(cols)]
            draw.rectangle([x1, y1, x2, y2], outline=c + (255,), width=line_w)
            label = f"{prompt} · {float(scores[i]):.2f}"
            tb = draw.textbbox((0, 0), label, font=font)
            tw, th = tb[2] - tb[0], tb[3] - tb[1]
            ty = max(0, y1 - th - 6)
            draw.rectangle([x1, ty, x1 + tw + 10, ty + th + 6], fill=c + (235,))
            draw.text((x1 + 5, ty + 3), label, fill=(20, 24, 31, 255), font=font)

    return overlay_img, mask_only_img


# --------------------------------------------------------------------------------------
# Status banner HTML
# --------------------------------------------------------------------------------------
def status_html(count: int, ms: float, model_id: str, device: str) -> str:
    plural = "" if count == 1 else "es"
    return (
        "<div class='status'>"
        f"<span class='chip chip-count'>{count} match{plural}</span>"
        f"<span class='chip chip-ms'>{ms:.0f} ms</span>"
        f"<span class='chip chip-dim'>{model_id} · {device}</span>"
        "</div>"
    )


def empty_status_html(prompt: str) -> str:
    return (
        "<div class='status'>"
        f"<span class='chip chip-empty'>No matches for &ldquo;{prompt}&rdquo;</span>"
        "<span class='chip chip-dim'>Try a simpler noun, or lower the threshold</span>"
        "</div>"
    )


def info_status_html(message: str) -> str:
    return f"<div class='status'><span class='chip chip-empty'>{message}</span></div>"


IDLE_STATUS = (
    "<div class='status'><span class='chip chip-dim'>"
    "Upload an image, type a prompt, then run.</span></div>"
)


# --------------------------------------------------------------------------------------
# Gradio callbacks
# --------------------------------------------------------------------------------------
def _noop(status_md, history):
    # leave images & gallery untouched
    return (gr.update(), gr.update(), gr.update(), status_md, gr.update(), history)


def run_single(image, prompt, threshold, show_boxes, history):
    history = history or []
    if image is None:
        return _noop(info_status_html("Upload an image to start."), history)
    prompt = (prompt or "").strip()
    if not prompt:
        return _noop(info_status_html("Type a prompt or pick an example."), history)

    try:
        masks, boxes, scores, ms, mid, dev = _infer_single(image, prompt, threshold)
    except Exception as e:
        return _noop(info_status_html(_friendly_error(e)), history)

    overlay, mask_only = render(image, masks, boxes, scores, prompt,
                                show_boxes=show_boxes)
    count = 0 if scores is None else int(len(scores))
    status = status_html(count, ms, mid, dev) if count else empty_status_html(prompt)

    history = ([(overlay, f"{prompt} · {count}")] + history)[:12]
    return overlay, mask_only, image, status, history, history


def run_many(image, multi_text, threshold, show_boxes, history):
    history = history or []
    if image is None:
        return _noop(info_status_html("Upload an image to start."), history)

    prompts, seen = [], set()
    for chunk in (multi_text or "").replace(",", "\n").splitlines():
        p = chunk.strip()
        if p and p.lower() not in seen:
            prompts.append(p)
            seen.add(p.lower())
    prompts = prompts[:6]
    if not prompts:
        return _noop(info_status_html("Add one prompt per line first."), history)

    try:
        results, ms, mid, dev = _infer_many(image, prompts, threshold)
    except Exception as e:
        return _noop(info_status_html(_friendly_error(e)), history)

    first_overlay = first_mask = None
    total = 0
    new_entries = []
    for idx, (prompt, masks, boxes, scores) in enumerate(results):
        overlay, mask_only = render(image, masks, boxes, scores, prompt,
                                    show_boxes=show_boxes)
        count = 0 if scores is None else int(len(scores))
        total += count
        new_entries.append((overlay, f"{prompt} · {count}"))
        if idx == 0:
            first_overlay, first_mask = overlay, mask_only

    history = (new_entries + history)[:12]
    status = (
        "<div class='status'>"
        f"<span class='chip chip-count'>{len(prompts)} prompts · {total} matches</span>"
        f"<span class='chip chip-ms'>{ms:.0f} ms total</span>"
        f"<span class='chip chip-dim'>{mid} · {dev} · vision features reused</span>"
        "</div>"
    )
    return first_overlay, first_mask, image, status, history, history


def fill_prompt(choice):
    return choice or ""


def reset_all():
    return (
        None,            # image
        "",              # prompt
        None,            # example dropdown
        None,            # overlay
        None,            # mask only
        None,            # original
        IDLE_STATUS,     # status
    )


def clear_history():
    return [], []


# --------------------------------------------------------------------------------------
# UI
# --------------------------------------------------------------------------------------
CSS = """
@import url('https://fonts.googleapis.com/css2?family=Space+Grotesk:wght@500;600;700&display=swap');

.gradio-container { max-width: 1280px !important; }

#app-header { display:flex; align-items:center; gap:.65rem; margin:.2rem 0 0; }
#app-header .logo {
  width:34px; height:34px; border-radius:9px;
  background:linear-gradient(135deg,#5145E5,#00B3A4);
  box-shadow:0 2px 10px rgba(81,69,229,.35);
}
#app-header h1 {
  font-family:'Space Grotesk', Inter, system-ui, sans-serif;
  font-size:1.55rem; font-weight:700; letter-spacing:-0.015em; margin:0;
}
#app-sub { color:#5B6472; margin:.15rem 0 0; font-size:.96rem; }

#key-banner {
  margin:.5rem 0 1rem; padding:.7rem 1rem; border-radius:12px; color:#fff;
  background:linear-gradient(90deg,#5145E5,#00B3A4); font-weight:600;
  display:flex; gap:.6rem; align-items:center; line-height:1.3;
}
#key-banner .dot {
  width:8px; height:8px; border-radius:50%; background:#fff;
  box-shadow:0 0 0 4px rgba(255,255,255,.28); flex:none;
}

.status { display:flex; gap:.4rem; flex-wrap:wrap; align-items:center; min-height:34px; }
.chip { font-size:.8rem; padding:.2rem .6rem; border-radius:999px; font-weight:600;
  white-space:nowrap; }
.chip-count { background:#ECEAFE; color:#3F33CF; }
.chip-ms    { background:#E1F6F2; color:#00897B; }
.chip-dim   { background:#F0F2F5; color:#5B6472; font-weight:500; }
.chip-empty { background:#FFF4E5; color:#B26A00; }

.fade img { animation: sam-fade .45s ease both; }
@keyframes sam-fade { from { opacity:0; transform:scale(.992); } to { opacity:1; transform:none; } }
@media (prefers-reduced-motion: reduce) { .fade img { animation:none; } }
"""

THEME = gr.themes.Soft(
    primary_hue=gr.themes.colors.indigo,
    secondary_hue=gr.themes.colors.teal,
    neutral_hue=gr.themes.colors.slate,
    font=[gr.themes.GoogleFont("Inter"), "system-ui", "sans-serif"],
)


def build_demo():
    with gr.Blocks(theme=THEME, css=CSS, title="SAM 3.1 · Concept Segmentation") as demo:
        history_state = gr.State([])

        gr.HTML(
            '<div id="app-header"><div class="logo"></div>'
            "<div><h1>SAM 3.1 · Concept Segmentation</h1></div></div>"
            '<p id="app-sub">Type what you want to find — the model segments every '
            "matching instance. No boxes, no clicks, no retraining.</p>"
        )
        gr.HTML(
            f'<div id="key-banner"><span class="dot"></span><span>{KEY_MESSAGE}</span></div>'
        )

        with gr.Row(equal_height=False):
            # ---------------- Inputs ----------------
            with gr.Column(scale=5, min_width=360):
                image_in = gr.Image(
                    type="pil",
                    label="Image",
                    sources=["upload", "clipboard"],
                    height=420,
                    elem_classes=["fade"],
                )
                prompt_tb = gr.Textbox(
                    label="Prompt",
                    placeholder="e.g. horse",
                    info="Short noun phrases work best, e.g. \u201chorse\u201d or \u201csaddle\u201d.",
                    autofocus=True,
                )
                example_dd = gr.Dropdown(
                    choices=EXAMPLE_PROMPTS,
                    label="Example prompts",
                    value=None,
                    interactive=True,
                )
                with gr.Row():
                    run_btn = gr.Button("Run segmentation", variant="primary", scale=3)
                    reset_btn = gr.Button("Reset", variant="secondary", scale=1)

                status = gr.HTML(IDLE_STATUS)

                with gr.Accordion("Advanced", open=False):
                    threshold = gr.Slider(
                        minimum=0.05, maximum=0.95, value=0.5, step=0.05,
                        label="Confidence threshold",
                        info="Lower to reveal more instances; higher to keep only strong matches.",
                    )
                    show_boxes = gr.Checkbox(
                        value=False, label="Show boxes & confidence scores"
                    )
                    gr.Markdown(
                        "**Multiple prompts** — one per line. They share a single vision "
                        "pass, so adding prompts is fast."
                    )
                    multi_tb = gr.Textbox(
                        label="Prompts (one per line)",
                        placeholder="horse\nsaddle\nperson",
                        lines=3,
                    )
                    run_many_btn = gr.Button("Run all prompts", variant="secondary")

            # ---------------- Outputs (the hero) ----------------
            with gr.Column(scale=7, min_width=420):
                with gr.Tabs():
                    with gr.Tab("Overlay"):
                        overlay_out = gr.Image(
                            label=None, height=540, interactive=False,
                            show_label=False, elem_classes=["fade"],
                        )
                    with gr.Tab("Mask only"):
                        mask_out = gr.Image(
                            label=None, height=540, interactive=False,
                            show_label=False, elem_classes=["fade"],
                        )
                    with gr.Tab("Original"):
                        original_out = gr.Image(
                            label=None, height=540, interactive=False,
                            show_label=False, elem_classes=["fade"],
                        )

        with gr.Accordion("Prompt history", open=False):
            history_gallery = gr.Gallery(
                label=None, show_label=False, columns=4, height=240,
                object_fit="cover", preview=False,
            )
            clear_btn = gr.Button("Clear history", variant="secondary", size="sm")

        # Optional bundled examples (image + prompt). Lights up only if files exist,
        # so the Space runs fine without any image assets checked in.
        ex_dir = "examples"
        ex_pairs = []
        if os.path.isdir(ex_dir):
            for fn, pr in [("horse.jpg", "horse"), ("street.jpg", "person"),
                           ("kitchen.jpg", "handle")]:
                p = os.path.join(ex_dir, fn)
                if os.path.exists(p):
                    ex_pairs.append([p, pr])
        if ex_pairs:
            gr.Examples(examples=ex_pairs, inputs=[image_in, prompt_tb],
                        label="Try an example")

        gr.Markdown(
            "<sub>Built on Meta's Segment Anything Model 3.1 (Promptable Concept "
            "Segmentation). SAM 3 / 3.1 weights are gated on Hugging Face. "
            "Very descriptive phrases are less reliable than short nouns — for the "
            "reins, \u201cbridle\u201d or \u201creins\u201d will usually beat "
            "\u201cobject used for riding control\u201d.</sub>"
        )

        # ----- wiring -----
        out_targets = [overlay_out, mask_out, original_out, status,
                       history_gallery, history_state]

        run_btn.click(
            run_single,
            inputs=[image_in, prompt_tb, threshold, show_boxes, history_state],
            outputs=out_targets,
        )
        prompt_tb.submit(
            run_single,
            inputs=[image_in, prompt_tb, threshold, show_boxes, history_state],
            outputs=out_targets,
        )
        run_many_btn.click(
            run_many,
            inputs=[image_in, multi_tb, threshold, show_boxes, history_state],
            outputs=out_targets,
        )
        example_dd.change(fill_prompt, inputs=example_dd, outputs=prompt_tb)
        reset_btn.click(
            reset_all,
            outputs=[image_in, prompt_tb, example_dd, overlay_out, mask_out,
                     original_out, status],
        )
        clear_btn.click(clear_history, outputs=[history_gallery, history_state])

    return demo


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