# bubble_detect.py — comic-text-detector ONNX -> per-panel TEXT-REGION boxes, for set-of-mark
# speaker attribution (see spec: decompose the VLM's implicit tail-tracing into grounded marks).
# Runs on the CPU EP so it never contends with gemma for the GPU. The model's `blk` head is a
# YOLO-style text-line detector; validated pixel-accurate on real panels + synthetic bubbles.
#
# We use ONLY the text-region boxes: a drawn, numbered region is the attribution anchor, and because
# region #k lives on panel N's own image it structurally prevents the window from smearing a line onto
# a neighbouring panel (the p108 bleed) or inventing dialogue that isn't drawn (the p108 phantom line).
# ponytail: full balloon-fill mask + tail-tip geometry are also in the `det`/`seg` heads but noisy and
#   webtoon bubbles are often tailless anyway — add them only if region marks prove insufficient.
import os
import numpy as np
import cv2

MODEL = os.environ.get("CTD_MODEL",
                       os.path.join(os.path.dirname(__file__), "models", "comictextdetector.pt.onnx"))
CONF = float(os.environ.get("CTD_CONF", "0.20"))   # blk head confidences run low (~0.55 max); tune per title
_sess = None


def _load():
    global _sess
    if _sess is None:
        import onnxruntime as ort
        _sess = ort.InferenceSession(MODEL, providers=["CPUExecutionProvider"])
    return _sess


def _letterbox(img, sz=1024):
    """resize keeping aspect ratio, pad to sz*sz with 114-grey (the model was trained letterboxed;
    a plain squash distorts tall webtoon panels and wrecks the masks). returns canvas + inverse params."""
    h, w = img.shape[:2]
    r = min(sz / h, sz / w)
    nh, nw = int(round(h * r)), int(round(w * r))
    canvas = np.full((sz, sz, 3), 114, np.uint8)
    py, px = (sz - nh) // 2, (sz - nw) // 2
    canvas[py:py + nh, px:px + nw] = cv2.resize(img, (nw, nh))
    return canvas, r, px, py


def detect_text_regions(img, conf: float = None) -> list:
    """img: BGR ndarray (one panel). -> [{"bbox":[x1,y1,x2,y2], "conf":float}] in ORIGINAL pixels,
    reading order (top->bottom, left->right). Empty list on no detections (caller falls back to the
    holistic prompt). Load + inference are lazy so importing this never touches the GPU or the model."""
    conf = CONF if conf is None else conf
    h, w = img.shape[:2]
    canvas, r, px, py = _letterbox(img)
    x = cv2.cvtColor(canvas, cv2.COLOR_BGR2RGB).astype(np.float32).transpose(2, 0, 1)[None] / 255.0
    blk = _load().run(None, {"images": x})[0][0]   # [N,7] = cx,cy,w,h,conf,cls0,cls1 in 1024 space
    m = blk[:, 4] > conf
    if not m.any():
        return []
    b = blk[m]
    rects = [[float(cx - bw / 2), float(cy - bh / 2), float(bw), float(bh)] for cx, cy, bw, bh in b[:, :4]]
    idx = cv2.dnn.NMSBoxes(rects, b[:, 4].tolist(), conf, 0.5)
    if len(idx) == 0:
        return []
    out = []
    for i in np.array(idx).flatten():
        cx, cy, bw, bh = b[i, :4]
        x1, y1 = max(0, int((cx - bw / 2 - px) / r)), max(0, int((cy - bh / 2 - py) / r))
        x2, y2 = min(w, int((cx + bw / 2 - px) / r)), min(h, int((cy + bh / 2 - py) / r))
        if x2 > x1 and y2 > y1:
            out.append({"bbox": [x1, y1, x2, y2], "conf": round(float(b[i, 4]), 3)})
    out.sort(key=lambda d: (d["bbox"][1], d["bbox"][0]))
    return out


def draw_region_marks(img, regions: list):
    """draw a numbered red box (#1..#N) over each text region on a COPY of img; return the marked
    image. the number is what gemma cites in set-of-mark attribution ("region 2 -> P1")."""
    vis = img.copy()
    for n, reg in enumerate(regions, 1):
        x1, y1, x2, y2 = reg["bbox"]
        cv2.rectangle(vis, (x1, y1), (x2, y2), (0, 0, 255), 2)
        cv2.rectangle(vis, (x1, max(0, y1 - 22)), (x1 + 26, y1), (0, 0, 255), -1)
        cv2.putText(vis, str(n), (x1 + 3, y1 - 5), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    return vis


def draw_face_marks(img, faces: list):
    """draw a labelled green box for each character face/body, IN PLACE, return img. faces:
    [{"label": "P1", "bbox": [x1,y1,x2,y2]}]. gemma names the label (P1/P2) as the region's speaker,
    which the worker maps back to that character's local_id."""
    for f in faces:
        x1, y1, x2, y2 = f["bbox"]
        cv2.rectangle(img, (x1, y1), (x2, y2), (0, 170, 0), 2)
        cv2.rectangle(img, (x1, y1), (min(x2, x1 + 40), y1 + 22), (0, 170, 0), -1)
        cv2.putText(img, f["label"], (x1 + 3, y1 + 17), cv2.FONT_HERSHEY_SIMPLEX, 0.7, (255, 255, 255), 2)
    return img


if __name__ == "__main__":
    # self-check: the real ONNX model must find the text inside a synthetic speech bubble, and the
    # returned box must overlap the text we drew. Exercises letterbox + decode + NMS end to end.
    panel = np.full((600, 400, 3), 200, np.uint8)
    cv2.ellipse(panel, (200, 200), (150, 90), 0, 0, 360, (255, 255, 255), -1)
    cv2.ellipse(panel, (200, 200), (150, 90), 0, 0, 360, (0, 0, 0), 3)
    cv2.putText(panel, "HELLO", (110, 215), cv2.FONT_HERSHEY_SIMPLEX, 1.8, (0, 0, 0), 5)
    regs = detect_text_regions(panel)
    assert regs, "no text region detected in synthetic bubble"
    # at least one box should contain the drawn text centre (~200,200)
    hit = any(x1 <= 200 <= x2 and y1 <= 200 <= y2 for (x1, y1, x2, y2) in (r["bbox"] for r in regs))
    assert hit, f"no region covers the text centre: {regs}"
    marked = draw_region_marks(panel, regs)
    assert marked.shape == panel.shape and marked is not panel
    # face marks: labelled box drawn in place; sample the left border low enough to miss the bubble
    draw_face_marks(marked, [{"label": "P1", "bbox": [50, 180, 350, 500]}])
    assert tuple(int(v) for v in marked[490, 51]) == (0, 170, 0), f"face mark not green: {marked[490, 51]}"
    print(f"bubble_detect self-check ok ({len(regs)} region(s), top conf {regs[0]['conf']})")