# face_detect.py — anime face detector (YOLOv8, ONNX) -> real character-face boxes for set-of-mark
# speaker attribution. The vision DETECT stage returns gemma-guessed bboxes which are too imprecise to
# draw a mark on; this gives grounded boxes the same way bubble_detect gives grounded text regions.
# Model: deepghs/anime_face_detection face_detect_v1.4_s (single class "face"). CPU EP -> no GPU contention.
# ponytail: face boxes only (not full body). Attribution just needs "which face said this"; add a
#   person head only if off-panel/back-turned speakers need body grounding.
import os
import numpy as np
import cv2

MODEL = os.environ.get("FACE_MODEL",
                       os.path.join(os.path.dirname(__file__), "models", "anime_face_v1.4_s.onnx"))
CONF = float(os.environ.get("FACE_CONF", "0.30"))
_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=640):
    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_faces(img, conf: float | None = None) -> list:
    """img: BGR ndarray. -> [{"bbox":[x1,y1,x2,y2], "conf":float}] in ORIGINAL pixels, sorted
    reading order. Empty on no faces. Lazy load so import never touches 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
    out = _load().run(None, {"images": x})[0][0].T   # (8400,5) = cx,cy,w,h,conf in 640 space
    m = out[:, 4] > conf
    if not m.any():
        return []
    b = out[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.45)
    if len(idx) == 0:
        return []
    res = []
    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:
            res.append({"bbox": [x1, y1, x2, y2], "conf": round(float(b[i, 4]), 3)})
    res.sort(key=lambda d: (d["bbox"][1], d["bbox"][0]))
    return res


if __name__ == "__main__":
    # self-check needs a real anime face; use a page from the sample series if present, else skip loudly.
    import glob
    pages = glob.glob("/mnt/server/mnt/hdd2/manga-raw/*/*.png")
    if not pages:
        print("face_detect: no sample page available, skipping live check")
    else:
        faces = detect_faces(cv2.imread(sorted(pages)[13]))
        assert faces, "no face detected on a page that should have characters"
        for f in faces:
            x1, y1, x2, y2 = f["bbox"]
            assert x2 > x1 and y2 > y1 and f["conf"] >= CONF
        print(f"face_detect self-check ok ({len(faces)} face(s), top conf {faces[0]['conf']})")