import base64
import gc
import io
import json
import os
import re

IS_ZERO_GPU = os.environ.get("SPACES_ZERO_GPU") is not None

# ZeroGPU requires `spaces` to be imported before any CUDA-initializing package (importing torch /
# diffusers touches torch.cuda), so this block must stay above the torch/diffusers imports below.
if IS_ZERO_GPU:
    import spaces
else:

    class _NoSpaces:
        @staticmethod
        def GPU(*args, **kwargs):
            if len(args) == 1 and callable(args[0]) and not kwargs:
                return args[0]

            def deco(fn):
                return fn

            return deco

    spaces = _NoSpaces()


import gradio as gr  # noqa: E402
import torch  # noqa: E402

from diffusers import ModularPipeline  # noqa: E402


# Captioning logic lives in the published modular block; the Space is just a UI + model host over it.
BLOCK_REPO = "OzzyGT/ideogram4_caption_blocks"

# Selectable checkpoints fed to the block (no on-the-fly quantization — the SDNQ repos are already quantized):
#   bf16 full ~15GB (ZeroGPU / >=24GB GPU), SDNQ 8-bit ~9GB, SDNQ 4-bit ~6GB (both fit a 16GB GPU).
MODELS = {
    "bf16 (full · google/gemma-4-E4B-it)": "google/gemma-4-E4B-it",
    "SDNQ 8-bit (OzzyGT)": "OzzyGT/gemma_4_E4B_it_sdnq_dynamic_8bit",
    "SDNQ 4-bit (OzzyGT)": "OzzyGT/gemma_4_E4B_it_sdnq_dynamic_4bit",
}
# bf16 on ZeroGPU (48GB, fastest); 8-bit locally (better quality than 4-bit, still fits a 16GB GPU).
DEFAULT_PRECISION = (
    "bf16 (full · google/gemma-4-E4B-it)" if IS_ZERO_GPU else "SDNQ 8-bit (OzzyGT)"
)

# Ideogram 4's native caption schema (from the official ideogram-oss/ideogram4 docs). Key ORDER matters for
# quality, bbox is [ymin, xmin, ymax, xmax] in 0-1000 normalized coords, colors are UPPERCASE #RRGGBB. The
# verbatim example anchors the exact format.
DEFAULT_INSTRUCTION = (
    "You are a captioner for the Ideogram 4 image model. Look at the image and output ONE JSON object (no prose, "
    "no markdown fences) that reconstructs it in Ideogram 4's native caption schema. Keep keys in the exact order "
    "shown.\n\n"
    "SCHEMA (keys in order):\n"
    "- high_level_description: one or two sentences summarizing the whole image (subject + medium/style).\n"
    "- style_description: an object with EXACTLY ONE of `photo` or `art_style`:\n"
    "    photo caption order:     aesthetics, lighting, photo, medium, color_palette\n"
    "    non-photo caption order: aesthetics, lighting, medium, art_style, color_palette\n"
    '    color_palette: optional, up to 16 UPPERCASE hex colors (e.g. "#1B1B2F").\n'
    "- compositional_deconstruction: {background, elements}.\n"
    "    background: the global setting/backdrop.\n"
    "    elements: list; keys in order per type:\n"
    "      obj:  type, bbox, desc, color_palette\n"
    "      text: type, bbox, text, desc, color_palette\n\n"
    "RULES:\n"
    "- bbox is [ymin, xmin, ymax, xmax] in 0-1000 normalized coordinates (origin top-left), independent of "
    "resolution. bbox and per-element color_palette are optional (<=5 colors per element).\n"
    "- Main subject first. A coherent subject (one person/animal/object) is exactly ONE element; its parts go in "
    "its `desc`. `desc` is identity-first and concrete (color, material, pose, distinguishing features).\n"
    "- Commit to one concrete value each (no 'various'/'or similar'). `text` elements carry the literal in-image "
    "characters, verbatim. Colors UPPERCASE #RRGGBB. Describe only what is visible; do not invent.\n\n"
    "EXAMPLE (format reference only):\n"
    '{"high_level_description":"A medium-shot photograph of Formula 1 driver Max Verstappen wearing his Red Bull '
    "Racing suit and cap, smiling as he holds his helmet and talks to a man in a white shirt and black vest at a "
    'race track.","style_description":{"aesthetics":"saturated primary colors, rule of thirds, joyful and '
    'triumphant","lighting":"overcast daylight, diffused, soft subtle shadows","photo":"shallow depth of field, '
    'sharp focus, eye-level, telephoto","medium":"photograph"},"compositional_deconstruction":{"background":"An '
    "out-of-focus racing paddock; blurred figures, a purple and white structure with a red F1 logo on the left, "
    'daylight, sky not visible.","elements":[{"type":"obj","bbox":[55,642,1000,937],"desc":"An older man in '
    "profile facing left, grey hair, fair skin, white long-sleeved button-down shirt under a navy blue quilted "
    'vest, a slight smile."}]}}'
)


DEVICE = (
    torch.device("cuda")
    if (IS_ZERO_GPU or torch.cuda.is_available())
    else torch.device("cpu")
)

# One pipeline resident at a time: ZeroGPU serves the bf16 default; locally the user picks one and switching reloads.
_current = {"id": None, "pipe": None}


def _load(checkpoint):
    if "sdnq" in checkpoint.lower():
        import sdnq  # noqa: F401  registers the SDNQ backend so prebuilt quantized checkpoints load

    pipe = ModularPipeline.from_pretrained(BLOCK_REPO, trust_remote_code=True)
    pipe.load_components(
        pretrained_model_name_or_path=checkpoint, torch_dtype=torch.bfloat16
    )
    pipe.to(DEVICE)
    return pipe


def get_pipe(precision):
    """Return the caption pipeline, (re)loading the selected checkpoint if it changed (one resident at a time)."""
    if IS_ZERO_GPU:
        precision = DEFAULT_PRECISION  # ZeroGPU serves the bf16 default only, whatever the caller passes
    checkpoint = MODELS[precision]
    if _current["id"] != checkpoint:
        _current.update(id=None, pipe=None)
        gc.collect()
        if torch.cuda.is_available():
            torch.cuda.empty_cache()
        _current.update(id=checkpoint, pipe=_load(checkpoint))
        print(f"loaded {checkpoint}")
    return _current["pipe"]


# Two paths:
#   ZeroGPU  -> the platform hands out the GPU only for the duration of a @spaces.GPU call and takes
#              it right back, so there's nothing to "hog". We just preload bf16 at module scope (the
#              canonical ZeroGPU pattern) and skip the whole load/unload + quantization machinery.
#   Local    -> we hold a real GPU, so this support app must not squat on VRAM while the user runs the
#              Ideogram generation app. Nothing loads until the Load button; Unload frees it again.
if IS_ZERO_GPU:
    get_pipe(DEFAULT_PRECISION)


@spaces.GPU(duration=120)
def load_model(precision):
    """Load the selected model into VRAM (local only; nothing is resident until this runs)."""
    get_pipe(precision)
    return f"✅ Loaded: {precision}"


def unload_model():
    """Evict the resident model from VRAM and RAM so it stops hogging the GPU (local only)."""
    _current.update(id=None, pipe=None)
    gc.collect()
    if torch.cuda.is_available():
        torch.cuda.empty_cache()
    return "⚪ No model loaded"


def _parse_json(text: str):
    """Parse the model output; return (dict|None, repaired: bool).

    repaired=True means the JSON was incomplete (unbalanced brackets) and we closed it to parse —
    a sign the output was cut short. Robust to markdown fences, stray text/braces, and control
    chars in strings (strict=False).
    """
    text = text.strip()
    text = re.sub(r"^```(?:json)?\s*", "", text)
    text = re.sub(r"\s*```$", "", text).strip()
    start = text.find("{")
    if start == -1:
        return None, False
    frag = text[start:]
    # 1) first '{' .. last '}' (handles the common clean case + trailing prose without braces)
    try:
        obj = json.loads(frag[: frag.rfind("}") + 1], strict=False)
        if isinstance(obj, dict):
            return obj, False
    except json.JSONDecodeError:
        pass
    # 2) decode just the first complete object, ignoring any trailing content (stray braces, prose)
    try:
        obj, _ = json.JSONDecoder(strict=False).raw_decode(frag)
        if isinstance(obj, dict):
            return obj, False
    except json.JSONDecodeError:
        pass
    # 3) repair truncated/short output (e.g. a missing final '}') by balancing brackets, then parse
    try:
        obj = json.loads(_balance_json(frag), strict=False)
        return (obj, True) if isinstance(obj, dict) else (None, False)
    except json.JSONDecodeError:
        return None, False


def _balance_json(frag: str) -> str:
    """Close unbalanced strings/brackets in a truncated JSON fragment so it can be parsed."""
    stack, in_str, esc = [], False, False
    for ch in frag:
        if in_str:
            if esc:
                esc = False
            elif ch == "\\":
                esc = True
            elif ch == '"':
                in_str = False
        elif ch == '"':
            in_str = True
        elif ch in "{[":
            stack.append(ch)
        elif ch == "}" and stack and stack[-1] == "{":
            stack.pop()
        elif ch == "]" and stack and stack[-1] == "[":
            stack.pop()
    out = frag
    if in_str:
        out += '"'  # close a string cut off mid-value
    out = out.rstrip()
    if out.endswith(","):
        out = out[:-1]  # drop a dangling comma before the closers
    for ch in reversed(stack):
        out += "}" if ch == "{" else "]"
    return out


def _parse_palette(text: str) -> list:
    """Parse a free-form list of hex colors into normalized UPPERCASE #RRGGBB (max 16)."""
    out = []
    for tok in re.split(r"[,\s]+", (text or "").strip()):
        tok = tok.strip().upper().lstrip("#")
        if re.fullmatch(r"[0-9A-F]{6}", tok):
            out.append("#" + tok)
    return out[:16]


_push_counter = {"v": 0}


def _nonce() -> int:
    _push_counter["v"] += 1
    return _push_counter["v"]


def _data_url(image, max_side=1400) -> str:
    """Downscaled JPEG data URL for the canvas. Full-res PNG of a real photo can be tens of MB,
    which bloats the #bbox-in payload and makes the JS JSON.parse fail (so nothing renders). A
    display-sized JPEG keeps it small; bboxes are 0-1000 normalized so resolution doesn't matter."""
    img = image.convert("RGB")
    w, h = img.size
    if max(w, h) > max_side:
        scale = max_side / max(w, h)
        img = img.resize((max(1, round(w * scale)), max(1, round(h * scale))))
    buf = io.BytesIO()
    img.save(buf, format="JPEG", quality=85)
    return "data:image/jpeg;base64," + base64.b64encode(buf.getvalue()).decode()


def _ensure_ids(state):
    """Give every element a stable _id (used to reconcile canvas edits back to elements)."""
    cd = state.get("compositional_deconstruction") or {}
    els = cd.get("elements") or []
    n = state.get("_next_id", 0)
    for el in els:
        if not el.get("_id"):
            el["_id"] = f"e{n}"
            n += 1
    state["_next_id"] = n
    return state


def _elements(state) -> list:
    cd = (state or {}).get("compositional_deconstruction") or {}
    return cd.get("elements") or []


# ---- JSON export (schema-clean, ordered) ---------------------------------------------------------
def _clean_element(el: dict) -> dict:
    """Serialize one element in Ideogram4 key order, dropping empty optionals."""
    out = {"type": el.get("type", "obj")}
    if el.get("bbox"):
        out["bbox"] = el["bbox"]
    if el.get("type") == "text":
        out["text"] = el.get("text", "")
        if el.get("desc"):
            out["desc"] = el["desc"]
    else:
        if el.get("desc"):
            out["desc"] = el["desc"]
    if el.get("color_palette"):
        out["color_palette"] = el["color_palette"]
    return out


def _export(state) -> dict:
    """Normalize the working state into the clean Ideogram4 JSON object."""
    state = state or {}
    out = {}
    if state.get("high_level_description"):
        out["high_level_description"] = state["high_level_description"]
    if state.get("style_description"):
        out["style_description"] = state["style_description"]
    cd = state.get("compositional_deconstruction") or {}
    cdd = {}
    if cd.get("background"):
        cdd["background"] = cd["background"]
    cdd["elements"] = [_clean_element(e) for e in (cd.get("elements") or [])]
    out["compositional_deconstruction"] = cdd
    return out


def _dumps(state) -> str:
    return json.dumps(_export(state), indent=2, ensure_ascii=False)


def canvas_payload(image, state, include_image=True) -> str:
    """Build the JSON payload for the JS editor (#bbox-in): the source image + the full element list
    (type / bbox / desc / text / palette), each keyed by a stable id. The JS renders both the boxes
    and the right-side cards from this. include_image=False = elements only (no image reload)."""
    _ensure_ids(state)
    elements = []
    for el in _elements(state):
        bbox = el.get("bbox")
        if not (isinstance(bbox, (list, tuple)) and len(bbox) == 4):
            bbox = [350, 350, 650, 650]  # give box-less elements a default so they're editable
        ymin, xmin, ymax, xmax = bbox
        elements.append(
            {
                "id": el["_id"],
                "type": el.get("type", "obj"),
                "ymin": ymin, "xmin": xmin, "ymax": ymax, "xmax": xmax,
                "desc": el.get("desc", ""),
                "text": el.get("text", ""),
                "colors": el.get("color_palette") or [],
            }
        )
    payload = {"nonce": _nonce(), "elements": elements}
    payload["style_colors"] = (state.get("style_description") or {}).get("color_palette") or []
    payload["image"] = _data_url(image) if (include_image and image is not None) else None
    return json.dumps(payload)


def reconcile_canvas(state, out_json):
    """Fold JS editor edits (#bbox-out) back into element state by stable id, preserving top-level
    fields. The JS owns the element list, so this just rebuilds compositional_deconstruction.elements
    and returns the new state (no re-render signalling needed — the JS already drew everything)."""
    try:
        data = json.loads(out_json or "{}")
    except json.JSONDecodeError:
        return state
    out_els = data.get("elements") or []
    by_id = {el.get("_id"): el for el in _elements(state) if el.get("_id")}
    new_elems = []
    for oe in out_els:
        el = dict(by_id.get(oe.get("id")) or {})
        el["_id"] = oe.get("id")
        el["type"] = oe.get("type", "obj")
        el["bbox"] = oe.get("bbox")
        el["desc"] = oe.get("desc", "")
        el["text"] = oe.get("text", "")
        el["color_palette"] = _parse_palette(" ".join(oe.get("colors") or []))
        new_elems.append(el)
    new_state = dict(state or {})
    # Style palette is edited on the canvas side too — fold it back, preserving other style fields.
    if "style_colors" in data:
        style = dict(new_state.get("style_description") or {})
        colors = _parse_palette(" ".join(data.get("style_colors") or []))
        if colors:
            style["color_palette"] = colors
        else:
            style.pop("color_palette", None)
        if style:
            new_state["style_description"] = style
        else:
            new_state.pop("style_description", None)
    cd = dict(new_state.get("compositional_deconstruction") or {})
    cd["elements"] = new_elems
    cd.setdefault("background", "")
    new_state["compositional_deconstruction"] = cd
    return new_state


def sync_from_state(image, state):
    """From the caption state: return (state-with-ids, top fields..., canvas payload)."""
    state = dict(state or {})
    cd = dict(state.get("compositional_deconstruction") or {})
    cd["elements"] = [dict(e) for e in (cd.get("elements") or [])]
    state["compositional_deconstruction"] = cd
    _ensure_ids(state)
    style = state.get("style_description") or {}
    if "art_style" in style and "photo" not in style:
        kind, kind_val = "art_style", style.get("art_style", "")
    else:
        kind, kind_val = "photo", style.get("photo", "")
    background = cd.get("background", "")
    return (
        state,
        state.get("high_level_description", ""),
        background,
        kind,
        style.get("aesthetics", ""),
        style.get("lighting", ""),
        style.get("medium", ""),
        kind_val,
        canvas_payload(image, state, include_image=True),
    )


def rebuild_from_fields(
    state, hld, background, kind, aesthetics, lighting, medium, kind_val
):
    """Write the non-label fields back into the caption state (preserving elements + the JS-managed
    color palette) and refresh JSON."""
    state = dict(state or {})
    state["high_level_description"] = hld
    # Key order matters for quality: photo -> aesthetics, lighting, photo, medium, color_palette;
    # non-photo -> aesthetics, lighting, medium, art_style, color_palette.
    colors = (state.get("style_description") or {}).get("color_palette") or []  # owned by the canvas
    style = {}
    if aesthetics:
        style["aesthetics"] = aesthetics
    if lighting:
        style["lighting"] = lighting
    if kind == "photo":
        if kind_val:
            style["photo"] = kind_val
        if medium:
            style["medium"] = medium
    else:
        if medium:
            style["medium"] = medium
        if kind_val:
            style["art_style"] = kind_val
    if colors:
        style["color_palette"] = colors
    if style:
        state["style_description"] = style
    else:
        state.pop("style_description", None)
    cd = dict(state.get("compositional_deconstruction") or {})
    cd["background"] = background
    cd.setdefault("elements", [])
    state["compositional_deconstruction"] = cd
    return state, _dumps(state)


def apply_json(image, text):
    """Parse the raw-JSON tab back into state, refreshing every field + the canvas."""
    data, _ = _parse_json(text)
    if data is None:
        raise gr.Error("Invalid JSON — could not parse.")
    return sync_from_state(image, data)  # (state, fields..., payload)


@spaces.GPU(duration=120)
def caption(
    image,
    instruction,
    precision,
    max_new_tokens,
    temperature,
):
    if image is None:
        raise gr.Error("An image is required.")

    pipe = get_pipe(precision)
    text = pipe(
        image=image,
        instruction=instruction,
        max_new_tokens=int(max_new_tokens),
        temperature=float(temperature),
        output="caption_raw",
    )

    data, repaired = _parse_json(text)
    if data is None:
        return {}, text, (
            "⚠️ Couldn't parse the model output as JSON — it may have been truncated at the token "
            "limit. Raise **Max new tokens** and re-run, or fix it in the JSON tab and click Apply JSON."
        )
    # Ended naturally but dropped trailing brackets = complete content we just closed.
    status = (
        "✅ Caption ready (minor auto-repair of the closing brackets)."
        if repaired
        else "✅ Caption ready."
    )
    return data, _dumps(data), status


CSS = """
#caption-btn {
    background: linear-gradient(180deg, #2563eb 0%, #1e3a8a 100%) !important;
    color: #fff !important; border: none !important;
    box-shadow: inset 0 1px 0 rgba(255,255,255,0.15), 0 1px 2px rgba(0,0,0,0.25) !important;
}
#load-btn {
    background: linear-gradient(180deg, #15803d 0%, #14532d 100%) !important;
    color: #fff !important; border: none !important;
    box-shadow: inset 0 1px 0 rgba(255,255,255,0.15), 0 1px 2px rgba(0,0,0,0.25) !important;
}
#unload-btn {
    background: linear-gradient(180deg, #b91c1c 0%, #7f1d1d 100%) !important;
    color: #fff !important; border: none !important;
    box-shadow: inset 0 1px 0 rgba(255,255,255,0.15), 0 1px 2px rgba(0,0,0,0.25) !important;
}
#caption-btn, #load-btn, #unload-btn { min-height: 56px !important; height: 56px !important; }

/* ---- Scrollless, full-viewport layout: the page never scrolls; only the fields column does. ---- */
html, body { height: 100%; margin: 0; overflow: hidden; }
gradio-app { height: var(--app-h, 100vh); overflow: hidden; }
.gradio-container {
    height: var(--app-h, 100vh) !important;
    max-width: 100% !important;
    padding: 6px 14px 4px !important;
    overflow: hidden !important;
}
#app-title { margin: 0 0 2px !important; }
#app-title h1 { font-size: 1.2rem !important; margin: 0 !important; }
#app-title small {
    font-size: 0.68rem; line-height: 1.35;
    color: var(--body-text-color-subdued, #9aa0ac);
}
/* The right-hand editor column is the only thing that scrolls (like the desktop app). */
#fields-col {
    max-height: calc(var(--app-h, 100vh) - 118px);
    overflow-y: auto;
    overflow-x: hidden;
    padding-right: 6px;
}
#canvas-col { min-height: 0; }

/* Theme tokens for the custom editor — light by default, dark when Gradio adds .dark to <body>
   (inherited by the body-appended popover too, so everything follows the HF light/dark toggle). */
:root {
    --ib-input-bg: #ffffff; --ib-input-fg: #1a1b25; --ib-border: #d7dae1;
    --ib-panel-bg: #ffffff; --ib-card-bg: #f6f7f9; --ib-eye-bg: #eef0f4; --ib-subdued: #6b7280;
}
.dark {
    --ib-input-bg: #15181f; --ib-input-fg: #e6e6e6; --ib-border: #3a3f4b;
    --ib-panel-bg: #1b1e25; --ib-card-bg: #1f2229; --ib-eye-bg: #2a2f3a; --ib-subdued: #9aa0ac;
}

/* ---- Hand-rolled bbox canvas ---- */
#bbox-root {
    width: 100%;
    height: calc(var(--app-h, 100vh) - 120px);
    display: flex;
    align-items: center;
    justify-content: center;
    overflow: hidden;
    touch-action: none;
}
.ib-hint { color: var(--ib-subdued); font-size: 0.9rem; padding: 12px; }
.ib-stage { position: relative; display: inline-block; max-width: 100%; max-height: 100%; line-height: 0; }
.ib-img { display: block; max-width: 100%; max-height: calc(var(--app-h, 100vh) - 120px); width: auto; height: auto; user-select: none; }
.ib-box {
    position: absolute;
    border: 2px solid #4da460;
    background: rgba(77, 164, 96, 0.15);
    box-sizing: border-box;
    cursor: move;
}
.ib-box.sel { border-color: #ff8c00; background: rgba(255, 140, 0, 0.18); z-index: 5; }
/* eyedrop mode: crosshair over the image, and boxes must not intercept the sample click */
.ib-stage.ib-eyedrop { cursor: crosshair !important; }
.ib-stage.ib-eyedrop .ib-box { pointer-events: none !important; }
.ib-label {
    position: absolute; top: 0; left: 0;
    font: bold 12px/1.5 sans-serif; color: #fff;
    background: rgba(0, 0, 0, 0.6); padding: 0 4px;
    white-space: nowrap; pointer-events: none;
}
.ib-h { position: absolute; width: 11px; height: 11px; background: #fff; border: 1px solid #ff8c00; box-sizing: border-box; display: none; }
.ib-box.sel .ib-h { display: block; }
.ib-nw { left: -6px; top: -6px; cursor: nwse-resize; }
.ib-ne { right: -6px; top: -6px; cursor: nesw-resize; }
.ib-sw { left: -6px; bottom: -6px; cursor: nesw-resize; }
.ib-se { right: -6px; bottom: -6px; cursor: nwse-resize; }

/* group-transform box (shown when 2+ boxes selected): outline is click-through, handles aren't */
.ib-group {
    position: absolute; border: 1px dashed #ff8c00; box-sizing: border-box;
    pointer-events: none; z-index: 6;
}
.ib-gh {
    position: absolute; width: 11px; height: 11px; background: #fff; border: 1px solid #ff8c00;
    box-sizing: border-box; pointer-events: auto;
}
.ib-gh.ib-n { left: 50%; top: -6px; transform: translateX(-50%); cursor: ns-resize; }
.ib-gh.ib-s { left: 50%; bottom: -6px; transform: translateX(-50%); cursor: ns-resize; }
.ib-gh.ib-e { right: -6px; top: 50%; transform: translateY(-50%); cursor: ew-resize; }
.ib-gh.ib-w { left: -6px; top: 50%; transform: translateY(-50%); cursor: ew-resize; }

/* ---- Right-side element cards (rendered by JS into #panel-root) ---- */
.ib-panel-head { display: flex; align-items: center; gap: 8px; margin: 4px 0 6px; }
.ib-add-wrap { margin-left: auto; display: flex; gap: 6px; align-items: center; }
/* !important throughout: Gradio's global button/select styles otherwise win on specificity. */
.ib-add {
    display: inline-flex !important; align-items: center !important; justify-content: center !important;
    background: linear-gradient(180deg, #15803d 0%, #14532d 100%) !important;
    color: #fff !important; border: none !important; border-radius: 6px !important;
    padding: 0 10px !important; font-size: 0.75rem !important; line-height: 1 !important;
    height: 26px !important; min-width: 0 !important; box-shadow: none !important;
    margin: 0 !important; box-sizing: border-box !important; align-self: center !important;
    vertical-align: middle !important; cursor: pointer;
}
.ib-add:hover { filter: brightness(1.12); }
#panel-root { display: flex; flex-direction: column; gap: 6px; }
.ib-empty { color: var(--ib-subdued); font-size: 0.85rem; padding: 8px 2px; }
.ib-card {
    border: 1px solid var(--ib-border);
    border-radius: 8px; padding: 6px 8px; background: var(--ib-card-bg);
    display: flex; flex-direction: column; gap: 5px;
}
.ib-card.sel { border-color: #ff8c00; box-shadow: 0 0 0 1px #ff8c00 inset; }
.ib-card-head { display: flex; align-items: center; gap: 6px; }
.ib-badge {
    display: inline-flex; align-items: center; justify-content: center;
    width: 19px; height: 19px; border-radius: 50%; background: #4da460; color: #fff;
    font-size: 0.7rem; font-weight: 700; flex: 0 0 auto;
}
/* Native <select> ignores theme bg unless appearance:none — that was the white box. */
.ib-type {
    appearance: none !important; -webkit-appearance: none !important; -moz-appearance: none !important;
    background-color: var(--ib-input-bg) !important;
    background-image: url('data:image/svg+xml;utf8,<svg xmlns="http://www.w3.org/2000/svg" width="10" height="10" viewBox="0 0 10 10"><path d="M1 3l4 4 4-4" fill="none" stroke="%23888" stroke-width="1.5"/></svg>') !important;
    background-repeat: no-repeat !important; background-position: right 6px center !important;
    color: var(--ib-input-fg) !important; border: 1px solid var(--ib-border) !important; border-radius: 6px !important;
    padding: 2px 20px 2px 8px !important; font-size: 0.76rem !important; height: 24px !important;
    min-width: 0 !important; width: 74px !important; box-shadow: none !important; cursor: pointer;
}
.ib-type option { background: var(--ib-input-bg); color: var(--ib-input-fg); }
.ib-del {
    margin-left: auto; background: #b91c1c !important; color: #fff !important; border: none !important;
    border-radius: 5px !important; width: 22px !important; height: 22px !important; cursor: pointer;
    font-size: 0.72rem !important; line-height: 1 !important; min-width: 0 !important;
    padding: 0 !important; box-shadow: none !important;
    display: inline-flex !important; align-items: center; justify-content: center; flex: 0 0 auto;
}
.ib-del:hover { filter: brightness(1.15); }
.ib-f {
    width: 100%; box-sizing: border-box; background: var(--ib-input-bg);
    color: var(--ib-input-fg); border: 1px solid var(--ib-border);
    border-radius: 6px; padding: 5px 7px; font-size: 0.8rem; font-family: inherit;
}
.ib-desc { resize: vertical; min-height: 38px; }
.ib-f:focus { outline: none; border-color: #4da460; }
.ib-f::placeholder { color: #6b7280; }

/* ---- color palette chips (swatch tags + picker). !important: bare <button>s still pick up
   Gradio's base button sizing, which is what was bloating the chips. ---- */
.ib-pal { display: flex; flex-wrap: wrap; gap: 4px; align-items: center; }
.ib-pal-label { font-size: 0.8rem; color: var(--ib-subdued); margin: 2px 0; }
#style-pal-root { min-height: 20px; }
.ib-chip {
    display: inline-flex !important; align-items: center; gap: 2px;
    border-radius: 4px; padding: 0 3px 0 6px; font-size: 0.65rem !important; font-weight: 600;
    border: 1px solid rgba(0, 0, 0, 0.35); cursor: pointer; height: 18px; box-sizing: border-box;
}
.ib-chip-hex { letter-spacing: 0.2px; }
.ib-chip-x {
    background: none !important; border: none !important; box-shadow: none !important;
    color: inherit !important; cursor: pointer; opacity: 0.7;
    font-size: 0.8rem !important; line-height: 1 !important;
    padding: 0 1px !important; margin: 0 !important;
    min-width: 0 !important; width: auto !important; height: auto !important;
}
.ib-chip-x:hover { opacity: 1; }
.ib-chip-add {
    display: inline-flex !important; align-items: center !important;
    background: none !important; color: var(--ib-subdued) !important; cursor: pointer;
    border: 1px dashed var(--ib-border) !important; border-radius: 4px !important;
    padding: 0 7px !important; font-size: 0.65rem !important; line-height: 1 !important;
    height: 18px !important; min-width: 0 !important; box-shadow: none !important; margin: 0 !important;
}
.ib-chip-add:hover { color: var(--ib-input-fg) !important; border-color: var(--ib-subdued) !important; }

/* ---- inline color-picker popover ---- */
.ib-cp {
    position: fixed; z-index: 9999; width: 188px; box-sizing: border-box;
    background: var(--ib-panel-bg); border: 1px solid var(--ib-border); border-radius: 8px; padding: 8px;
    box-shadow: 0 8px 26px rgba(0, 0, 0, 0.55); display: flex; flex-direction: column; gap: 8px;
    user-select: none;
}
.ib-cp-sv { position: relative; width: 100%; height: 120px; border-radius: 5px; cursor: crosshair; }
.ib-cp-svknob {
    position: absolute; width: 12px; height: 12px; border-radius: 50%; border: 2px solid #fff;
    box-shadow: 0 0 0 1px rgba(0, 0, 0, 0.6); transform: translate(-50%, -50%); pointer-events: none;
}
.ib-cp-hue {
    position: relative; width: 100%; height: 14px; border-radius: 7px; cursor: pointer;
    background: linear-gradient(to right, #f00, #ff0, #0f0, #0ff, #00f, #f0f, #f00);
}
.ib-cp-hueknob {
    position: absolute; top: 50%; width: 14px; height: 14px; border-radius: 50%; border: 2px solid #fff;
    box-shadow: 0 0 0 1px rgba(0, 0, 0, 0.6); transform: translate(-50%, -50%); pointer-events: none;
}
.ib-cp-row { display: flex; align-items: center; gap: 6px; }
.ib-cp-prev { width: 24px; height: 24px; border-radius: 5px; border: 1px solid var(--ib-border); flex: 0 0 auto; }
.ib-cp-hex {
    flex: 1; min-width: 0 !important; background: var(--ib-input-bg) !important; color: var(--ib-input-fg) !important;
    border: 1px solid var(--ib-border) !important; border-radius: 5px !important; padding: 3px 6px !important;
    font-size: 0.72rem !important; font-family: monospace !important; height: 24px !important;
    box-sizing: border-box !important;
}
.ib-cp-ok {
    background: #15803d !important; color: #fff !important; border: none !important;
    border-radius: 5px !important; padding: 0 10px !important; font-size: 0.72rem !important;
    height: 24px !important; min-width: 0 !important; box-shadow: none !important; cursor: pointer;
    display: inline-flex !important; align-items: center !important;
}
.ib-cp-ok:hover { filter: brightness(1.12); }
.ib-cp-eye {
    background: var(--ib-eye-bg) !important; color: var(--ib-input-fg) !important; border: 1px solid var(--ib-border) !important;
    border-radius: 5px !important; height: 24px !important; width: 26px !important;
    min-width: 0 !important; padding: 0 !important; box-shadow: none !important; cursor: pointer;
    display: inline-flex !important; align-items: center !important; justify-content: center !important;
    flex: 0 0 auto;
}
.ib-cp-eye:hover { border-color: var(--ib-subdued) !important; color: var(--ib-input-fg) !important; }
.ib-cp-eye.active {
    background: #15803d !important; border-color: #15803d !important; color: #fff !important;
    box-shadow: 0 0 0 2px rgba(21, 128, 61, 0.4) !important;
}

/* ---- Compact the Gradio top-level fields (denser than default, but consistent padding/width) ---- */
#fields-col .block { padding: 6px 8px !important; box-sizing: border-box !important; }
#fields-col .form { gap: 6px !important; }
#fields-col .gap { gap: 6px !important; }
#fields-col span[data-testid="block-info"], #fields-col label > span:first-child {
    font-size: 0.8rem !important; margin-bottom: 2px !important;
}
#fields-col textarea, #fields-col input[type="text"] { font-size: 0.85rem !important; }
/* The elements live in a gr.HTML .block (already padded 6px 8px like the fields), so its content
   needs no extra inset — that keeps every panel aligned on the same left/right edges. */
#fields-col .ib-panel-head { padding: 0; }
#panel-root { padding: 0; }
"""

# The whole interactive canvas (image + draggable/resizable boxes), injected once via demo.load.
# JS owns box identity (each box has a stable id); it emits edits to the hidden #bbox-out textbox
# and receives loads/relabels via #bbox-in (a `.change(js=...)` calls window.IB.load).
CANVAS_JS = r"""
() => {
  function setup() {
    if (window.IB) return true;
    const canvasRoot = document.getElementById("bbox-root");
    if (!canvasRoot) return false; // canvas mounts independently of the (optional) right panel

    const stage = document.createElement("div");
    stage.className = "ib-stage";
    const img = document.createElement("img");
    img.className = "ib-img";
    img.draggable = false;
    const hint = document.createElement("div");
    hint.className = "ib-hint";
    hint.textContent = "Caption an image to start editing boxes here.";
    stage.appendChild(img);
    canvasRoot.appendChild(hint);
    canvasRoot.appendChild(stage);

    let els = []; // {id,type,ymin,xmin,ymax,xmax,desc,text,palette}
    let styleColors = []; // the top-level style_description.color_palette
    let sel = new Set(); // selected ids (multi-select via ctrl/cmd-click)
    let drag = null;
    let jsCounter = 0;
    let emitTimer = null;
    let eyedropSample = null; // active eyedrop callback, else null
    let sampleCtx = null, sampleW = 0, sampleH = 0; // offscreen copy of the image for pixel sampling

    // Keep an offscreen canvas of the (same-origin data-URL) image so we can sample pixels in any
    // browser — cross-browser eyedropper without the Chromium-only EyeDropper API.
    img.addEventListener("load", () => {
      try {
        const c = document.createElement("canvas");
        c.width = img.naturalWidth;
        c.height = img.naturalHeight;
        const cx = c.getContext("2d");
        cx.drawImage(img, 0, 0);
        sampleCtx = cx;
        sampleW = img.naturalWidth;
        sampleH = img.naturalHeight;
      } catch (e) {
        sampleCtx = null;
      }
    });

    const clamp = (v) => Math.max(0, Math.min(1000, v));
    const outEl = () => document.querySelector("#bbox-out textarea, #bbox-out input");
    const byId = (id) => els.find((e) => e.id === id);

    function setNativeValue(el, value) {
      const proto = el.tagName === "TEXTAREA"
        ? window.HTMLTextAreaElement.prototype
        : window.HTMLInputElement.prototype;
      Object.getOwnPropertyDescriptor(proto, "value").set.call(el, value);
    }

    function emitNow() {
      const el = outEl();
      if (!el) return;
      const payload = {
        elements: els.map((e) => ({
          id: e.id,
          type: e.type || "obj",
          bbox: [Math.round(e.ymin), Math.round(e.xmin), Math.round(e.ymax), Math.round(e.xmax)],
          desc: e.desc || "",
          text: e.text || "",
          colors: e.colors || [],
        })),
        style_colors: styleColors,
        selected: Array.from(sel),
      };
      setNativeValue(el, JSON.stringify(payload));
      el.dispatchEvent(new Event("input", { bubbles: true }));
    }
    function emit() {
      clearTimeout(emitTimer);
      emitTimer = setTimeout(emitNow, 150);
    }

    function labelText(e, i) {
      if (e.type === "text") return (i + 1) + '. "' + (e.text || "").slice(0, 18) + '"';
      return (i + 1) + ". " + (e.desc || "").slice(0, 18);
    }

    function renderBoxes() {
      stage.querySelectorAll(".ib-box, .ib-group").forEach((n) => n.remove());
      const single = sel.size === 1;
      els.forEach((e, i) => {
        const d = document.createElement("div");
        d.className = "ib-box" + (sel.has(e.id) ? " sel" : "");
        d.style.left = e.xmin / 10 + "%";
        d.style.top = e.ymin / 10 + "%";
        d.style.width = (e.xmax - e.xmin) / 10 + "%";
        d.style.height = (e.ymax - e.ymin) / 10 + "%";
        d.dataset.id = e.id;
        const lab = document.createElement("span");
        lab.className = "ib-label";
        lab.textContent = labelText(e, i);
        d.appendChild(lab);
        if (single && sel.has(e.id)) {
          ["nw", "ne", "sw", "se"].forEach((h) => {
            const hd = document.createElement("div");
            hd.className = "ib-h ib-" + h;
            hd.dataset.h = h;
            d.appendChild(hd);
          });
        }
        stage.appendChild(d);
      });
      // A single group box with 8 handles when 2+ are selected — drag a handle to scale them together.
      if (sel.size >= 2) {
        const g = groupBBox();
        const go = document.createElement("div");
        go.className = "ib-group";
        go.style.left = g.x0 / 10 + "%";
        go.style.top = g.y0 / 10 + "%";
        go.style.width = (g.x1 - g.x0) / 10 + "%";
        go.style.height = (g.y1 - g.y0) / 10 + "%";
        ["nw", "n", "ne", "e", "se", "s", "sw", "w"].forEach((h) => {
          const hd = document.createElement("div");
          hd.className = "ib-gh ib-" + h;
          hd.dataset.gh = h;
          go.appendChild(hd);
        });
        stage.appendChild(go);
      }
    }

    function cardHTML(e, i) {
      const isText = e.type === "text";
      return (
        '<div class="ib-card' + (sel.has(e.id) ? " sel" : "") + '" data-id="' + e.id + '">' +
        '<div class="ib-card-head">' +
        '<span class="ib-badge">' + (i + 1) + "</span>" +
        '<select class="ib-type" data-f="type">' +
        '<option value="obj"' + (isText ? "" : " selected") + ">obj</option>" +
        '<option value="text"' + (isText ? " selected" : "") + ">text</option>" +
        "</select>" +
        '<button class="ib-del" title="Delete">✕</button>' +
        "</div>" +
        (isText
          ? '<input class="ib-f" data-f="text" placeholder="text (verbatim)" value="' +
            escapeAttr(e.text || "") + '">'
          : "") +
        '<textarea class="ib-f ib-desc" data-f="desc" rows="2" placeholder="description">' +
        escapeHtml(e.desc || "") + "</textarea>" +
        '<div class="ib-pal">' + paletteHTML(e) + "</div>" +
        "</div>"
      );
    }
    function escapeHtml(s) {
      return String(s).replace(/&/g, "&amp;").replace(/</g, "&lt;").replace(/>/g, "&gt;");
    }
    function escapeAttr(s) {
      return escapeHtml(s).replace(/"/g, "&quot;");
    }

    // ---- per-element color palette (swatch chips + native color picker) ----
    function textColor(hex) {
      const m = /^#?([0-9a-f]{6})$/i.exec(hex || "");
      if (!m) return "#fff";
      const n = parseInt(m[1], 16);
      const r = (n >> 16) & 255, g = (n >> 8) & 255, b = n & 255;
      return 0.299 * r + 0.587 * g + 0.114 * b > 140 ? "#000" : "#fff";
    }
    function paletteChipsHTML(colors) {
      const chips = (colors || []).map((c, ci) =>
        '<span class="ib-chip" data-ci="' + ci + '" style="background:' + escapeAttr(c) +
        ";color:" + textColor(c) + '">' +
        '<span class="ib-chip-hex">' + escapeHtml(c) + "</span>" +
        '<button class="ib-chip-x" data-ci="' + ci + '" title="remove" type="button">×</button>' +
        "</span>"
      ).join("");
      return chips + '<button class="ib-chip-add" type="button">+ color</button>';
    }
    const paletteHTML = (e) => paletteChipsHTML(e.colors || []);
    function refreshPalette(id) {
      const card = document.querySelector('#panel-root .ib-card[data-id="' + id + '"]');
      if (!card) return;
      const pal = card.querySelector(".ib-pal");
      if (pal) pal.innerHTML = paletteHTML(byId(id));
    }
    function renderStylePalette() {
      const root = document.getElementById("style-pal-root");
      if (root) root.innerHTML = paletteChipsHTML(styleColors);
    }
    // ---- self-contained color picker popover (native <input type=color> can't be positioned and
    //      opens an OS dialog with no OK on Linux) ----
    const clamp2 = (x, a, b) => Math.max(a, Math.min(b, x));
    function hexToRgb(hex) {
      const m = /^#?([0-9a-f]{6})$/i.exec(hex || "");
      if (!m) return { r: 77, g: 164, b: 96 };
      const n = parseInt(m[1], 16);
      return { r: (n >> 16) & 255, g: (n >> 8) & 255, b: n & 255 };
    }
    function rgbToHex(r, g, b) {
      const h = (x) => ("0" + Math.round(x).toString(16)).slice(-2);
      return ("#" + h(r) + h(g) + h(b)).toUpperCase();
    }
    function rgbToHsv(r, g, b) {
      r /= 255; g /= 255; b /= 255;
      const mx = Math.max(r, g, b), mn = Math.min(r, g, b), d = mx - mn;
      let h = 0;
      if (d) {
        if (mx === r) h = ((g - b) / d) % 6;
        else if (mx === g) h = (b - r) / d + 2;
        else h = (r - g) / d + 4;
        h *= 60; if (h < 0) h += 360;
      }
      return { h: h, s: (mx ? d / mx : 0) * 100, v: mx * 100 };
    }
    function hsvToRgb(h, s, v) {
      s /= 100; v /= 100;
      const c = v * s, x = c * (1 - Math.abs(((h / 60) % 2) - 1)), m = v - c;
      let r = 0, g = 0, b = 0;
      if (h < 60) { r = c; g = x; } else if (h < 120) { r = x; g = c; }
      else if (h < 180) { g = c; b = x; } else if (h < 240) { g = x; b = c; }
      else if (h < 300) { r = x; b = c; } else { r = c; b = x; }
      return { r: (r + m) * 255, g: (g + m) * 255, b: (b + m) * 255 };
    }
    function sampleAt(clientX, clientY) {
      if (!sampleCtx) return null;
      const r = stage.getBoundingClientRect();
      const fx = (clientX - r.left) / r.width, fy = (clientY - r.top) / r.height;
      if (fx < 0 || fx > 1 || fy < 0 || fy > 1) return null;
      const px = Math.min(sampleW - 1, Math.max(0, Math.floor(fx * sampleW)));
      const py = Math.min(sampleH - 1, Math.max(0, Math.floor(fy * sampleH)));
      const d = sampleCtx.getImageData(px, py, 1, 1).data;
      return rgbToHex(d[0], d[1], d[2]);
    }
    // Enter "pick from image" mode: hover previews, click commits, Esc cancels. onSample(hex|null, committed).
    function startEyedrop(onSample) {
      stage.classList.add("ib-eyedrop");
      const move = (ev) => { const hx = sampleAt(ev.clientX, ev.clientY); if (hx) onSample(hx, false); };
      const down = (ev) => {
        if (ev.button !== 0) return;
        ev.preventDefault();
        ev.stopPropagation();
        const hx = sampleAt(ev.clientX, ev.clientY);
        end();
        onSample(hx, true);
      };
      const key = (ev) => { if (ev.key === "Escape") { end(); onSample(null, true); } };
      function end() {
        eyedropSample = null;
        stage.classList.remove("ib-eyedrop");
        stage.removeEventListener("pointermove", move, true);
        stage.removeEventListener("pointerdown", down, true);
        document.removeEventListener("keydown", key, true);
      }
      eyedropSample = onSample;
      stage.addEventListener("pointermove", move, true);
      stage.addEventListener("pointerdown", down, true);
      document.addEventListener("keydown", key, true);
    }

    function openColorPopover(initial, anchor, cb) {
      document.querySelectorAll(".ib-cp").forEach((n) => n.remove());
      const rgb = hexToRgb(initial);
      const hsv = rgbToHsv(rgb.r, rgb.g, rgb.b);
      let h = hsv.h, s = hsv.s, v = hsv.v;
      const eyeBtn =
        '<button class="ib-cp-eye" type="button" title="Pick color from the image">' +
        '<svg width="13" height="13" viewBox="0 0 24 24" fill="none" stroke="currentColor" ' +
        'stroke-width="2" stroke-linecap="round" stroke-linejoin="round">' +
        '<path d="m2 22 1-1h3l9-9"/><path d="M3 21v-3l9-9"/>' +
        '<path d="m15 6 3.4-3.4a2.1 2.1 0 1 1 3 3L18 9l.4.4a2.1 2.1 0 1 1-3 3l-3.8-3.8a2.1 ' +
        '2.1 0 1 1 3-3l.4.4Z"/></svg></button>';
      const pop = document.createElement("div");
      pop.className = "ib-cp";
      pop.innerHTML =
        '<div class="ib-cp-sv"><div class="ib-cp-svknob"></div></div>' +
        '<div class="ib-cp-hue"><div class="ib-cp-hueknob"></div></div>' +
        '<div class="ib-cp-row"><span class="ib-cp-prev"></span>' + eyeBtn +
        '<input class="ib-cp-hex" maxlength="7" spellcheck="false">' +
        '<button class="ib-cp-ok" type="button">OK</button></div>';
      document.body.appendChild(pop);
      const r0 = anchor ? anchor.getBoundingClientRect() : { left: 100, top: 100, bottom: 100 };
      const pw = 188, ph = 210;
      pop.style.left = Math.max(8, Math.min(r0.left, window.innerWidth - pw - 8)) + "px";
      let top = r0.bottom + 6;
      if (top + ph > window.innerHeight) top = Math.max(8, r0.top - ph - 6);
      pop.style.top = top + "px";
      const sv = pop.querySelector(".ib-cp-sv");
      const svk = pop.querySelector(".ib-cp-svknob");
      const hue = pop.querySelector(".ib-cp-hue");
      const huek = pop.querySelector(".ib-cp-hueknob");
      const prev = pop.querySelector(".ib-cp-prev");
      const hexIn = pop.querySelector(".ib-cp-hex");
      const curHex = () => { const c = hsvToRgb(h, s, v); return rgbToHex(c.r, c.g, c.b); };
      function paint() {
        sv.style.background =
          "linear-gradient(to top,#000,rgba(0,0,0,0)),linear-gradient(to right,#fff,hsl(" +
          h + ",100%,50%))";
        svk.style.left = s + "%";
        svk.style.top = 100 - v + "%";
        huek.style.left = (h / 360) * 100 + "%";
        const hx = curHex();
        prev.style.background = hx;
        if (document.activeElement !== hexIn) hexIn.value = hx;
      }
      paint();
      const svFrom = (ev) => {
        const r = sv.getBoundingClientRect();
        s = clamp2(((ev.clientX - r.left) / r.width) * 100, 0, 100);
        v = clamp2((1 - (ev.clientY - r.top) / r.height) * 100, 0, 100);
        paint();
      };
      const hueFrom = (ev) => {
        const r = hue.getBoundingClientRect();
        h = clamp2(((ev.clientX - r.left) / r.width) * 360, 0, 360);
        paint();
      };
      function drag(el, fn) {
        el.addEventListener("pointerdown", (ev) => {
          el.setPointerCapture(ev.pointerId);
          fn(ev);
          const mv = (e) => fn(e);
          const up = () => {
            el.removeEventListener("pointermove", mv);
            el.removeEventListener("pointerup", up);
          };
          el.addEventListener("pointermove", mv);
          el.addEventListener("pointerup", up);
          ev.preventDefault();
        });
      }
      drag(sv, svFrom);
      drag(hue, hueFrom);
      hexIn.addEventListener("input", () => {
        if (/^#?[0-9a-f]{6}$/i.test(hexIn.value)) {
          const c = hexToRgb(hexIn.value);
          const hs = rgbToHsv(c.r, c.g, c.b);
          h = hs.h; s = hs.s; v = hs.v;
          paint();
        }
      });
      function close() {
        pop.remove();
        document.removeEventListener("pointerdown", outside, true);
      }
      function outside(ev) { if (!pop.contains(ev.target)) close(); }
      setTimeout(() => document.addEventListener("pointerdown", outside, true), 0);
      pop.querySelector(".ib-cp-ok").addEventListener("click", () => {
        const hx = curHex();
        close();
        cb(hx);
      });
      const eyeEl = pop.querySelector(".ib-cp-eye");
      eyeEl.addEventListener("click", (ev) => {
        ev.preventDefault();
        ev.stopPropagation();
        // Suspend click-outside close while sampling over the image, then re-arm it.
        document.removeEventListener("pointerdown", outside, true);
        eyeEl.classList.add("active"); // show the pipette is armed
        startEyedrop((hx, committed) => {
          if (hx) {
            const c = hexToRgb(hx);
            const hs = rgbToHsv(c.r, c.g, c.b);
            h = hs.h; s = hs.s; v = hs.v;
            paint();
          }
          if (committed) {
            eyeEl.classList.remove("active");
            if (pop.isConnected) {
              setTimeout(() => document.addEventListener("pointerdown", outside, true), 50);
            }
          }
        });
      });
    }

    function renderPanel() {
      const panelRoot = document.getElementById("panel-root");
      if (!panelRoot) return; // panel may not be mounted yet; canvas still works
      if (!els.length) {
        panelRoot.innerHTML = '<div class="ib-empty">No elements yet — draw a box on the ' +
          "image, or use + Object / + Text.</div>";
        return;
      }
      panelRoot.innerHTML = els.map((e, i) => cardHTML(e, i)).join("");
    }

    function renderAll() {
      renderBoxes();
      renderPanel();
      renderStylePalette();
    }

    const selArr = () => els.filter((e) => sel.has(e.id));
    function groupBBox() {
      const s = selArr();
      if (!s.length) return null;
      let x0 = 1e9, y0 = 1e9, x1 = -1e9, y1 = -1e9;
      s.forEach((e) => {
        x0 = Math.min(x0, e.xmin); y0 = Math.min(y0, e.ymin);
        x1 = Math.max(x1, e.xmax); y1 = Math.max(y1, e.ymax);
      });
      return { x0: x0, y0: y0, x1: x1, y1: y1 };
    }
    function applySelClasses() {
      stage.querySelectorAll(".ib-box").forEach((n) => n.classList.toggle("sel", sel.has(n.dataset.id)));
      const pr = document.getElementById("panel-root");
      if (pr) pr.querySelectorAll(".ib-card").forEach((n) => n.classList.toggle("sel", sel.has(n.dataset.id)));
    }
    function scrollToCard(id) {
      const pr = document.getElementById("panel-root");
      const c = pr && pr.querySelector('.ib-card[data-id="' + id + '"]');
      if (c) c.scrollIntoView({ block: "nearest" });
    }
    function selectOnly(id) {
      sel = new Set(id ? [id] : []);
      renderBoxes();
      applySelClasses();
      if (id) scrollToCard(id);
    }
    function toggleSelect(id) {
      if (sel.has(id)) sel.delete(id); else sel.add(id);
      renderBoxes();
      applySelClasses();
    }

    function toNorm(ev) {
      const r = stage.getBoundingClientRect();
      return {
        x: clamp(((ev.clientX - r.left) / r.width) * 1000),
        y: clamp(((ev.clientY - r.top) / r.height) * 1000),
      };
    }

    // snapshot the current geometry of the selected boxes (for move/group-resize deltas)
    const snapshot = () => {
      const o = {};
      selArr().forEach((e) => { o[e.id] = { xmin: e.xmin, ymin: e.ymin, xmax: e.xmax, ymax: e.ymax }; });
      return o;
    };

    function resizeSingle(dr, p, shift) {
      const e = dr.e, o = dr.orig, h = dr.h;
      let xmin = o.xmin, ymin = o.ymin, xmax = o.xmax, ymax = o.ymax;
      if (h.includes("w")) xmin = Math.min(p.x, o.xmax - 1);
      if (h.includes("e")) xmax = Math.max(p.x, o.xmin + 1);
      if (h.includes("n")) ymin = Math.min(p.y, o.ymax - 1);
      if (h.includes("s")) ymax = Math.max(p.y, o.ymin + 1);
      if (shift) { // preserve aspect ratio, anchored at the opposite corner
        const ow = o.xmax - o.xmin, oh = o.ymax - o.ymin;
        const scale = Math.max((xmax - xmin) / ow, (ymax - ymin) / oh);
        const nw = ow * scale, nh = oh * scale;
        if (h.includes("w")) xmin = o.xmax - nw; else xmax = o.xmin + nw;
        if (h.includes("n")) ymin = o.ymax - nh; else ymax = o.ymin + nh;
      }
      e.xmin = clamp(xmin); e.ymin = clamp(ymin); e.xmax = clamp(xmax); e.ymax = clamp(ymax);
    }

    function groupResize(dr, p, shift) {
      const h = dr.h, g = dr.g;
      let nx0 = g.x0, ny0 = g.y0, nx1 = g.x1, ny1 = g.y1;
      if (h.includes("w")) nx0 = Math.min(p.x, g.x1 - 1);
      if (h.includes("e")) nx1 = Math.max(p.x, g.x0 + 1);
      if (h.includes("n")) ny0 = Math.min(p.y, g.y1 - 1);
      if (h.includes("s")) ny1 = Math.max(p.y, g.y0 + 1);
      let sx = (nx1 - nx0) / (g.x1 - g.x0);
      let sy = (ny1 - ny0) / (g.y1 - g.y0);
      if (h === "n" || h === "s") sx = 1;
      if (h === "e" || h === "w") sy = 1;
      if (shift && h.length === 2) { const f = Math.max(sx, sy); sx = f; sy = f; }
      const ax = h.includes("w") ? g.x1 : g.x0; // anchor = opposite side
      const ay = h.includes("n") ? g.y1 : g.y0;
      selArr().forEach((e) => {
        const o = dr.orig[e.id];
        e.xmin = clamp(ax + (o.xmin - ax) * sx);
        e.xmax = clamp(ax + (o.xmax - ax) * sx);
        e.ymin = clamp(ay + (o.ymin - ay) * sy);
        e.ymax = clamp(ay + (o.ymax - ay) * sy);
      });
    }

    // ---- canvas interaction ----
    // Pick the SMALLEST box under a point, so overlapping big boxes don't block small ones.
    function hitBox(p) {
      let best = null, bestArea = Infinity;
      for (const e of els) {
        if (p.x >= e.xmin && p.x <= e.xmax && p.y >= e.ymin && p.y <= e.ymax) {
          const area = (e.xmax - e.xmin) * (e.ymax - e.ymin);
          if (area < bestArea) { bestArea = area; best = e; }
        }
      }
      return best;
    }

    stage.addEventListener("pointerdown", (ev) => {
      if (ev.button !== 0 || !img.src || eyedropSample) return;
      const p = toNorm(ev);
      const gh = ev.target.classList.contains("ib-gh") ? ev.target.dataset.gh : null;
      const boxEl = ev.target.closest(".ib-box");
      const handle = ev.target.classList.contains("ib-h") ? ev.target.dataset.h : null;
      const additive = ev.ctrlKey || ev.metaKey;

      if (gh) {
        drag = { mode: "gresize", h: gh, g: groupBBox(), orig: snapshot() };
      } else if (handle && boxEl) {
        const e = byId(boxEl.dataset.id);
        drag = { mode: "resize", h: handle, e: e,
          orig: { xmin: e.xmin, ymin: e.ymin, xmax: e.xmax, ymax: e.ymax } };
      } else if (hitBox(p)) {
        const id = hitBox(p).id; // smallest box under the cursor
        if (additive) {
          toggleSelect(id);
          stage.setPointerCapture(ev.pointerId);
          ev.preventDefault();
          return;
        }
        if (!sel.has(id)) selectOnly(id);
        drag = { mode: "move", start: p, orig: snapshot(), ids: Array.from(sel) };
      } else {
        sel = new Set();
        const id = "j" + jsCounter++;
        const e = { id: id, type: "obj", ymin: p.y, xmin: p.x, ymax: p.y, xmax: p.x,
          desc: "", text: "", colors: [] };
        els.push(e);
        sel = new Set([id]);
        drag = { mode: "draw", e: e, start: p };
        renderAll();
      }
      stage.setPointerCapture(ev.pointerId);
      ev.preventDefault();
    });

    stage.addEventListener("pointermove", (ev) => {
      if (!drag) return;
      const p = toNorm(ev);
      const shift = ev.shiftKey;
      if (drag.mode === "move") {
        let dx = p.x - drag.start.x, dy = p.y - drag.start.y;
        let gx0 = 1e9, gy0 = 1e9, gx1 = -1e9, gy1 = -1e9;
        drag.ids.forEach((id) => {
          const o = drag.orig[id];
          gx0 = Math.min(gx0, o.xmin); gy0 = Math.min(gy0, o.ymin);
          gx1 = Math.max(gx1, o.xmax); gy1 = Math.max(gy1, o.ymax);
        });
        dx = Math.max(-gx0, Math.min(1000 - gx1, dx));
        dy = Math.max(-gy0, Math.min(1000 - gy1, dy));
        drag.ids.forEach((id) => {
          const o = drag.orig[id], e = byId(id);
          if (!e) return;
          e.xmin = o.xmin + dx; e.ymin = o.ymin + dy;
          e.xmax = o.xmax + dx; e.ymax = o.ymax + dy;
        });
      } else if (drag.mode === "draw") {
        const e = drag.e;
        e.xmin = Math.min(drag.start.x, p.x);
        e.xmax = Math.max(drag.start.x, p.x);
        e.ymin = Math.min(drag.start.y, p.y);
        e.ymax = Math.max(drag.start.y, p.y);
      } else if (drag.mode === "resize") {
        resizeSingle(drag, p, shift);
      } else if (drag.mode === "gresize") {
        groupResize(drag, p, shift);
      }
      renderBoxes();
    });

    function endDrag() {
      if (!drag) return;
      if (drag.mode === "draw") {
        const e = drag.e;
        if (e.xmax - e.xmin < 8 || e.ymax - e.ymin < 8) {
          els = els.filter((x) => x !== e);
          sel = new Set();
        }
      }
      const wasDraw = drag.mode === "draw";
      drag = null;
      renderAll();
      applySelClasses();
      emitNow();
      if (wasDraw) {
        const pr = document.getElementById("panel-root");
        const card = pr && pr.querySelector(".ib-card.sel .ib-desc");
        if (card) card.focus();
      }
    }
    stage.addEventListener("pointerup", endDrag);
    stage.addEventListener("pointercancel", endDrag);

    window.addEventListener("keydown", (ev) => {
      if (ev.key === "Delete" && sel.size) {
        const t = document.activeElement;
        if (t && (t.tagName === "TEXTAREA" || t.tagName === "INPUT")) return;
        els = els.filter((x) => !sel.has(x.id));
        sel = new Set();
        renderAll();
        emitNow();
      }
    });

    // ---- panel interaction: document-level delegation so it works no matter when #panel-root
    //      renders and survives re-renders (the panel div is owned by Gradio's HTML component) ----
    document.addEventListener("input", (ev) => {
      const card = ev.target.closest && ev.target.closest("#panel-root .ib-card");
      if (!card) return;
      const e = byId(card.dataset.id);
      const f = ev.target.dataset.f;
      if (!e || !f) return;
      e[f] = ev.target.value;
      if (f === "desc" || f === "text") {
        const i = els.indexOf(e);
        const lab = stage.querySelector('.ib-box[data-id="' + e.id + '"] .ib-label');
        if (lab) lab.textContent = labelText(e, i);
      }
      emit();
    });
    document.addEventListener("change", (ev) => {
      if (!ev.target.classList || !ev.target.classList.contains("ib-type")) return;
      if (!ev.target.closest("#panel-root")) return;
      const card = ev.target.closest(".ib-card");
      const e = byId(card.dataset.id);
      if (!e) return;
      e.type = ev.target.value;
      renderAll();
      selectOnly(e.id);
      emitNow();
    });
    document.addEventListener("click", (ev) => {
      if (!ev.target.closest) return;
      if (ev.target.id === "ib-add-obj") { addElement("obj"); return; }
      if (ev.target.id === "ib-add-text") { addElement("text"); return; }
      // top-level style palette (same chip UI as elements, but backed by styleColors)
      if (ev.target.closest("#style-pal-root")) {
        const addB = ev.target.closest(".ib-chip-add");
        if (addB) {
          const last = styleColors.length ? styleColors[styleColors.length - 1] : "#4DA460";
          openColorPopover(last, addB, (hex) => {
            styleColors.push(hex); renderStylePalette(); emitNow();
          });
          return;
        }
        const xB = ev.target.closest(".ib-chip-x");
        if (xB) { styleColors.splice(parseInt(xB.dataset.ci, 10), 1); renderStylePalette(); emitNow(); return; }
        const ch = ev.target.closest(".ib-chip");
        if (ch) {
          const ci = parseInt(ch.dataset.ci, 10);
          openColorPopover(styleColors[ci], ch, (hex) => {
            styleColors[ci] = hex; renderStylePalette(); emitNow();
          });
        }
        return;
      }
      const card = ev.target.closest("#panel-root .ib-card");
      if (!card) return;
      const e = byId(card.dataset.id);
      if (ev.target.classList.contains("ib-del")) {
        els = els.filter((x) => x.id !== card.dataset.id);
        sel.delete(card.dataset.id);
        renderAll();
        emitNow();
        return;
      }
      // ---- color palette chips ----
      const addBtn = ev.target.closest(".ib-chip-add");
      if (addBtn) {
        const last = e.colors && e.colors.length ? e.colors[e.colors.length - 1] : "#4DA460";
        openColorPopover(last, addBtn, (hex) => {
          e.colors = (e.colors || []).concat(hex);
          refreshPalette(e.id);
          emitNow();
        });
        return;
      }
      const xBtn = ev.target.closest(".ib-chip-x");
      if (xBtn) {
        e.colors.splice(parseInt(xBtn.dataset.ci, 10), 1);
        refreshPalette(e.id);
        emitNow();
        return;
      }
      const chip = ev.target.closest(".ib-chip");
      if (chip) {
        const ci = parseInt(chip.dataset.ci, 10);
        openColorPopover(e.colors[ci], chip, (hex) => {
          e.colors[ci] = hex;
          refreshPalette(e.id);
          emitNow();
        });
        return;
      }
      if (!ev.target.closest(".ib-f, .ib-type, .ib-pal")) {
        if (ev.ctrlKey || ev.metaKey) toggleSelect(card.dataset.id);
        else selectOnly(card.dataset.id);
      }
    });

    function addElement(type) {
      const id = "j" + jsCounter++;
      els.push({ id: id, type: type, ymin: 350, xmin: 350, ymax: 650, xmax: 650,
        desc: "", text: "", colors: [] });
      renderAll();
      selectOnly(id);
      emitNow();
    }

    window.IB = {
      load(v) {
        try {
          console.log("[IB] load called, payload length", (v || "").length);
          const p = typeof v === "string" ? JSON.parse(v || "{}") : v;
          if (!p) { console.warn("[IB] load: empty payload"); return; }
          console.log("[IB] load: hasImage", !!p.image, "elements",
            Array.isArray(p.elements) ? p.elements.length : "none");
          if (p.image) { img.src = p.image; hint.style.display = "none"; }
          if (Array.isArray(p.elements)) {
            els = p.elements.map((e) => Object.assign({}, e));
          }
          styleColors = Array.isArray(p.style_colors) ? p.style_colors.slice() : [];
          sel = new Set(Array.isArray(p.selected) ? p.selected : []);
          renderAll();
        } catch (e) {
          console.error("[IB] load error", e, "payload head:", (v || "").slice(0, 120));
        }
      },
    };
    console.log("[IB] setup complete");
    if (window.__bboxPending) window.IB.load(window.__bboxPending);
    return true;
  }

  // Gradio 6 lazily renders inactive tab content, so #bbox-root may not exist until the Editor tab
  // is shown (which can be minutes later, after inference + auto-switch). Mount the instant it
  // appears via a MutationObserver, with a slow poll as backup. setup() applies any pending payload.
  if (!setup()) {
    const obs = new MutationObserver(() => { if (setup()) obs.disconnect(); });
    obs.observe(document.body, { childList: true, subtree: true });
    let tries = 0;
    const t = setInterval(() => {
      if (setup() || ++tries > 3000) clearInterval(t);
    }, 200);
  }
}
"""

# Inject CSS + the canvas script straight into <head> as literal tags. This bypasses Gradio 6's
# config-js path (which needs `new Function` and may be blocked/uninvoked), so the browser just runs
# the <script> on load. CANVAS_JS is an arrow fn; `(...)()` self-invokes it.
# --app-h is a *fixed pixel* stand-in for 100vh. On HF Spaces the app runs in a <gradio-app>
# iframe that resizes to its content, so `vh` units feed back and grow forever. Pinning a pixel
# height from window.innerHeight (updated only on real resize) breaks that loop.
# --app-h is a fixed-px stand-in for the app height.
#   Direct URL (top-level page): use the real available viewport -> full-height, scrollless.
#   Embedded in an iframe (e.g. the HF Space page): the iframe is content-sized, so innerHeight is
#   circular and unrelated to the visible area — we can't know it, so we use a sensible fixed height.
EMBED_HEIGHT_PX = 720
HEIGHT_SYNC_JS = (
    "(function(){var embedded=window.top!==window.self;var h=0;function s(){"
    "if(embedded){h=" + str(EMBED_HEIGHT_PX) + ";}"
    "else{var el=document.querySelector('.gradio-container');"
    "var top=el?el.getBoundingClientRect().top:0;"
    "var a=Math.max(320, window.innerHeight - Math.max(0, top));"
    "if(!h||a<h)h=a;}"  # only ever shrink on the direct URL (never chase an auto-sizing parent)
    "document.documentElement.style.setProperty('--app-h', h + 'px');}"
    "s();window.addEventListener('resize', s);window.addEventListener('load', s);"
    "setTimeout(s,300);setTimeout(s,900);})();"
)
HEAD = (
    "<style>\n" + CSS + "\n</style>\n"
    "<script>\nconsole.log('[IB] head script running');\n"
    + HEIGHT_SYNC_JS + "\n(" + CANVAS_JS + ")();\n</script>\n"
)

intro = (
    "# 🏷️ Ideogram 4 Caption Studio\n"
    "<small>Gemma-4 describes your image in Ideogram 4's native JSON schema, then edit it visually. "
    "&nbsp;·&nbsp; <b>Caption</b> tab: upload → Caption. "
    "&nbsp;·&nbsp; <b>Editor</b> tab: drag on the image to draw a box; drag it to move, handles to "
    "resize; <b>Ctrl-click</b> to multi-select (group-scale, hold <b>Shift</b> to keep aspect); "
    "<b>Del</b> removes. Edit each box's text &amp; colors on the right (the pipette samples the "
    "image). The <b>JSON</b> tab holds the raw caption to feed the generator.</small>"
)

with gr.Blocks(title="Ideogram4 Caption (Gemma4)") as demo:
    gr.Markdown(intro, elem_id="app-title")

    # Single source of truth for the whole caption; the fields, JSON tab and canvas are views of it.
    caption_state = gr.State({})
    # The Ideogram4 schema prompt is fixed and hidden; kept as a state so it's still passed to caption().
    instruction = gr.State(value=DEFAULT_INSTRUCTION)

    with gr.Tabs(elem_id="main-tabs") as main_tabs:
        # ---- Caption tab: upload + model + generation controls -------------------------------------
        with gr.Tab("Caption", id="caption"):
            with gr.Row():
                with gr.Column(scale=2):
                    image = gr.Image(
                        label="Source image",
                        type="pil",
                        height="calc(var(--app-h, 100vh) - 190px)",
                    )
                with gr.Column(scale=1):
                    run = gr.Button(
                        "Caption", variant="primary", elem_id="caption-btn"
                    )
                    run_status = gr.Markdown("")

                    if IS_ZERO_GPU:
                        # ZeroGPU: bf16 preloaded, GPU is per-request — no selector / load / unload.
                        precision = gr.State(value=DEFAULT_PRECISION)
                    else:
                        # Local: pick a model, load on demand, free it so we don't squat on VRAM.
                        with gr.Row(equal_height=True):
                            precision = gr.Dropdown(
                                choices=list(MODELS),
                                value=DEFAULT_PRECISION,
                                show_label=False,
                                scale=2,
                            )
                            load_btn = gr.Button("Load model", elem_id="load-btn")
                            unload_btn = gr.Button("Unload model", elem_id="unload-btn")
                        model_status = gr.Textbox(
                            label="Model status",
                            value="⚪ No model loaded",
                            interactive=False,
                        )

                    max_new_tokens = gr.Slider(
                        256, 4096, value=2048, step=64, label="Max new tokens"
                    )
                    temperature = gr.Slider(
                        0.0, 2.0, value=0.0, step=0.1, label="Temperature (0 = greedy)"
                    )

        # ---- Editor tab: interactive box canvas (left) + fields / JSON (right) ---------------------
        with gr.Tab("Editor", id="editor"):
            # Hidden bridge to the JS canvas: Python writes #bbox-in (load/relabel), JS writes
            # #bbox-out (user edits). Neither is user-visible.
            bbox_in = gr.Textbox(visible=False, elem_id="bbox-in")
            bbox_out = gr.Textbox(visible=False, elem_id="bbox-out")
            with gr.Row(elem_id="editor-row"):
                with gr.Column(scale=3, elem_id="canvas-col"):
                    gr.HTML(
                        '<div id="bbox-root"></div>',
                        label="Boxes — drag on empty space to draw · drag box to move · "
                        "handles to resize · Del to remove",
                    )
                with gr.Column(scale=2, elem_id="fields-col"):
                    with gr.Tabs():
                        with gr.Tab("Fields"):
                            hld = gr.Textbox(label="High-level description", lines=2)
                            background = gr.Textbox(label="Background", lines=2)
                            with gr.Group():
                                gr.Markdown("**Style**")
                                medium_kind = gr.Radio(
                                    ["photo", "art_style"],
                                    value="photo",
                                    label="Medium kind",
                                )
                                aesthetics = gr.Textbox(label="Aesthetics")
                                lighting = gr.Textbox(label="Lighting")
                                medium = gr.Textbox(label="Medium")
                                kind_value = gr.Textbox(
                                    label="Photo / art_style value",
                                    info="camera & lens for photo, or the named art style",
                                )
                                # Style color palette — same JS chip editor (picker + eyedropper)
                                # as the elements; backed by styleColors in the canvas script.
                                gr.HTML(
                                    '<div class="ib-pal-label">Color palette</div>'
                                    '<div id="style-pal-root" class="ib-pal"></div>'
                                )
                            # Element cards are rendered by the JS editor (kept in lockstep with the
                            # boxes: draw/select/edit sync both ways). #ib-add-* buttons are wired in JS.
                            gr.HTML(
                                '<div class="ib-panel-head">'
                                "<strong>Elements</strong>"
                                '<span class="ib-add-wrap">'
                                '<button id="ib-add-obj" class="ib-add" type="button">+ Object</button>'
                                '<button id="ib-add-text" class="ib-add" type="button">+ Text</button>'
                                "</span></div>"
                                '<div id="panel-root"></div>'
                            )

                        with gr.Tab("JSON"):
                            json_code = gr.Code(
                                label="Ideogram4 JSON caption", language="json"
                            )
                            apply_btn = gr.Button("Apply JSON", variant="secondary")

    # sync_from_state returns (state-with-ids, top fields..., canvas payload) in this exact order.
    # (color_palette is now a JS chip editor, not a Gradio field.)
    sync_outputs = [
        caption_state,
        hld,
        background,
        medium_kind,
        aesthetics,
        lighting,
        medium,
        kind_value,
        bbox_in,
    ]

    # Non-label fields (they never change the boxes) -> just rewrite state + JSON.
    field_inputs = [
        caption_state,
        hld,
        background,
        medium_kind,
        aesthetics,
        lighting,
        medium,
        kind_value,
    ]
    for field in [
        hld,
        background,
        medium_kind,
        aesthetics,
        lighting,
        medium,
        kind_value,
    ]:
        field.input(
            rebuild_from_fields,
            inputs=field_inputs,
            outputs=[caption_state, json_code],
        )

    # All editor edits (draw / move / resize / delete / add / type / desc / text / palette / select)
    # happen in JS and arrive together via #bbox-out. We just fold them into state + refresh JSON;
    # the JS already rendered everything, so there's nothing to push back.
    def on_canvas_out(out, st):
        new_st = reconcile_canvas(st, out)
        return new_st, _dumps(new_st)

    bbox_out.input(
        on_canvas_out,
        [bbox_out, caption_state],
        [caption_state, json_code],
    )

    if not IS_ZERO_GPU:
        load_btn.click(load_model, inputs=precision, outputs=model_status, api_name="load")
        unload_btn.click(
            unload_model, inputs=None, outputs=model_status, api_name="unload"
        )

    # Show the image on the canvas as soon as it's uploaded (boxes appear once captioned/edited).
    image.change(
        lambda img, st: canvas_payload(img, st, include_image=True),
        inputs=[image, caption_state],
        outputs=bbox_in,
    )

    # Push #bbox-in payloads into the JS canvas.
    bbox_in.change(
        None,
        inputs=bbox_in,
        outputs=None,
        js="(v) => { console.log('[IB] bbox_in change, len', (v||'').length, 'IB?', !!window.IB); window.__bboxPending = v; if (window.IB) window.IB.load(v); }",
    )

    run.click(
        lambda: "⏳ Running the model…",
        outputs=run_status,
    ).then(
        caption,
        inputs=[image, instruction, precision, max_new_tokens, temperature],
        outputs=[caption_state, json_code, run_status],
        api_name="caption",
    ).then(
        sync_from_state,
        inputs=[image, caption_state],
        outputs=sync_outputs,
    ).then(
        lambda: gr.Tabs(selected="editor"),
        outputs=main_tabs,
    )

    apply_btn.click(
        apply_json,
        inputs=[image, json_code],
        outputs=sync_outputs,
    )


# In Gradio 6, css/js/head live on launch(). We inject via head (literal <style>/<script>) so it
# doesn't depend on Gradio's config-js execution.
if __name__ == "__main__":
    demo.queue().launch(head=HEAD)