#!/usr/bin/env python3
"""
Create a NEXT11 keypose image-source plan from an IK/exercise index.

This does not download copyrighted images. It creates query strings, per-exercise
source folders, and a contract scaffold so the next chat/operator can collect
candidate stills and then MediaPipe them.
"""
from __future__ import annotations

import argparse
import json
import re
from pathlib import Path
from typing import Any, Dict, List

DEFAULT_KEYPOSES_BY_FAMILY = [
    {"phaseId":"start_or_setup", "visualIntent":"exercise start/setup; full body visible; contacts visible"},
    {"phaseId":"load_or_descent", "visualIntent":"loading/descent phase; joints beginning to flex/prepare"},
    {"phaseId":"deep_or_bottom", "visualIntent":"maximum range/key position; contacts and limb angles visible"},
    {"phaseId":"return_or_ascent", "visualIntent":"return/ascent phase; compare joint path and drift"},
    {"phaseId":"finish", "visualIntent":"finish position; compare against start for drift"},
]

SPECIAL_KEYPOSES = {
    "push_up": [
        {"phaseId":"top_plank_start", "visualIntent":"side view; straight arms; rigid plank; hands and toes visible"},
        {"phaseId":"descent_mid", "visualIntent":"side view; elbows flexing; shoulders tracking over wrists; body line preserved"},
        {"phaseId":"bottom_chest_near_floor", "visualIntent":"side view; chest close to floor; elbows flexed; pelvis not sagging"},
        {"phaseId":"press_mid", "visualIntent":"side view; pressing up; body line and hand/foot contacts unchanged"},
        {"phaseId":"top_plank_finish", "visualIntent":"side view; returned to straight-arm plank; compare drift vs start"},
    ],
    "pushups": [
        {"phaseId":"top_plank_start", "visualIntent":"side view; straight arms; rigid plank; hands and toes visible"},
        {"phaseId":"descent_mid", "visualIntent":"side view; elbows flexing; shoulders tracking over wrists; body line preserved"},
        {"phaseId":"bottom_chest_near_floor", "visualIntent":"side view; chest close to floor; elbows flexed; pelvis not sagging"},
        {"phaseId":"press_mid", "visualIntent":"side view; pressing up; body line and hand/foot contacts unchanged"},
        {"phaseId":"top_plank_finish", "visualIntent":"side view; returned to straight-arm plank; compare drift vs start"},
    ],
    "linear_pogo_jumps": [
        {"phaseId":"ground_load", "visualIntent":"front/side view; forefoot contact; slight ankle/knee flex; upright torso"},
        {"phaseId":"takeoff", "visualIntent":"ankle plantarflexion; body rising; minimal knee bend"},
        {"phaseId":"flight_apex", "visualIntent":"both feet airborne; upright torso; hip height peak"},
        {"phaseId":"landing_contact", "visualIntent":"forefoot landing contact; soft ankle/knee; no collapse"},
        {"phaseId":"rebound", "visualIntent":"immediate next takeoff; low ground contact time; minimal drift"},
    ],
    "goblet_squat": [
        {"phaseId":"standing_top", "visualIntent":"front/45 view; kettlebell/ball close to sternum; feet visible"},
        {"phaseId":"descent_mid", "visualIntent":"hips/knees flexing; torso angle; object fixed to torso"},
        {"phaseId":"bottom", "visualIntent":"deep squat; pelvis/root low; knees/feet geometry; object not penetrating body"},
        {"phaseId":"ascent_mid", "visualIntent":"returning up; object remains fixed; no left/right drift"},
        {"phaseId":"standing_finish", "visualIntent":"standing top; compare feet/root/object against start"},
    ],
    "plank": [
        {"phaseId":"setup", "visualIntent":"side view; elbows/hands and toes visible; body line visible"},
        {"phaseId":"hold_mid", "visualIntent":"stable plank hold; pelvis, thorax, head in line"},
        {"phaseId":"hold_late", "visualIntent":"late hold; detect sag/pike/drift"},
    ],
}


def norm(s: str) -> str:
    return re.sub(r"[^a-z0-9]+", "_", s.lower()).strip("_")


def display_to_query(display: str) -> str:
    q = re.sub(r"[_-]+", " ", display).strip()
    return q


def choose_keyposes(ex: Dict[str, Any]) -> List[Dict[str, str]]:
    key = norm(str(ex.get("exerciseKey") or ""))
    display = norm(str(ex.get("displayName") or ""))
    joined = key + " " + display
    for special, poses in SPECIAL_KEYPOSES.items():
        if special in joined:
            return poses
    phases = ex.get("phases") or []
    if phases:
        return [{"phaseId": norm(str(p)), "visualIntent": f"existing IK phase '{p}'; full body visible; contacts visible"} for p in phases]
    return DEFAULT_KEYPOSES_BY_FAMILY


def build_plan(index: Dict[str, Any], limit: int | None = None, only: str | None = None) -> Dict[str, Any]:
    exercises = index.get("exercises", [])
    if only:
        needle = norm(only)
        exercises = [e for e in exercises if needle in norm(str(e.get("exerciseKey", "") + " " + e.get("displayName", "")))]
    if limit:
        exercises = exercises[:limit]
    plan = []
    for ex in exercises:
        display = ex.get("displayName") or ex.get("exerciseKey")
        key = ex.get("exerciseKey") or norm(display)
        keyposes = choose_keyposes(ex)
        base_q = display_to_query(str(display))
        queries = [
            f"{base_q} exercise start bottom finish side view",
            f"{base_q} proper form sequence images",
            f"{base_q} exercise key positions",
        ]
        if "push" in norm(str(display)):
            queries.append("push up exercise side view top bottom position")
        plan.append({
            "exerciseKey": key,
            "displayName": display,
            "sourceQualityTarget": "A/B candidate images; not IK truth",
            "sourceFolder": f"tools/keypose_image_sampler/source_images/{key}",
            "recommendedQueries": queries,
            "keyposes": keyposes,
            "ikProvenance": ex.get("sources", []),
            "status": "needs_candidate_images"
        })
    return {
        "schema": "next11.keypose_image.source_plan.v1",
        "mode": "candidate_sources_only_no_runtime_patch",
        "inputSummary": index.get("summary", {}),
        "exerciseCount": len(plan),
        "exercises": plan,
    }


def main() -> None:
    ap = argparse.ArgumentParser()
    ap.add_argument("--exercise-index", required=True)
    ap.add_argument("--out", required=True)
    ap.add_argument("--limit", type=int)
    ap.add_argument("--only", help="filter by exercise key/display, e.g. pushup")
    args = ap.parse_args()
    index_path = Path(args.exercise_index)
    index = json.loads(index_path.read_text(encoding="utf-8"))
    plan = build_plan(index, limit=args.limit, only=args.only)
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    out.write_text(json.dumps(plan, indent=2, ensure_ascii=False), encoding="utf-8")
    print(json.dumps({"exerciseCount": plan["exerciseCount"], "out": str(out)}, indent=2))


if __name__ == "__main__":
    main()