"""
Public HF Space — pure HTTP glue to the private Inference Endpoint.

This file is intentionally minimal. The real model code lives in the
private Endpoint repo at hf_endpoint/. Anyone browsing this Space sees
only the Gradio UI and `requests.post(...)` calls — no inference logic,
no checkpoints, no algorithm details.

Required Space secrets (Settings → Variables and secrets):
    ENDPOINT_URL   https://<your-endpoint>.endpoints.huggingface.cloud
    HF_API_TOKEN   read token with access to the endpoint
"""
# ── Monkey-patch gradio_client's broken JSON-schema serializer ────────────
# Several `gradio_client` releases (incl. the ones pinned by gradio 4.40+)
# crash on bare-bool JSON-Schemas (`True`/`False`, JSON-Schema for "any" /
# "never"). `ImagePrompter`'s API schema contains exactly one of those, so
# /api_info raises `TypeError: argument of type 'bool' is not iterable` and
# the Space container fails its health probe. Patch both code paths before
# importing gradio so they never crash.
import gradio_client.utils as _gc_utils

_orig_get_type = _gc_utils.get_type
def _safe_get_type(schema):
    if isinstance(schema, bool):
        return "Any"
    return _orig_get_type(schema)
_gc_utils.get_type = _safe_get_type

_orig_jsts_pt = _gc_utils._json_schema_to_python_type
def _safe_jsts_pt(schema, defs=None):
    if isinstance(schema, bool):
        return "Any"
    return _orig_jsts_pt(schema, defs)
_gc_utils._json_schema_to_python_type = _safe_jsts_pt
# ── end patch ──────────────────────────────────────────────────────────────

import base64
import io
import os
import time
from pathlib import Path

import gradio as gr
import requests
from gradio_image_prompter import ImagePrompter
from PIL import Image, ImageDraw

ENDPOINT = os.environ.get("ENDPOINT_URL", "").rstrip("/")
TOKEN    = os.environ.get("HF_API_TOKEN", "")
if not ENDPOINT:
    raise RuntimeError("ENDPOINT_URL must be set as a Space secret.")

HEADERS = {"Authorization": f"Bearer {TOKEN}"} if TOKEN else {}
TIMEOUT = 180  # seconds; covers cold-start + diffusion


# ---------- transport ----------

def pil_to_data_url(img: Image.Image) -> str:
    buf = io.BytesIO()
    img.convert("RGB").save(buf, format="PNG")
    return "data:image/png;base64," + base64.b64encode(buf.getvalue()).decode()


def data_url_to_pil(data_url: str) -> Image.Image:
    if "," in data_url:
        data_url = data_url.split(",", 1)[1]
    return Image.open(io.BytesIO(base64.b64decode(data_url))).convert("RGB")


def post(path: str, body: dict) -> dict:
    """The backend is a private Docker Space hosting a FastAPI app; we hit
    its real per-task paths (/api/segment, /api/depth_synth, /api/move)."""
    r = requests.post(f"{ENDPOINT}{path}", headers=HEADERS, json=body,
                      timeout=TIMEOUT)
    if not r.ok:
        print(f"[error] {path} {r.status_code}: {r.text[:500]}", flush=True)
        if r.status_code in (503, 504):
            raise gr.Error("Model is warming up (~45 s on first call). Try again in a moment.")
        raise gr.Error(f"Pipeline error ({r.status_code}). Please try again.")
    return r.json()


# ---------- visual helpers ----------

def draw_target_dot(img: Image.Image, tgt_pt, color=(229, 57, 53)):
    """Compose a red target dot at the given normalized (x, y) on a copy of img."""
    if img is None or tgt_pt is None:
        return img
    out = img.copy()
    draw = ImageDraw.Draw(out, "RGBA")
    W, H = out.size
    r = max(10, min(W, H) // 50)
    halo = r + 7
    x, y = int(tgt_pt[0] * W), int(tgt_pt[1] * H)
    draw.ellipse([x - halo, y - halo, x + halo, y + halo],
                 fill=(255, 255, 255, 80))
    draw.ellipse([x - r - 3, y - r - 3, x + r + 3, y + r + 3],
                 fill=(255, 255, 255, 255))
    draw.ellipse([x - r, y - r, x + r, y + r], fill=(*color, 255))
    return out


def make_mask_overlay(image_orig, mask_pil,
                      tint=(40, 220, 90), opacity=0.45):
    """Composite a tinted (default soft-green) layer over the masked region
    of the original image. Returns an RGB PIL image; the binary mask is kept
    separately in `mask_state` for the backend."""
    if image_orig is None or mask_pil is None:
        return None
    base = image_orig.convert("RGBA")
    W, H = base.size
    mask_l = mask_pil.convert("L").resize((W, H), Image.NEAREST)
    alpha = Image.eval(mask_l, lambda v: int(v * opacity))
    overlay = Image.new("RGBA", (W, H), (*tint, 0))
    overlay.putalpha(alpha)
    return Image.alpha_composite(base, overlay).convert("RGB")


# ---------- prompter <-> box helpers ----------
#
# ImagePrompter stores prompts as a `points` list where each entry is
# [x1, y1, label_top_left, x2, y2, label_bot_right]. Label codes:
#   0.0 negative point · 1.0 positive point ·
#   2.0 top-left box corner · 3.0 bottom-right box corner · 4.0 empty.
# A box is therefore a single entry like [x0, y0, 2.0, x1, y1, 3.0].

def extract_box_norm(prompter_value):
    """Return the LATEST box drawn in the prompter as a normalized
    (x0, y0, x1, y1) tuple in [0, 1], or None if no box present."""
    if not prompter_value:
        return None
    img = prompter_value.get("image")
    points = prompter_value.get("points") or []
    if img is None or not points:
        return None
    if hasattr(img, "size"):
        W, H = img.size
    else:
        H, W = img.shape[:2]
    for p in reversed(points):
        if len(p) >= 6 and int(p[2]) == 2 and int(p[5]) == 3:
            x_lo, x_hi = sorted((p[0], p[3]))
            y_lo, y_hi = sorted((p[1], p[4]))
            return (x_lo / W, y_lo / H, x_hi / W, y_hi / H)
    return None


def mask_centroid_norm(mask_pil):
    """Bounding-box center of the mask in normalized coords. Cheap proxy for
    the true centroid that avoids a numpy dependency at runtime."""
    bbox = mask_pil.getbbox()
    if not bbox:
        return 0.5, 0.5
    W, H = mask_pil.size
    return ((bbox[0] + bbox[2]) / 2 / W,
            (bbox[1] + bbox[3]) / 2 / H)


def _prompter_value(img):
    """Initial ImagePrompter value: just the image, no points yet."""
    return {"image": img, "points": []}


def keep_only_latest_box(prompter_value):
    """Trim the prompter's `points` list down to the most recent box AND
    drop any non-box entries (positive/negative point clicks). The source
    canvas is for box drawing only — point clicks are not part of our
    workflow and shouldn't visibly stick around."""
    if not prompter_value:
        return gr.update()
    points = prompter_value.get("points") or []
    boxes = [p for p in points
             if len(p) >= 6 and int(p[2]) == 2 and int(p[5]) == 3]
    # Already at most one entry and it's a box (or empty) — nothing to trim.
    if len(points) <= 1 and len(boxes) == len(points):
        return gr.update()
    new_points = [boxes[-1]] if boxes else []
    return {"image": prompter_value.get("image"), "points": new_points}


# ---------- example library ----------
#
# Curated image+mask pairs shipped with the Space (see ./examples/). When
# the user picks an example, the mask is already precomputed so they can
# skip the box-draw / segment step and jump straight to picking the target.
# They can still draw a new box in the source panel + Segment to override.

EXAMPLES_DIR  = Path(__file__).parent / "examples"
EXAMPLE_NAMES = ["000", "003", "005", "006", "007", "008", "009", "010", "011"]
EXAMPLE_THUMBS = [
    (str(EXAMPLES_DIR / f"{n}.png"), f"#{i}")
    for i, n in enumerate(EXAMPLE_NAMES)
]


def do_load_example(evt: gr.SelectData):
    """Returns the same 7-tuple as do_upload, but with the mask + overlay
    pre-filled and has_mask=True so the user can skip segmentation."""
    idx  = evt.index
    name = EXAMPLE_NAMES[idx]
    img  = Image.open(EXAMPLES_DIR / f"{name}.png").convert("RGB")
    mask = Image.open(EXAMPLES_DIR / f"{name}_mask.png").convert("L")
    overlay = make_mask_overlay(img, mask)
    status = (f"📚 Loaded example #{idx} ({img.size[0]}×{img.size[1]}) with a "
              "precomputed mask. **Click the mask panel** to set the target, "
              "tune ΔZ / Scale / prompt, then press **▶ Run**. (Or drag a new "
              "rectangle in the source panel and press ✂ Segment to override.)")
    return (img, _prompter_value(img), mask, overlay, status, True, None)


def do_upload(file):
    """Returns: image_orig, prompter_value, mask_state, mask_view,
                seg_status, has_mask, tgt_pt."""
    if file is None:
        return (None, None, None, None,
                "Upload an image to begin.", False, None)
    img = Image.open(file.name).convert("RGB")
    status = (f"Loaded ({img.size[0]}×{img.size[1]}). "
              "**Drag a rectangle on the source panel** around the object — "
              "you can re-drag to adjust — then press **✂ Segment**.")
    return (img, _prompter_value(img), None, None, status, False, None)


def do_segment(image_orig, prompter_value):
    """Run SAM2 with the latest box drawn in the prompter. Also write a
    trimmed prompter value back so the canvas shows only the box that was
    actually used (drops accumulated stale boxes from earlier drags)."""
    if image_orig is None:
        raise gr.Error("Upload an image first.")
    box = extract_box_norm(prompter_value)
    if box is None:
        raise gr.Error("Draw a rectangle on the source panel (drag from one "
                       "corner to the opposite).")
    if (box[2] - box[0]) < 0.005 or (box[3] - box[1]) < 0.005:
        raise gr.Error("Box is too small. Drag a larger rectangle.")
    j = post("/api/segment", {
        "image": pil_to_data_url(image_orig),
        "box":   list(box),
    })
    binary_mask = data_url_to_pil(j["mask"]).convert("L")
    overlay = make_mask_overlay(image_orig, binary_mask)
    status = (f"✅ Segmented ({j['mask_pixels']} mask px in {j['elapsed']} s). "
              "If you don't like the mask, **re-drag the rectangle** in the "
              "source panel and press ✂ Segment again. When happy, **click the "
              "mask panel** to set the target.")
    # Replay the prompter value with only the box we used, dropping any
    # older accumulated drags.
    trimmed = keep_only_latest_box(prompter_value)
    if isinstance(trimmed, dict):
        prompter_out = trimmed
    else:
        prompter_out = prompter_value
    # Reset the target whenever we regenerate the mask.
    return binary_mask, overlay, True, status, None, prompter_out


def do_redraw_box(image_orig):
    """Re-mask flow: hard-reset the prompter's canvas so previously drawn
    rectangles are physically wiped.

    Why a generator and not just `return {"image": img, "points": []}`?
    gradio_image_prompter never reads `value.points` back into its internal
    draw arrays — value flow is component → Gradio, not the reverse. So a
    one-shot return clears the BACKEND state but leaves the canvas painted
    with the old rectangle. Yielding `image=None` first forces the Svelte
    canvas to dismount (the internal `o`/`u` arrays are scoped to that
    instance), then the second yield re-mounts it fresh."""
    if image_orig is None:
        raise gr.Error("Upload an image first.")
    status = ("📐 Draw a new rectangle on the source panel, then press "
              "**✂ Segment object** to re-run SAM2.")
    # Step 1 — wipe the prompter entirely. The canvas component unmounts.
    yield (
        None,                # prompter — fully cleared (forces dismount)
        None,                # mask_state
        None,                # mask_view
        False,               # has_mask
        None,                # tgt_pt
        "Clearing prompter…",
    )
    # Brief pause so the dismount actually happens before we re-mount.
    time.sleep(0.15)
    # Step 2 — re-mount with the original image and an empty points list.
    yield (
        _prompter_value(image_orig),   # prompter — fresh canvas
        None,                          # mask_state
        None,                          # mask_view
        False,                         # has_mask
        None,                          # tgt_pt
        status,                        # seg_status
    )


def on_target_click(image_orig, mask_state, has_mask_flag,
                    evt: gr.SelectData):
    """Click on the mask panel → set the target, redraw with the red dot."""
    if image_orig is None:
        raise gr.Error("Upload an image first.")
    if not has_mask_flag or mask_state is None:
        raise gr.Error("Press ✂ Segment first to compute the object mask.")
    W, H = image_orig.size
    x_px, y_px = evt.index
    x_norm, y_norm = x_px / W, y_px / H
    new_tgt = (x_norm, y_norm)
    overlay = make_mask_overlay(image_orig, mask_state)
    rendered = draw_target_dot(overlay, new_tgt)
    status = (f"🎯 Target set at ({x_norm:.2f}, {y_norm:.2f}). "
              "Tune ΔZ / Scale / prompt and press **▶ Run**. "
              "Click again to retarget; press 🔄 Reset to start over.")
    return rendered, new_tgt, status


def do_reset():
    return (
        None,                          # image_orig
        None,                          # prompter (cleared)
        None,                          # mask_state
        None,                          # mask_view
        "Upload an image to begin.",   # seg_status
        False, None,                   # has_mask, tgt_pt
        None, None, None, None, None,  # 5 output panels
        "",                            # timings
    )


# ---------- final pipeline: depth synth + diffusion ----------
#
# Implemented as a generator so the 5 panels populate one-by-one in the
# order: source depth → background depth → target mask → target depth →
# final result. Gradio streams every `yield` to the client immediately.

REVEAL_DELAY = 0.5   # seconds between intermediate-image reveals


def do_move(image_orig, mask_state, tgt_pt, delta_z, scale, prompt):
    if image_orig is None:
        raise gr.Error("Upload an image first.")
    if mask_state is None:
        raise gr.Error("Press ✂ Segment first to compute the object mask.")
    if tgt_pt is None:
        raise gr.Error("Click the mask panel to set the target location.")

    W, H = image_orig.size
    src_cx, src_cy = mask_centroid_norm(mask_state)
    tx, ty = tgt_pt
    drag_x = (tx - src_cx) * W
    drag_y = (ty - src_cy) * H

    body_synth = {
        "image":  pil_to_data_url(image_orig),
        "mask":   pil_to_data_url(mask_state),
        "drag_x": drag_x, "drag_y": drag_y, "delta_z": delta_z,
    }

    # Output order: out_result, out_depth_src, out_depth_bg,
    #               out_mask_tgt, out_depth_tgt, timings
    yield (None, None, None, None, None, "Synthesizing depth…")

    j_synth = post("/api/depth_synth", body_synth)
    depth_src = data_url_to_pil(j_synth["depth_source"])
    depth_bg  = data_url_to_pil(j_synth["depth_background"])
    mask_tgt  = data_url_to_pil(j_synth["mask_target"])
    depth_tgt = data_url_to_pil(j_synth["depth_target"])

    yield (None, depth_src, None, None, None,
           "✅ Source depth (MoGe-2)")
    time.sleep(REVEAL_DELAY)
    yield (None, depth_src, depth_bg, None, None,
           "✅ Background depth (Laplacian inpaint)")
    time.sleep(REVEAL_DELAY)
    yield (None, depth_src, depth_bg, mask_tgt, None,
           "✅ Target mask (shifted by drag)")
    time.sleep(REVEAL_DELAY)
    yield (None, depth_src, depth_bg, mask_tgt, depth_tgt,
           "✅ Target depth (Z-buffer composite) — running diffusion pipeline (~30 s)…")

    body_move = {**body_synth, "scale": scale, "prompt": prompt}
    t0 = time.perf_counter()
    j_move = post("/api/move", body_move)
    total = time.perf_counter() - t0
    final_result = data_url_to_pil(j_move["result"])

    yield (final_result, depth_src, depth_bg, mask_tgt, depth_tgt,
           f"✅ Done · MoGe+synth {j_move['depth_seconds']}s · "
           f"pipe {j_move['pipe_seconds']}s ({j_move['steps']} steps) · "
           f"total {total:.1f}s")


# ---------- UI ----------

DARK_THEME = gr.themes.Default(
    primary_hue="blue",
    neutral_hue="slate",
).set(
    body_background_fill="#0d1117",
    body_text_color="#e6edf3",
    background_fill_primary="#161b22",
    background_fill_secondary="#0d1117",
    block_background_fill="#161b22",
    block_border_color="#30363d",
    border_color_primary="#30363d",
    input_background_fill="#0d1117",
    input_border_color="#30363d",
    button_primary_background_fill="#1f6feb",
    button_primary_text_color="#ffffff",
    button_secondary_background_fill="#21262d",
    button_secondary_text_color="#e6edf3",
)

# CSS belt-and-suspenders in case some Gradio components ignore the theme
# (Image/Markdown sometimes default to the light palette).
DARK_CSS = """
.gradio-container, body { background: #0d1117 !important; color: #e6edf3 !important; }
.prose, .prose * { color: #e6edf3 !important; }

/* Single-row, horizontally-scrollable examples gallery. */
#examples-gallery .grid-wrap {
    display: flex !important;
    flex-wrap: nowrap !important;
    overflow-x: auto !important;
    overflow-y: hidden !important;
    grid-template-columns: none !important;
    grid-template-rows: none !important;
    gap: 8px !important;
    padding-bottom: 6px;
}
#examples-gallery .grid-wrap > * {
    flex: 0 0 auto !important;
    width: 160px !important;
    height: 160px !important;
}
#examples-gallery .grid-wrap::-webkit-scrollbar { height: 8px; }
#examples-gallery .grid-wrap::-webkit-scrollbar-thumb {
    background: #30363d; border-radius: 4px;
}
#examples-gallery .grid-wrap::-webkit-scrollbar-thumb:hover { background: #484f58; }

"""

# ImagePrompter implicitly maps a no-drag click to a "positive point" and a
# click-and-drag to a "box" — there's no toolbar to hide. Two things to fix
# at the JS layer:
#  1) Click-without-drag would otherwise drop a dot. Capture-phase listeners
#     swallow the event so the prompter's `handle_draw_end` never sees it.
#  2) The prompter's Svelte canvas keeps its own internal point list, so a
#     server-side trim of `value.points` doesn't actually wipe earlier
#     rectangles. To enforce "one rectangle only", after the user's first
#     successful drag we set a per-canvas `boxDrawn` flag and swallow every
#     subsequent mousedown/mouseup. The user has to press 🔁 Re-mask (or
#     🔄 Reset, 📤 Upload, or pick a new example) to reset the flag.
PROMPTER_NO_DOTS_JS = r"""
() => {
    const DRAG_PX = 4;
    const states = new WeakMap();   // canvas -> per-canvas state

    function setupCanvas(canvas) {
        if (states.has(canvas)) return;
        const s = {boxDrawn: false, downX: 0, downY: 0, dragged: false};
        states.set(canvas, s);

        canvas.addEventListener('mousedown', (e) => {
            if (s.boxDrawn) {
                e.stopPropagation();
                e.stopImmediatePropagation();
                e.preventDefault();
                return;
            }
            s.downX = e.clientX;
            s.downY = e.clientY;
            s.dragged = false;
        }, true);

        canvas.addEventListener('mousemove', (e) => {
            if (Math.abs(e.clientX - s.downX) > DRAG_PX ||
                Math.abs(e.clientY - s.downY) > DRAG_PX) {
                s.dragged = true;
            }
        }, true);

        canvas.addEventListener('mouseup', (e) => {
            if (s.boxDrawn || !s.dragged) {
                e.stopPropagation();
                e.stopImmediatePropagation();
                e.preventDefault();
                return;
            }
            // Real drag finished: lock further drags until reset.
            s.boxDrawn = true;
        }, true);

        // Belt-and-suspenders: synthetic click + right-click never produce
        // useful prompts in our pipeline.
        const swallow = (e) => {
            e.stopPropagation();
            e.stopImmediatePropagation();
            e.preventDefault();
        };
        canvas.addEventListener('click', swallow, true);
        canvas.addEventListener('contextmenu', swallow, true);
    }

    function resetAll() {
        document
            .querySelectorAll('#source-prompter canvas')
            .forEach(c => {
                const s = states.get(c);
                if (s) s.boxDrawn = false;
            });
    }

    function scan() {
        document
            .querySelectorAll('#source-prompter canvas')
            .forEach(setupCanvas);
    }

    // Re-arm the "draw one rectangle" lock when the user clicks any of the
    // entry points that reset the prompter to a fresh image. setTimeout
    // delays the reset slightly so it runs *after* Gradio has finished
    // pushing the new prompter value to the canvas.
    document.addEventListener('click', (e) => {
        const btn = e.target.closest('button');
        if (btn) {
            const txt = btn.textContent || '';
            if (txt.includes('Re-mask') ||
                txt.includes('Reset')   ||
                txt.includes('Upload')) {
                setTimeout(resetAll, 150);
                return;
            }
        }
        if (e.target.closest('#examples-gallery .grid-wrap > *')) {
            setTimeout(resetAll, 150);
        }
    }, true);

    scan();
    const obs = new MutationObserver(scan);
    obs.observe(document.body, {childList: true, subtree: true});
}
"""

with gr.Blocks(title="Object Moving Demo", theme=DARK_THEME, css=DARK_CSS,
               js=PROMPTER_NO_DOTS_JS) as demo:
    gr.Markdown(
        "# Depth-Aware Object Moving — Live Demo\n"
        "**Pick an example below** (image + mask preloaded — jump to step 4), "
        "or **Upload** your own and follow all the steps:\n"
        "1) **Upload** an image · "
        "2) **Drag a rectangle** on the source panel around the object · "
        "3) Press **✂ Segment object** · "
        "4) **Click the mask panel** to set the target · "
        "5) Tune ΔZ / Scale / prompt · "
        "6) Press **▶ Run**.\n"
        "First request after a quiet period takes ~60–90 s (GPU cold-start)."
    )

    image_orig = gr.State(value=None)
    has_mask   = gr.State(value=False)
    tgt_point  = gr.State(value=None)
    # mask_state holds the BINARY mask returned by SAM2 — what the backend
    # needs. `mask_view` below is the user-facing overlay (RGB composite).
    mask_state = gr.State(value=None)

    with gr.Row():
        with gr.Column(scale=1):
            examples_gallery = gr.Gallery(
                value=EXAMPLE_THUMBS,
                label="📚 Examples — click any to load image + precomputed mask "
                      "(scroll →)",
                columns=len(EXAMPLE_THUMBS), rows=1, height=200,
                object_fit="cover",
                allow_preview=False,
                show_label=True,
                elem_id="examples-gallery",
            )
            upload_btn = gr.UploadButton("📤 Upload image",
                                          file_types=["image"],
                                          variant="primary")
            prompter = ImagePrompter(
                show_label=True,
                label="Source — drag a rectangle to mark the object "
                      "(press 🔁 Re-mask to start over)",
                type="pil",
                elem_id="source-prompter",
            )
            mask_view = gr.Image(
                label="Segmentation mask (green overlay). After ✂ Segment, "
                      "**click here to set the target**.",
                type="pil", interactive=False,
            )
            seg_status = gr.Markdown("Upload an image to begin.")

            delta_z = gr.Slider(-1.0, 1.0, 0.0, step=0.05, label="ΔZ (depth offset)")
            scale   = gr.Slider(0.6, 1.4, 1.0, step=0.02, label="Scale factor")
            prompt = gr.Textbox(
                value=("Move the object indicated by the mask to the new location, "
                       "keeping the rest of the scene unchanged."),
                lines=2, label="Prompt")

            with gr.Row():
                segment_btn = gr.Button("✂ Segment object", variant="secondary")
                redraw_btn  = gr.Button("🔁 Re-mask (draw new box)")
                run   = gr.Button("▶ Run move pipeline", variant="primary")
                reset = gr.Button("🔄 Reset")

        with gr.Column(scale=1):
            out_result = gr.Image(label="Final result", type="pil",
                                   interactive=False)
            with gr.Row():
                out_depth_src = gr.Image(label="Source depth (MoGe-2)",
                                          type="pil", interactive=False)
                out_depth_bg  = gr.Image(label="Background depth (Laplace)",
                                          type="pil", interactive=False)
            with gr.Row():
                out_mask_tgt  = gr.Image(label="Target mask (shifted)",
                                          type="pil", interactive=False)
                out_depth_tgt = gr.Image(label="Target depth (Z-buffer)",
                                          type="pil", interactive=False)
            timings = gr.Markdown("")

    upload_btn.upload(
        fn=do_upload,
        inputs=upload_btn,
        outputs=[image_orig, prompter, mask_state, mask_view, seg_status,
                 has_mask, tgt_point],
    )
    examples_gallery.select(
        fn=do_load_example,
        inputs=[],
        outputs=[image_orig, prompter, mask_state, mask_view, seg_status,
                 has_mask, tgt_point],
    )
    # Trim the prompter to a single box every time its value changes — a
    # second drag visually replaces the first instead of stacking on top.
    # `change` fires for programmatic updates too (upload / example load /
    # re-mask) but those return `gr.update()` (no-op) inside the handler.
    # `show_progress='hidden'` keeps the UI from flashing a loading state.
    prompter.change(
        fn=keep_only_latest_box,
        inputs=[prompter],
        outputs=[prompter],
        show_progress="hidden",
    )
    segment_btn.click(
        fn=do_segment,
        inputs=[image_orig, prompter],
        outputs=[mask_state, mask_view, has_mask, seg_status, tgt_point,
                 prompter],
    )
    redraw_btn.click(
        fn=do_redraw_box,
        inputs=[image_orig],
        outputs=[prompter, mask_state, mask_view, has_mask, tgt_point,
                 seg_status],
    )
    mask_view.select(
        fn=on_target_click,
        inputs=[image_orig, mask_state, has_mask],
        outputs=[mask_view, tgt_point, seg_status],
    )
    reset.click(
        fn=do_reset,
        inputs=[],
        outputs=[image_orig, prompter, mask_state, mask_view, seg_status,
                 has_mask, tgt_point,
                 out_result, out_depth_src, out_depth_bg,
                 out_mask_tgt, out_depth_tgt,
                 timings],
    )
    run.click(
        fn=do_move,
        inputs=[image_orig, mask_state, tgt_point, delta_z, scale, prompt],
        outputs=[out_result, out_depth_src, out_depth_bg,
                 out_mask_tgt, out_depth_tgt, timings],
    )

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