"""app.py — HuggingFace ZeroGPU Space: the deterministic 12-task vision
extraction + fusion pipeline as an interactive showcase.

Stick an image in → get the full JSON readout (12 task JSONs + FusedScene +
deterministic fused prompt + overlays). Every possibility in the system is a
selectable toggle.

ZeroGPU teardown-friendly design
--------------------------------
* PYTORCH_CUDA_ALLOC_CONF is set BEFORE torch imports (OOM-probing batched path).
* The always-on specialist models load ONCE at module level (CUDA-emulation
  outside `@spaces.GPU`; real CUDA inside) — the efficient, fork-friendly residency.
* Optional structurer (0.8B / 9B) + age gate load on demand, single-resident.
* GPU functions return only picklable CPU data (task/digest dicts + rendered PIL
  overlays). fuse()/fused_prompt()/build_semantic_association() run on the CPU in
  the main process — no GPU is held during fusion.

The pipeline modules themselves are the real `qwen_test_runner` package, vendored
verbatim by ../build_space.py.
"""
from __future__ import annotations

import os

# ── ZeroGPU rule: set the allocator conf BEFORE torch is imported ────────────
os.environ.setdefault("PYTORCH_CUDA_ALLOC_CONF", "expandable_segments:True")

import json
import tempfile
import time

import numpy as np
from PIL import Image, ImageDraw

import gradio as gr

# spaces (ZeroGPU). Degrade to a no-op decorator when running off-platform.
try:
    import spaces

    _HAS_SPACES = True
except Exception:  # pragma: no cover - local/CPU dev
    _HAS_SPACES = False

    class _NoSpaces:
        @staticmethod
        def GPU(*_a, **_k):
            def deco(fn):
                return fn

            return deco

    spaces = _NoSpaces()  # type: ignore

import torch

# ── real pipeline (vendored package) ─────────────────────────────────────────
import qwen_test_runner.vision.specialists_gpu as g
from qwen_test_runner.vision.specialists import Solids
from qwen_test_runner.vision import derive
from qwen_test_runner.vision.fuse import (
    solids_digest,
    fuse,
    phrases_for_grounding,
    build_semantic_association,
)
from qwen_test_runner.vision.fuse_prompt import fused_prompt
from qwen_test_runner.vision.tasks_vision import get_task, model_for
from qwen_test_runner.vision.coords import CoordSpace
from qwen_test_runner.model_runner import SYSTEM_PROMPT_JSON
from qwen_test_runner.evaluator import parse_safely


DEVICE = "cuda" if torch.cuda.is_available() else "cpu"
IS_GPU_ENV = bool(os.environ.get("SPACES_ZERO_GPU")) or DEVICE == "cuda"
MAX_DIM = 1024                      # match production DECODE_MAX_DIM
BATCH_CAP = 24                      # interactive batch ceiling (ZeroGPU quota)

# 12 deterministic tasks (11 from _build_tasks + semantic_association from fusion)
DET_TASKS = [
    "image_classification", "bbox_grounding", "ocr_text",
    "data_type_differentiation", "data_type_utilization",
    "structural_spatial_awareness", "depth_analysis", "subject_fixation",
    "segmentation", "outline_association", "style_structural_awareness",
    "semantic_association",
]
# registry entries with no deterministic builder (shown, disabled)
VLM_TASKS = ["vit_accuracy_to_prompt", "geometric_3d_object_id", "camera_rotational_offset"]

VOCABS = {"COCO-80": g.COCO_CLASSES, "shapes": g.SHAPE_CLASSES}
STRUCTURERS = {"off": None, "Qwen3.5-0.8B": "Qwen/Qwen3.5-0.8B", "Qwen3.5-9B": "Qwen/Qwen3.5-9B"}
COORD_SPACES = ["norm_0_1000", "norm_0_1", "pixel_abs"]

_PALETTE = [
    (239, 71, 111), (17, 138, 178), (6, 214, 160), (255, 209, 102),
    (155, 93, 229), (241, 91, 181), (0, 187, 249), (254, 127, 45),
]


# ═════════════════════════════════════════════════════════════════════════════
# Model residency (teardown-friendly)
# ═════════════════════════════════════════════════════════════════════════════

_PIPE = None
_OCR = None
_AGE = None
_STRUCT: dict = {}


def get_pipe():
    """The always-on specialist pipeline (GroundingDINO/SAM/Depth/SigLIP2[/OCR])."""
    global _PIPE
    if _PIPE is None:
        with_ocr = os.environ.get("SPACE_WITH_OCR", "1") == "1"
        _PIPE = g.SpecialistPipeline(device=DEVICE, with_ocr=with_ocr)
    return _PIPE


def _get_ocr(pipe):
    """OCR reader — from the pipeline if it loaded there, else a lazy singleton
    (the teardown-safe fallback when EasyOCR misbehaves at module level)."""
    global _OCR
    if pipe.ocr is not None:
        return pipe.ocr
    if _OCR is None:
        _OCR = g.load_ocr(DEVICE)
    return _OCR


def _get_age_filter():
    """Age-gate pre-filter — imported lazily (the module loads its model at import)."""
    global _AGE
    if _AGE is None:
        import importlib

        faf = importlib.import_module("face_age_filter")
        _AGE = faf.FaceAgeFilter(decision_mode="strict", batch_size=32)
    return _AGE


class _Structurer:
    """Caption→struct (slot-registry JSON), mirroring the production ModelPack."""

    def __init__(self, model_id: str):
        from transformers import AutoProcessor

        try:
            from transformers import AutoModelForMultimodalLM as _M
        except ImportError:  # pragma: no cover
            from transformers import AutoModelForImageTextToText as _M

        self.proc = AutoProcessor.from_pretrained(model_id)
        tok = getattr(self.proc, "tokenizer", self.proc)
        tok.padding_side = "left"
        if tok.pad_token_id is None:
            tok.pad_token = tok.eos_token
        self.pad_id = tok.pad_token_id
        if DEVICE == "cuda":
            dtype = torch.bfloat16 if torch.cuda.is_bf16_supported() else torch.float16
            self.model = _M.from_pretrained(model_id, dtype=dtype, device_map="cuda").eval()
        else:
            self.model = _M.from_pretrained(model_id).to("cpu").eval()

    @torch.no_grad()
    def structure(self, captions: list, max_tok: int = 512) -> list:
        msgs = [[{"role": "system", "content": SYSTEM_PROMPT_JSON},
                 {"role": "user", "content": c}] for c in captions]
        enc = self.proc.apply_chat_template(
            msgs, add_generation_prompt=True, tokenize=True, return_dict=True,
            return_tensors="pt", padding=True, enable_thinking=False).to(self.model.device)
        n_in = enc["input_ids"].shape[1]
        gen = self.model.generate(**enc, max_new_tokens=max_tok, do_sample=False,
                                  pad_token_id=self.pad_id)
        outs = [self.proc.decode(s, skip_special_tokens=True).strip() for s in gen[:, n_in:]]
        structs = []
        for raw in outs:
            pr = parse_safely(raw)
            if pr.schema_valid and pr.parsed is not None:
                m = pr.parsed
                structs.append(m.model_dump() if hasattr(m, "model_dump") else m.dict())
            else:
                structs.append(None)
        return structs


def _get_structurer(model_id: str):
    if model_id in _STRUCT:
        return _STRUCT[model_id]
    _STRUCT.clear()                      # single-resident (evict on switch)
    if DEVICE == "cuda":
        torch.cuda.empty_cache()
    _STRUCT[model_id] = _Structurer(model_id)
    return _STRUCT[model_id]


# Preload the always-on specialists at module level on a GPU/ZeroGPU env
# (lazy on a CPU dev box so the module imports cheaply for tests).
if IS_GPU_ENV:
    try:
        get_pipe()
    except Exception as e:  # pragma: no cover
        print(f"[app] specialist preload deferred: {type(e).__name__}: {e}")


# ═════════════════════════════════════════════════════════════════════════════
# Solidify orchestration (public batched primitives + threshold / skip control)
# ═════════════════════════════════════════════════════════════════════════════

def _resolve_vocab(vocab_choice: str, custom: str) -> list:
    if vocab_choice == "custom":
        toks = [t.strip() for t in (custom or "").split(",") if t.strip()]
        return toks or g.COCO_CLASSES
    return VOCABS.get(vocab_choice, g.COCO_CLASSES)


def _solidify(pipe, images, vocab, phrases_list, box_thr, text_thr,
              use_ocr, use_masks, use_depth, batch, gdino_batch) -> list:
    """Mirror SpecialistPipeline.solidify_batch, but pass detection thresholds and
    honour the specialist on/off toggles."""
    images = list(images)
    solids = []
    ocr_reader = _get_ocr(pipe) if use_ocr else None
    for start in range(0, len(images), batch):
        chunk = images[start:start + batch]
        p_chunk = phrases_list[start:start + batch] if phrases_list is not None else None

        boxes_list = []
        for s2 in range(0, len(chunk), gdino_batch):
            boxes_list.extend(g.detect_batch(
                pipe.gdino, chunk[s2:s2 + gdino_batch], vocab,
                box_threshold=box_thr, text_threshold=text_thr, device=DEVICE))
        if use_masks and pipe.sam is not None:
            boxes_list = g.segment_batch(pipe.sam, chunk, boxes_list, device=DEVICE)
        depths = (g.depth_map_batch(pipe.depth, chunk)
                  if (use_depth and pipe.depth is not None) else [None] * len(chunk))
        classes = (g.zero_shot_batch(pipe.siglip, chunk, vocab, device=DEVICE)
                   if pipe.siglip is not None else [None] * len(chunk))
        styles = (g.zero_shot_batch(pipe.siglip, chunk, g.STYLE_LABELS, device=DEVICE)
                  if pipe.siglip is not None else [None] * len(chunk))
        if p_chunk is not None and any(p_chunk):
            attrs = []
            for s2 in range(0, len(chunk), gdino_batch):
                attrs.extend(g.ground_phrases_batch(
                    pipe.gdino, chunk[s2:s2 + gdino_batch],
                    p_chunk[s2:s2 + gdino_batch], device=DEVICE))
        else:
            attrs = [[] for _ in chunk]

        for k, im in enumerate(chunk):
            s = Solids(size=im.size)
            s.boxes = boxes_list[k]
            s.depth = depths[k]
            s.gray = np.asarray(im.convert("L"), dtype=np.float32)
            if classes[k] is not None:
                s.class_top = classes[k][:5]
                s.style = styles[k][0]["label"]
            if ocr_reader is not None:
                s.ocr = g.ocr_read(ocr_reader, im)
            s.attr_boxes = attrs[k]
            solids.append(s)
    return solids


def _solidify_oom(pipe, images, vocab, phrases_list, box_thr, text_thr,
                  use_ocr, use_masks, use_depth, batch, gdino_batch) -> list:
    """OOM-halving wrapper (mirrors produce_fused_dataset's guard)."""
    solids, i, bs = [], 0, batch
    while i < len(images):
        chunk = images[i:i + bs]
        p_chunk = phrases_list[i:i + bs] if phrases_list is not None else None
        try:
            solids.extend(_solidify(pipe, chunk, vocab, p_chunk, box_thr, text_thr,
                                    use_ocr, use_masks, use_depth, bs, gdino_batch))
            i += len(chunk)
            bs = batch
        except torch.cuda.OutOfMemoryError:  # pragma: no cover
            torch.cuda.empty_cache()
            if bs == 1:
                solids.append(Solids(size=images[i].size))
                i += 1
                bs = batch
            else:
                bs = max(1, bs // 2)
    return solids


# ═════════════════════════════════════════════════════════════════════════════
# GPU stage (everything that touches CUDA) — teardown-friendly
# ═════════════════════════════════════════════════════════════════════════════

def _gpu_duration(images, *_a, **_k):
    n = len(images) if images else 1
    return int(min(240, 25 + 9 * n))


@spaces.GPU(duration=_gpu_duration)
def gpu_extract(images, vocab, box_thr, text_thr, use_ocr, use_masks, use_depth,
                structurer_id, captions_list, use_age, batch, gdino_batch, render):
    """All CUDA work in one allocation. Returns picklable CPU data:
    per-image (tasks, digest, overlays) + caption structs + age audits + timing."""
    pipe = get_pipe()
    timing = {}
    n = len(images)

    audits = None
    if use_age:
        t = time.perf_counter()
        audits = [r.to_audit() for r in _get_age_filter().check_batch(images)]
        timing["age_s"] = round(time.perf_counter() - t, 3)

    structs_rows = [{} for _ in images]
    raws_rows = [{} for _ in images]
    if structurer_id and any(captions_list or []):
        t = time.perf_counter()
        st = _get_structurer(structurer_id)
        for idx, caps in enumerate(captions_list or []):
            caps = [c for c in (caps or []) if c and str(c).strip()]
            if not caps:
                continue
            got = st.structure(caps)
            structs_rows[idx] = {f"cap_{j}": s for j, s in enumerate(got)}
            raws_rows[idx] = {f"cap_{j}": c for j, c in enumerate(caps)}
        timing["struct_s"] = round(time.perf_counter() - t, 3)

    phrases_list = [phrases_for_grounding(sr) for sr in structs_rows]

    t = time.perf_counter()
    solids = _solidify_oom(pipe, images, vocab, phrases_list, box_thr, text_thr,
                           use_ocr, use_masks, use_depth, batch, gdino_batch)
    timing["extract_s"] = round(time.perf_counter() - t, 3)

    results = []
    for s, im in zip(solids, images):
        tasks = g.SpecialistPipeline._build_tasks(s)      # CPU, torch-free, fast
        digest = solids_digest(s)
        overlays = _render_overlays(im, s) if render else None
        results.append({"tasks": tasks, "digest": digest, "overlays": overlays})

    if DEVICE == "cuda":
        torch.cuda.empty_cache()
    timing["n_images"] = n
    return results, structs_rows, raws_rows, audits, timing


# ═════════════════════════════════════════════════════════════════════════════
# CPU fusion + assembly (no GPU held)
# ═════════════════════════════════════════════════════════════════════════════

def _task_valid(task: str, pred) -> bool:
    try:
        m = model_for(get_task(task))
        m.model_validate(pred) if hasattr(m, "model_validate") else m(**pred)
        return True
    except Exception:
        return False


def _assemble(results, structs_rows, raws_rows, audits, sizes,
              t_own, t_margin, dedup_iou, coord_space, task_filter):
    """Fuse each image's digest + structs → scene + prompt + one output row."""
    rows = []
    cs = CoordSpace(coord_space)
    for i, r in enumerate(results):
        tasks = dict(r["tasks"])
        try:
            scene = fuse(r["digest"], structs_rows[i] or {}, raws_rows[i] or {},
                         t_own=t_own, t_margin=t_margin, dedup_iou=dedup_iou, coord_space=cs)
            tasks["semantic_association"] = build_semantic_association(scene)
            prompt = fused_prompt(scene)
            conf = float(scene["quality"]["overall_confidence"])
        except Exception as e:                          # pragma: no cover
            scene, prompt, conf = {"__error__": f"{type(e).__name__}: {e}"}, "", 0.0
        valid = {t: _task_valid(t, p) for t, p in tasks.items() if t != "__error__"}
        shown = {t: tasks[t] for t in tasks if (not task_filter or t in task_filter)}
        W, H = sizes[i]
        rows.append({
            "tasks_json": shown, "tasks_valid": valid, "fused_json": scene,
            "prompt_fused": prompt, "fusion_confidence": round(conf, 4),
            "struct": structs_rows[i] or {}, "age_audit": (audits[i] if audits else None),
            "proc_width": W, "proc_height": H,
            "overlays": r.get("overlays"),
        })
    return rows


def _download_row(row: dict) -> str:
    payload = {
        "tasks_json": json.dumps(row["tasks_json"]),
        "tasks_valid": json.dumps(row["tasks_valid"]),
        "fused_json": json.dumps(row["fused_json"]),
        "prompt_fused": row["prompt_fused"],
        "fusion_confidence": row["fusion_confidence"],
        "struct": json.dumps(row["struct"]),
        "age_audit": json.dumps(row["age_audit"]),
        "proc_width": row["proc_width"], "proc_height": row["proc_height"],
    }
    f = tempfile.NamedTemporaryFile("w", suffix=".json", delete=False, encoding="utf-8")
    json.dump(payload, f, indent=2)
    f.close()
    return f.name


# ═════════════════════════════════════════════════════════════════════════════
# Overlay rendering (from Solids, pixel space)
# ═════════════════════════════════════════════════════════════════════════════

def _colorize_depth(depth: np.ndarray) -> Image.Image:
    d = np.asarray(depth, dtype=np.float32)
    lo, hi = float(d.min()), float(d.max())
    n = (d - lo) / (hi - lo + 1e-6)                     # 0=far, 1=near
    # 3-stop gradient far(indigo)→mid(teal)→near(amber)
    stops = np.array([[40, 30, 90], [17, 138, 178], [255, 209, 102]], dtype=np.float32)
    x = n * 2.0
    lo_i = np.clip(np.floor(x).astype(int), 0, 1)
    frac = (x - lo_i)[..., None]
    rgb = (stops[lo_i] * (1 - frac) + stops[lo_i + 1] * frac).astype(np.uint8)
    return Image.fromarray(rgb, "RGB")


def _render_overlays(image: Image.Image, s: Solids) -> dict:
    base = image.convert("RGB")
    annotated = base.copy()
    overlay = Image.new("RGBA", annotated.size, (0, 0, 0, 0))
    od = ImageDraw.Draw(overlay)
    dr = ImageDraw.Draw(annotated)

    for i, b in enumerate(s.boxes):
        color = _PALETTE[i % len(_PALETTE)]
        x1, y1, x2, y2 = [int(v) for v in b["box"]]
        mask = b.get("mask")
        if mask is not None:
            m = np.asarray(mask, dtype=bool)
            fill = np.zeros((*m.shape, 4), dtype=np.uint8)
            fill[m] = (*color, 90)
            overlay.alpha_composite(Image.fromarray(fill, "RGBA"))
        dr.rectangle([x1, y1, x2, y2], outline=color, width=3)
        label = f'{b.get("label", "?")} {b.get("score", 0):.2f}'
        dr.text((x1 + 3, max(0, y1 - 12)), label, fill=color)

    annotated = Image.alpha_composite(annotated.convert("RGBA"), overlay).convert("RGB")
    dr = ImageDraw.Draw(annotated)

    # subject box (thick white)
    subj = derive.subject_fixation(s.boxes, s.size).get("primary_subject", {})
    if subj.get("box"):
        x1, y1, x2, y2 = [int(v) for v in subj["box"]]
        dr.rectangle([x1, y1, x2, y2], outline=(255, 255, 255), width=4)

    # outline of the largest mask
    masked = [b for b in s.boxes if b.get("mask") is not None]
    if masked:
        big = max(masked, key=lambda b: np.asarray(b["mask"]).sum())
        poly = derive.outline_polygon(big["mask"], big["label"])["outline"]
        if len(poly) >= 6:
            pts = [(poly[j], poly[j + 1]) for j in range(0, len(poly) - 1, 2)]
            dr.line(pts + [pts[0]], fill=(255, 0, 128), width=2)

    depth_img = _colorize_depth(s.depth) if s.depth is not None else None
    return {"annotated": annotated, "depth": depth_img}


# ═════════════════════════════════════════════════════════════════════════════
# Gradio callbacks
# ═════════════════════════════════════════════════════════════════════════════

def _prep(image) -> Image.Image:
    im = image if isinstance(image, Image.Image) else Image.open(image)
    im = im.convert("RGB")
    if max(im.size) > MAX_DIM:
        im.thumbnail((MAX_DIM, MAX_DIM))
    return im


def run_single(image, vocab_choice, custom_vocab, tasks_sel, use_ocr, use_masks,
               use_depth, box_thr, text_thr, structurer_choice, captions_text,
               use_age, t_own, t_margin, dedup_iou, coord_space):
    if image is None:
        raise gr.Error("Upload or pick an image first.")
    im = _prep(image)
    vocab = _resolve_vocab(vocab_choice, custom_vocab)
    struct_id = STRUCTURERS.get(structurer_choice)
    caps = [c.strip() for c in (captions_text or "").splitlines() if c.strip()]

    results, structs_rows, raws_rows, audits, timing = gpu_extract(
        [im], vocab, float(box_thr), float(text_thr), bool(use_ocr), bool(use_masks),
        bool(use_depth), struct_id, [caps], bool(use_age), 1, 2, True)

    row = _assemble(results, structs_rows, raws_rows, audits, [im.size],
                    float(t_own), float(t_margin), float(dedup_iou), coord_space,
                    set(tasks_sel or []))[0]
    ov = row["overlays"] or {}
    return (
        ov.get("annotated"), ov.get("depth"),
        row["prompt_fused"], row["fusion_confidence"],
        row["tasks_json"], row["tasks_valid"], row["fused_json"],
        row["struct"], (row["age_audit"] or {}), timing,
        _download_row(row),
    )


def run_batch(files, vocab_choice, custom_vocab, use_ocr, use_masks, use_depth,
              box_thr, text_thr, structurer_choice, use_age, t_own, t_margin,
              dedup_iou, coord_space, batch, gdino_batch):
    if not files:
        raise gr.Error("Upload at least one image.")
    files = files[:BATCH_CAP]
    ims = [_prep(f) for f in files]
    vocab = _resolve_vocab(vocab_choice, custom_vocab)
    struct_id = STRUCTURERS.get(structurer_choice)

    t0 = time.perf_counter()
    results, structs_rows, raws_rows, audits, timing = gpu_extract(
        ims, vocab, float(box_thr), float(text_thr), bool(use_ocr), bool(use_masks),
        bool(use_depth), struct_id, [[] for _ in ims], bool(use_age),
        int(batch), int(gdino_batch), False)
    rows = _assemble(results, structs_rows, raws_rows, audits, [im.size for im in ims],
                     float(t_own), float(t_margin), float(dedup_iou), coord_space, None)
    wall = time.perf_counter() - t0

    table, jsonl = [], []
    for i, row in enumerate(rows):
        cls = row["tasks_json"].get("image_classification", {}) if row["tasks_json"] else {}
        n_ent = len(row["fused_json"].get("entities", [])) if isinstance(row["fused_json"], dict) else 0
        nvalid = sum(1 for v in row["tasks_valid"].values() if v)
        table.append([i, cls.get("label", ""), n_ent, row["fusion_confidence"],
                      f"{nvalid}/{len(row['tasks_valid'])}", (row["prompt_fused"] or "")[:90]])
        jsonl.append({
            "idx": i, "tasks_json": json.dumps(row["tasks_json"]),
            "tasks_valid": json.dumps(row["tasks_valid"]),
            "fused_json": json.dumps(row["fused_json"]),
            "prompt_fused": row["prompt_fused"], "fusion_confidence": row["fusion_confidence"],
            "proc_width": row["proc_width"], "proc_height": row["proc_height"],
        })

    f = tempfile.NamedTemporaryFile("w", suffix=".jsonl", delete=False, encoding="utf-8")
    for r in jsonl:
        f.write(json.dumps(r) + "\n")
    f.close()

    summary = {
        "images": len(ims), "wall_s": round(wall, 2),
        "img_per_s": round(len(ims) / max(0.001, wall), 2),
        **{k: v for k, v in timing.items() if k != "n_images"},
    }
    return table, summary, f.name


# ═════════════════════════════════════════════════════════════════════════════
# UI
# ═════════════════════════════════════════════════════════════════════════════

def _controls():
    """Shared control widgets — returned so both tabs can wire them."""
    vocab_choice = gr.Radio(list(VOCABS) + ["custom"], value="COCO-80", label="Detection vocab")
    custom_vocab = gr.Textbox(label="Custom phrases (comma-separated)", visible=False,
                              placeholder="person, red circle, laptop")
    with gr.Row():
        use_ocr = gr.Checkbox(True, label="OCR")
        use_masks = gr.Checkbox(True, label="SAM masks")
        use_depth = gr.Checkbox(True, label="Depth")
    with gr.Row():
        box_thr = gr.Slider(0.05, 0.6, 0.30, step=0.01, label="box threshold")
        text_thr = gr.Slider(0.05, 0.6, 0.25, step=0.01, label="text threshold")
    structurer = gr.Radio(list(STRUCTURERS), value="off", label="Caption structurer")
    with gr.Row():
        t_own = gr.Slider(0.0, 1.0, 0.60, step=0.01, label="t_own")
        t_margin = gr.Slider(0.0, 1.0, 0.25, step=0.01, label="t_margin")
        dedup_iou = gr.Slider(0.0, 1.0, 0.75, step=0.01, label="dedup_iou")
    coord_space = gr.Radio(COORD_SPACES, value="norm_0_1000", label="Fused-scene coord space")
    use_age = gr.Checkbox(False, label="Age-gate pre-filter (nateraw/vit-age-classifier)")

    vocab_choice.change(lambda c: gr.update(visible=(c == "custom")),
                        vocab_choice, custom_vocab)
    return (vocab_choice, custom_vocab, use_ocr, use_masks, use_depth, box_thr,
            text_thr, structurer, coord_space, use_age, t_own, t_margin, dedup_iou)


with gr.Blocks(title="Qwen Runner Vision — 12-task extraction + fusion") as demo:
    gr.Markdown(
        "# 🧩 Qwen Runner Vision\n"
        "Deterministic **12-task** extraction + **fusion** — stick an image in, get the "
        "full JSON readout (task JSONs + `FusedScene` + fused prompt). Specialists run on "
        "**ZeroGPU**; fusion is CPU-only. Every option below is a live toggle."
    )

    with gr.Tab("Single image"):
        with gr.Row():
            with gr.Column(scale=1):
                img_in = gr.Image(type="pil", label="Image", height=320)
                tasks_sel = gr.CheckboxGroup(DET_TASKS, value=DET_TASKS,
                                             label="Tasks to show (all 12 always computed)")
                gr.CheckboxGroup(VLM_TASKS, label="VLM/DEFER (no deterministic builder)",
                                 interactive=False)
                captions = gr.Textbox(lines=3, label="Captions (one per line — enrich fusion)",
                                      placeholder="a woman with long red hair in a blue coat")
                with gr.Accordion("Settings", open=False):
                    ctl = _controls()
                run_b = gr.Button("Extract", variant="primary")
            with gr.Column(scale=1):
                with gr.Row():
                    annotated = gr.Image(label="Detections · masks · subject · outline", height=280)
                    depth_img = gr.Image(label="Depth (near → far)", height=280)
                prompt_out = gr.Textbox(label="Fused prompt", lines=3)
                conf_out = gr.Number(label="Fusion confidence")
                dl = gr.File(label="Download row (JSON)")
        with gr.Accordion("Full JSON readout", open=True):
            tasks_out = gr.JSON(label="tasks_json (12 tasks)")
            with gr.Row():
                valid_out = gr.JSON(label="tasks_valid")
                struct_out = gr.JSON(label="caption structs")
            fused_out = gr.JSON(label="fused_json (FusedScene)")
            with gr.Row():
                age_out = gr.JSON(label="age_audit")
                timing_out = gr.JSON(label="timing")

        (vocab_choice, custom_vocab, use_ocr, use_masks, use_depth, box_thr,
         text_thr, structurer, coord_space, use_age, t_own, t_margin, dedup_iou) = ctl

        ex_dir = os.path.join(os.path.dirname(__file__), "examples")
        if os.path.isdir(ex_dir):
            ex_imgs = [[os.path.join(ex_dir, f)] for f in sorted(os.listdir(ex_dir))
                       if f.lower().endswith((".png", ".jpg", ".jpeg"))]
            if ex_imgs:
                gr.Examples(ex_imgs, inputs=img_in, label="Examples")

        run_b.click(
            run_single,
            inputs=[img_in, vocab_choice, custom_vocab, tasks_sel, use_ocr, use_masks,
                    use_depth, box_thr, text_thr, structurer, captions, use_age,
                    t_own, t_margin, dedup_iou, coord_space],
            outputs=[annotated, depth_img, prompt_out, conf_out, tasks_out, valid_out,
                     fused_out, struct_out, age_out, timing_out, dl],
        )

    with gr.Tab("Batch (the batched structure)"):
        gr.Markdown(
            f"Upload up to **{BATCH_CAP}** images → batched `solidify_batch` on ZeroGPU "
            "(`gdino_batch=2`, GDINO anti-scales) + per-image CPU fusion. Throughput mirrors "
            "`runs/extract_throughput_results.md`."
        )
        with gr.Row():
            with gr.Column(scale=1):
                files_in = gr.Files(label="Images", file_types=["image"])
                with gr.Accordion("Settings", open=False):
                    bctl = _controls()
                with gr.Row():
                    batch_sl = gr.Slider(1, 24, 16, step=1, label="extract_batch")
                    gdino_sl = gr.Slider(1, 8, 2, step=1, label="gdino_batch (keep ~2)")
                run_batch_b = gr.Button("Run batch", variant="primary")
            with gr.Column(scale=1):
                batch_table = gr.Dataframe(
                    headers=["#", "label", "entities", "fusion_conf", "valid", "prompt…"],
                    label="Per-image results", wrap=True)
                batch_summary = gr.JSON(label="Throughput")
                batch_dl = gr.File(label="Download rows (JSONL)")

        (b_vocab, b_custom, b_ocr, b_masks, b_depth, b_box, b_text, b_struct,
         b_coord, b_age, b_town, b_tmargin, b_dedup) = bctl

        run_batch_b.click(
            run_batch,
            inputs=[files_in, b_vocab, b_custom, b_ocr, b_masks, b_depth, b_box, b_text,
                    b_struct, b_age, b_town, b_tmargin, b_dedup, b_coord, batch_sl, gdino_sl],
            outputs=[batch_table, batch_summary, batch_dl],
        )


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