#!/usr/bin/env python3
"""Auto-generate hurtboxes from pose estimation over rendered frames.

The single biggest modernization over the 1990s digitized workflow: BKB's
.zon files carry 9 HAND-PLACED boxes per frame; here a skeleton pass
generates them. Uses rtmlib (RTMPose/DWPose family) if installed:

    pip install rtmlib onnxruntime

Falls back to alpha-channel bounding analysis when no pose model is
available — cruder (3 stacked boxes from the silhouette) but functional.

Writes/updates the "hurtboxes" key of each move's data.json in place.
Hit boxes (active attack frames) stay hand-authored — that's ~4 frames
per attack and it's game design, not labor.

Usage:
    python3 pipeline/autohitbox.py characters/hero            # all moves
    python3 pipeline/autohitbox.py characters/hero --move jab
"""
import argparse
import json
from pathlib import Path

import numpy as np
from PIL import Image

try:
    from rtmlib import Body  # type: ignore
    HAVE_POSE = True
except ImportError:
    HAVE_POSE = False

# COCO keypoint groups -> the three classic zones
HEAD_KPS = [0, 1, 2, 3, 4]          # nose, eyes, ears
TORSO_KPS = [5, 6, 11, 12]          # shoulders, hips
LEG_KPS = [11, 12, 13, 14, 15, 16]  # hips, knees, ankles
PAD = {"head": 14, "torso": 18, "legs": 12}


def rects_from_pose(kps, scores, w, h, ground_y, cx):
    """keypoints (COCO 17) -> up to 3 local-space rects [x,y,w,h]."""
    out = []
    for name, idxs in (("head", HEAD_KPS), ("torso", TORSO_KPS), ("legs", LEG_KPS)):
        pts = [kps[i] for i in idxs if scores[i] > 0.35]
        if len(pts) < 2:
            continue
        xs = [p[0] for p in pts]
        ys = [p[1] for p in pts]
        pad = PAD[name]
        x0, x1 = min(xs) - pad, max(xs) + pad
        y0, y1 = min(ys) - pad, max(ys) + pad
        if name == "legs":
            y1 = ground_y  # legs always reach the ground pivot
        out.append([round(x0 - cx), round(y0 - ground_y),
                    round(x1 - x0), round(y1 - y0)])
    return out


def rects_from_alpha(img: Image.Image, ground_y, cx):
    """Fallback: slice the alpha silhouette into head/torso/legs thirds."""
    a = np.array(img.split()[-1])
    rows = np.where(a.max(axis=1) > 32)[0]
    if rows.size == 0:
        return []
    top, bot = int(rows[0]), int(rows[-1])
    height = bot - top
    bands = [(top, top + height // 4), (top + height // 4, top + height * 5 // 8),
             (top + height * 5 // 8, ground_y)]
    out = []
    for y0, y1 in bands:
        band = a[y0:max(y1, y0 + 1)]
        cols = np.where(band.max(axis=0) > 32)[0]
        if cols.size == 0:
            continue
        out.append([int(cols[0]) - cx, y0 - ground_y,
                    int(cols[-1] - cols[0]), y1 - y0])
    return out


def process_move(move_dir: Path, body) -> None:
    frames = sorted((move_dir / "frames").glob("*.png"))
    if not frames:
        return
    data_path = move_dir / "data.json"
    data = json.loads(data_path.read_text())
    total = data.get("total", len(frames))

    per_frame = {}
    for img_path in frames:
        img = Image.open(img_path).convert("RGBA")
        w, h = img.size
        cx, ground_y = w // 2, h - 12
        if body is not None:
            arr = np.array(img.convert("RGB"))[:, :, ::-1]  # BGR
            kps, scores = body(arr)
            rects = rects_from_pose(kps[0], scores[0], w, h, ground_y, cx) if len(kps) else []
        else:
            rects = rects_from_alpha(img, ground_y, cx)
        if rects:
            per_frame[img_path.stem] = rects

    if not per_frame:
        print(f"  {move_dir.name}: no silhouette found, skipped")
        return
    # default = the median frame's boxes; per-tick overrides where they differ a lot
    keys = sorted(per_frame)
    default = per_frame[keys[len(keys) // 2]]
    frames_out = {}
    n = len(keys)
    for i, k in enumerate(keys):
        if per_frame[k] != default:
            t0 = int(i * total / n)
            t1 = max(t0, int((i + 1) * total / n) - 1)
            frames_out[f"{t0}-{t1}"] = per_frame[k]
    data["hurtboxes"] = {"default": default}
    if frames_out:
        data["hurtboxes"]["frames"] = frames_out
    data_path.write_text(json.dumps(data, indent=2))
    print(f"  {move_dir.name}: default {len(default)} boxes, {len(frames_out)} frame overrides")


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("char_dir", help="characters/<id>")
    ap.add_argument("--move", default=None)
    args = ap.parse_args()

    body = None
    if HAVE_POSE:
        body = Body(mode="balanced", backend="onnxruntime", device="cpu")
        print("pose model: rtmlib Body")
    else:
        print("rtmlib not installed - using alpha-silhouette fallback "
              "(pip install rtmlib onnxruntime for pose-based boxes)")

    moves_root = Path(args.char_dir) / "moves"
    targets = [moves_root / args.move] if args.move else sorted(moves_root.iterdir())
    for move_dir in targets:
        if move_dir.is_dir():
            process_move(move_dir, body)


if __name__ == "__main__":
    main()