src/poolcoach_cv/overlay.py

"""Render overlay.mp4 per cú — lát A2 phần 3 (13/08/2026, BRIEF A2).

Tái dụng cách vẽ ``render_overlay`` của ``scripts/broadcast/track_p0.py``
(bản Danh đã xem và chốt 13/08 vì trực quan — đuôi track + header trạng
thái), nâng cấp theo design §3.2: vẽ MỌI bi (không chỉ cue), màu theo số
bi (BallID hai tầng chạy per-frame trên chính crop bbox), chấm sự kiện
va chạm suy ra. Script P0 giữ NGUYÊN không sửa — đây là code path mới
cho app (nếp broadcast.py).

Render chạy TRONG job analyze, TRƯỚC khi worker xoá clip (bài học lần đo
A1: clip đã mất thì overlay không render nổi hồi tố). Detection per frame
KHÔNG detect lại: ``RecordingModel`` bọc YOLO quanh MỘT lần gọi
``analyze_clip`` và ghi lại đúng những bbox mà tầng số đã dùng (nếp
"adapter quanh một lần gọi, không sửa broadcast.py" — BG26/BG28/BG31);
frame đọc lại từ chính clip còn trên đĩa.

Chữ trên frame là ASCII không dấu — cv2.putText (font Hershey) không vẽ
được tiếng Việt có dấu; phần chữ có dấu sống ở FE/JSON.

Encode hai bước: cv2 VideoWriter mp4v (như P0) rồi ffmpeg re-encode H.264
yuv420p +faststart — mp4v mở được bằng player thường nhưng <video> của
trình duyệt KHÔNG phát được; ffmpeg đã là điều kiện sẵn của luồng video
đa cú (cut_clip). ffmpeg vắng → giữ bản mp4v + warning (player thường
vẫn mở — gate A2.2 vế "ngoài app" vẫn sống).

Import cv2 LƯỜI trong hàm (nếp poolcoach_cv) — venv app import module
này được cho unit test phần thuần.
"""

from __future__ import annotations

import os
import shutil
import subprocess
from pathlib import Path

import numpy as np

from poolcoach_cv.broadcast import BALL_CLASSES, BALL_R_M

TRAIL_N = 90            # số điểm track vẽ đuôi cue ball — như P0

# Màu bi theo số (BGR) — khớp bảng COLORS của FE (app.js) để hai nơi kể
# cùng một câu chuyện màu; None/không nhận ra → xám.
BALL_BGR = {
    1: (24, 197, 245), 2: (225, 105, 65), 3: (47, 52, 224),
    4: (191, 63, 125), 5: (19, 127, 240), 6: (70, 158, 46),
    7: (32, 32, 139), 8: (40, 40, 40), 9: (24, 197, 245),
}
CUE_BGR = (0, 220, 220)         # như P0
UNKNOWN_BGR = (180, 180, 180)
TRAIL_BGR = (0, 220, 0)
COLLISION_BGR = (0, 215, 255)   # vàng gold — khớp dấu ✕ trên FE


class RecordingModel:
    """Bọc YOLO model: forward nguyên ``predict``, ghi lại per-frame
    ``[(cls_name, conf, (x1, y1, x2, y2)), ...]`` theo ĐÚNG thứ tự frame
    mà analyze_clip đưa vào (batch tuần tự). Chỉ dùng quanh MỘT lần gọi
    analyze_clip — model gốc không đổi gì."""

    def __init__(self, model) -> None:
        self._model = model
        self.frames: list[list[tuple]] = []

    def predict(self, imgs, **kw):
        results = self._model.predict(imgs, **kw)
        seq = imgs if isinstance(imgs, (list, tuple)) else [imgs]
        for res in results[:len(seq)]:
            names = res.names
            dets = []
            for b in res.boxes:
                x1, y1, x2, y2 = (float(v) for v in b.xyxy[0])
                dets.append((names[int(b.cls)], float(b.conf),
                             (x1, y1, x2, y2)))
            self.frames.append(dets)
        return results


def _uncomp_xy(pts: np.ndarray, height_comp: dict | None) -> np.ndarray:
    """Đảo bù độ cao tâm bi: toạ độ P trong JSON (đã bù) → X' vị trí map
    mặt vải — chính là chỗ bi HIỆN RA trên frame khi chiếu ngược qua
    homography. Nghịch đảo đóng của camera.HeightCompensation.apply:
    X' = P·(1+k) − k·C, k = R/(h−R). Cờ off/thiếu camera → giữ nguyên."""
    if not height_comp or not height_comp.get("on"):
        return pts
    h = height_comp.get("h_m")
    cx, cy = height_comp.get("cx_m"), height_comp.get("cy_m")
    if h is None or cx is None or cy is None:
        return pts
    k = BALL_R_M / (h - BALL_R_M)
    return np.asarray(pts, dtype=np.float64) * (1.0 + k) \
        - k * np.array([cx, cy])


def _reencode_h264(src: Path, dst: Path) -> str | None:
    """mp4v → H.264 yuv420p +faststart để <video> trình duyệt phát được.
    Trả None khi OK; message warning khi ffmpeg vắng/hỏng (giữ bản mp4v)."""
    exe = shutil.which("ffmpeg")
    if exe is None:
        os.replace(src, dst)
        return ("máy không có ffmpeg — overlay giữ codec mp4v (player "
                "thường mở được, trình duyệt có thể không phát)")
    try:
        r = subprocess.run(
            [exe, "-y", "-v", "error", "-i", str(src), "-c:v", "libx264",
             "-preset", "veryfast", "-crf", "23", "-pix_fmt", "yuv420p",
             "-movflags", "+faststart", str(dst)],
            capture_output=True, text=True, timeout=300)
    except subprocess.TimeoutExpired:
        os.replace(src, dst)
        return "ffmpeg re-encode overlay quá 300s — overlay giữ codec mp4v"
    if r.returncode != 0 or not dst.exists():
        os.replace(src, dst)
        err = (r.stderr or "?").strip().splitlines()[-1][:120]
        return f"ffmpeg re-encode overlay lỗi ({err}) — giữ codec mp4v"
    src.unlink(missing_ok=True)
    return None


def render_shot_overlay(clip_path, corners_px, result: dict,
                        frames_dets: list[list[tuple]], out_path,
                        label: str = "", beat=None) -> str | None:
    """Render overlay MỘT cú từ clip còn trên đĩa + detection đã ghi.

    ``result``: JSON kết quả analyze (track/collisions/height_comp/
    suspect_not_shot — toạ độ bàn, t_s tương đối frame đầu clip).
    ``frames_dets``: từ RecordingModel — index = thứ tự frame.
    ``label``: chữ đầu header (ASCII), ví dụ "shot 03". ``beat``: callback
    gọi định kỳ để giữ heartbeat trong job dài. Trả None khi trơn, message
    warning (tiếng Việt) khi có chuyện nhưng file vẫn ra được; ném khi
    không ra nổi file.
    """
    import cv2

    from poolcoach_cv.ballid import identify_balls
    from poolcoach_cv.broadcast import TABLE_L_M, TABLE_W_M
    from poolcoach_cv.homography import TableHomography, orient_corners

    out_path = Path(out_path)
    th = TableHomography(orient_corners(np.asarray(corners_px,
                                                   dtype=np.float64)),
                         table_w=TABLE_W_M, table_l=TABLE_L_M)
    hc = result.get("height_comp")
    track = result.get("track") or []
    trail_t = np.array([float(p["t_s"]) for p in track])
    trail_px = (th.table_to_px(_uncomp_xy(
        np.array([[float(p["x_m"]), float(p["y_m"])] for p in track]), hc))
        if track else np.empty((0, 2)))
    cols = result.get("collisions") or []
    col_t = [float(c["t_s"]) for c in cols]
    col_px = (th.table_to_px(_uncomp_xy(
        np.array([[float(c["x_m"]), float(c["y_m"])] for c in cols]), hc))
        if cols else np.empty((0, 2)))
    suspect = bool((result.get("suspect_not_shot") or {}).get("flagged"))

    cap = cv2.VideoCapture(str(clip_path))
    if not cap.isOpened():
        raise ValueError("Không mở lại được clip để render overlay.")
    # temp CẠNH file đích (không phải system temp): os.replace cần cùng ổ
    # đĩa, và mkstemp giữ fd mở làm Windows khoá file không xoá được
    out_path.parent.mkdir(parents=True, exist_ok=True)
    tmp = out_path.with_name(out_path.stem + ".raw.mp4")
    vw = None
    try:
        fps = cap.get(cv2.CAP_PROP_FPS) or 30.0
        n_total = len(frames_dets)
        i = -1
        t_first = None
        while True:
            ok, frame = cap.read()
            if not ok:
                break
            i += 1
            t = cap.get(cv2.CAP_PROP_POS_MSEC) / 1000.0
            # POS_MSEC đọc SAU read() là PTS frame KẾ — xấp xỉ đủ cho vẽ;
            # số đo thật vẫn nằm ở JSON, overlay là lớp mắt người kiểm
            if t_first is None:
                t_first = t
            t_rel = t - t_first
            if vw is None:
                h, w = frame.shape[:2]
                vw = cv2.VideoWriter(str(tmp),
                                     cv2.VideoWriter_fourcc(*"mp4v"),
                                     fps, (w, h))
            if beat is not None and i % 32 == 0:
                beat()

            dets = frames_dets[i] if i < n_total else []
            balls = [(cls, conf, box) for cls, conf, box in dets
                     if cls in BALL_CLASSES]
            others = [b for b in balls if b[0] != "Cue"]
            cue = next((b for b in balls if b[0] == "Cue"), None)
            assigns = []
            if others:
                try:
                    assigns, _wb = identify_balls(
                        frame, [b[2] for b in others],
                        cue[2] if cue else None)
                except Exception:  # noqa: BLE001 — màu hỏng ≠ mất overlay
                    assigns = [(None, 0.0)] * len(others)

            # mọi bi: khung màu theo số + nhãn số (design §3.2)
            for (cls, _conf, (x1, y1, x2, y2)), (num, _nc) in zip(
                    others, assigns or [(None, 0.0)] * len(others)):
                col = BALL_BGR.get(num, UNKNOWN_BGR)
                cv2.rectangle(frame, (int(x1), int(y1)),
                              (int(x2), int(y2)), col, 2)
                if num is not None:
                    cv2.putText(frame, str(num), (int(x1), int(y1) - 4),
                                cv2.FONT_HERSHEY_SIMPLEX, 0.55, col, 2)
            if cue is not None:
                x1, y1, x2, y2 = cue[2]
                cv2.rectangle(frame, (int(x1), int(y1)),
                              (int(x2), int(y2)), CUE_BGR, 2)

            # đuôi track cue ball (như P0) — điểm đã đo tới thời điểm này
            m = trail_t <= t_rel + 1e-6
            pts = trail_px[m][-TRAIL_N:]
            if len(pts) >= 2:
                cv2.polylines(frame, [pts.astype(np.int32)], False,
                              TRAIL_BGR, 2)
            if len(pts):
                cv2.circle(frame, tuple(pts[-1].astype(int)), 14,
                           TRAIL_BGR, 2)

            # chấm sự kiện va chạm — hiện từ lúc xảy ra tới hết clip
            for tc, (px, py), c in zip(col_t, col_px, cols):
                if t_rel + 1e-6 < tc:
                    continue
                p = (int(px), int(py))
                s = 9
                cv2.line(frame, (p[0] - s, p[1] - s), (p[0] + s, p[1] + s),
                         COLLISION_BGR, 2)
                cv2.line(frame, (p[0] - s, p[1] + s), (p[0] + s, p[1] - s),
                         COLLISION_BGR, 2)
                cv2.putText(frame, f"{tc:.2f}s {c.get('contact') or ''}",
                            (p[0] + s + 3, p[1] - 3),
                            cv2.FONT_HERSHEY_SIMPLEX, 0.5, COLLISION_BGR, 1)

            # header trạng thái (ASCII — Hershey không có dấu)
            n_seen = int((trail_t <= t_rel + 1e-6).sum())
            head = (f"{label} t={t_rel:6.2f}s frame {i + 1}"
                    f" track {n_seen}/{len(trail_t)}")
            if suspect:
                head += "  [NGHI KHONG PHAI CU DANH]"
            cv2.rectangle(frame, (0, 0), (frame.shape[1], 30), (0, 0, 0), -1)
            cv2.putText(frame, head, (8, 22), cv2.FONT_HERSHEY_SIMPLEX,
                        0.6, (0, 0, 255) if suspect else (0, 220, 0), 2)
            vw.write(frame)
    finally:
        if vw is not None:
            vw.release()
        cap.release()
    if vw is None:
        tmp.unlink(missing_ok=True)
        raise ValueError("Clip không đọc được frame nào để render overlay.")
    return _reencode_h264(tmp, out_path)