#!/usr/bin/env python3
"""breakdown.py — a reference video → a MEASURED per-shot choreography breakdown.

Cuts the video into shots (its own scene detection, or a cut list you already have),
sends each shot's own footage to a multimodal model, and writes one record per shot:

  movesSentence   ONE present-tense sentence (≤25 words, no seconds, no camera verbs) —
                  what the body does. This is what a SUNG window gets: the audio
                  reference is the only clock there, and a timecoded block beside it is a
                  second clock (the rap lane's rhythm defects of 2026-08-29).
  blocks          "[a-bs: …]" timecoded blocks — for a window with NO voice, where a
                  clock in the prose is fine.
  framing / cameraMove / subjectCount / subjectVisible — what the shot IS, so a consumer
                  can pick the shot whose body is in frame.

Grounded on the frames, never on memory: the model sees the clip, and every record
carries the measured shot span from the cut list. Names, titles and "60fps / 4K /
hyper-real" lines are stripped (the render lane's copyright gate and its own look), and
the free-lane moderation nouns are flagged for the operator to reword.

Usage:
  breakdown.py --video V --out build/dance-breakdown.json
               [--cuts cuts.log | --detect] [--start S] [--end E] [--html …]
               [--model gemini-3.5-flash] [--max-shot 10] [--only 3,7] [--dry-run]

Exit 0 written · 2 bad input · 3 no API key · 4 the model failed on every shot.
"""
from __future__ import annotations

import argparse
import base64
import hashlib
import html
import json
import os
import re
import subprocess
import sys
import tempfile
import time
import urllib.error
import urllib.request

HERE = os.path.dirname(os.path.abspath(__file__))
REPO = os.path.abspath(os.path.join(HERE, "..", "..", ".."))
DEFAULT_MODEL = "gemini-3.5-flash"
ENDPOINT = "https://generativelanguage.googleapis.com/v1beta/models/{model}:generateContent"
FFMPEG = os.environ.get("FFMPEG_FULL") or os.environ.get("FFMPEG_BIN") or "ffmpeg"
SCENE_THRESHOLD = 0.15         # the same filter the rap lane's cut lists were made with
MIN_SHOT_S = 2.0
COOLDOWN_S = 60.0

# What the render lane will not accept in a prompt, or does not want. Replaced or flagged,
# never silently trusted.
TECH_NOISE = re.compile(r"\b(60\s?fps|4k|8k|hyper[- ]?real\w*|photo[- ]?real\w*|ultra[- ]?realistic|"
                        r"cinematic\s+lighting|zero\s+(limb\s+)?(distortion|jitter))\b", re.I)
# A MOVE of the camera is a second direction in a moves sentence; the dancer's orientation
# to it ("back to the camera", "facing the camera") is choreography and stays.
CAMERA_VERBS = re.compile(r"\bthe\s+camera\s+(pans|zooms|tracks|moves|pushes|pulls|tilts|orbits|dollies|cranes|follows)|"
                          r"\bcamera\s+(pan|zoom|move|movement|track|orbit)|\bdolly|\btracking\s+shot|\bcrane\s+(up|down)|"
                          r"\bzoom(s|ing)?\s+(in|out)|\bpan(s|ning)\s+(left|right|across)", re.I)
SECONDS_IN_PROSE = re.compile(r"\b\d+(\.\d+)?\s?s(ec(ond)?s?)?\b|\[\s*\d", re.I)


def say(msg: str) -> None:
    print(f"breakdown: {msg}", file=sys.stderr)


def sha256_file(path: str) -> str:
    h = hashlib.sha256()
    with open(path, "rb") as f:
        for chunk in iter(lambda: f.read(1 << 20), b""):
            h.update(chunk)
    return h.hexdigest()


def load_moderation_nouns() -> list[str]:
    p = os.path.join(REPO, "skills", "rap-avatar-mv", "references", "moderation-nouns.json")
    try:
        return list(json.load(open(p, encoding="utf-8")).get("nouns", []))
    except (OSError, ValueError):
        return []


def api_keys() -> list[str]:
    """GEMINI_API_KEYS → GOOGLE_API_KEYS → GOOGLE_API_KEY, from the environment, else from the
    repo's .env.local / .env (the runner exports these; a bare shell may not)."""
    def split(v: str) -> list[str]:
        return [k for k in re.split(r"[,;\s]+", v.strip()) if k]
    for name in ("GEMINI_API_KEYS", "GOOGLE_API_KEYS", "GOOGLE_API_KEY"):
        if os.environ.get(name):
            return split(os.environ[name])
    # A worktree has no .env.local of its own; the main checkout three levels up does.
    roots = [REPO]
    if "/.claude/worktrees/" in REPO:
        roots.append(os.path.abspath(os.path.join(REPO, "..", "..", "..")))
    for p in [os.path.join(r, fn) for r in roots for fn in (".env.local", ".env")]:
        if not os.path.isfile(p):
            continue
        for line in open(p, encoding="utf-8", errors="ignore"):
            m = re.match(r'^\s*(GEMINI_API_KEYS|GOOGLE_API_KEYS|GOOGLE_API_KEY)\s*=\s*"?([^"\n]+)"?', line)
            if m:
                return split(m.group(2))
    return []


# ---------------------------------------------------------------------------
# shots
# ---------------------------------------------------------------------------
def probe_seconds(path: str) -> float:
    r = subprocess.run(["ffprobe", "-v", "error", "-show_entries", "format=duration",
                        "-of", "csv=p=0", path], capture_output=True, text=True)
    try:
        return float(r.stdout.strip())
    except ValueError:
        return 0.0


def parse_cuts(path: str) -> list[float]:
    return sorted(float(m) for m in re.findall(r"pts_time:([0-9.]+)", open(path, encoding="utf-8").read()))


def detect_cuts(video: str, log_out: str) -> list[float]:
    subprocess.run([FFMPEG, "-hide_banner", "-v", "error", "-i", video,
                    "-vf", f"select='gt(scene,{SCENE_THRESHOLD})',metadata=print:file={log_out}",
                    "-an", "-f", "null", "-"], check=False, capture_output=True)
    return parse_cuts(log_out) if os.path.isfile(log_out) else []


def shots_from_cuts(cuts: list[float], start: float, end: float, max_shot: float) -> list[tuple[float, float]]:
    """[(a, b)] inside [start, end], no shot shorter than MIN_SHOT_S, none longer than
    max_shot (a long shot is chunked — a 17 s hold is several beats of movement)."""
    edges = [start] + [c for c in cuts if start < c < end] + [end]
    out: list[tuple[float, float]] = []
    a = edges[0]
    for b in edges[1:]:
        # A run of flash cuts (five cuts in 4.5 s of spins) is ONE beat of movement, not
        # five shots: keep accumulating until the span is at least MIN_SHOT_S.
        if b - a < MIN_SHOT_S and b != edges[-1]:
            continue
        while b - a > max_shot + 0.5:
            out.append((a, a + max_shot)); a += max_shot
        out.append((a, b)); a = b
    return [(round(x, 2), round(y, 2)) for x, y in out]


def cut_preview(video: str, a: float, b: float, out: str) -> None:
    subprocess.run([FFMPEG, "-v", "error", "-y", "-ss", f"{a:.3f}", "-t", f"{b - a:.3f}", "-i", video,
                    "-an", "-vf", "scale=640:-2,fps=12", "-c:v", "libx264", "-crf", "26",
                    "-pix_fmt", "yuv420p", "-movflags", "+faststart", out], check=True, capture_output=True)


def frames_b64(video: str, a: float, b: float, n: int = 8) -> list[str]:
    d = tempfile.mkdtemp(prefix="bd-frames-")
    subprocess.run([FFMPEG, "-v", "error", "-y", "-ss", f"{a:.3f}", "-t", f"{b - a:.3f}", "-i", video,
                    "-vf", f"fps={n / max(b - a, 0.5):.4f},scale=512:-2", "-frames:v", str(n),
                    os.path.join(d, "f-%02d.jpg")], check=True, capture_output=True)
    out = []
    for fn in sorted(os.listdir(d)):
        out.append(base64.b64encode(open(os.path.join(d, fn), "rb").read()).decode())
    return out


def strip_jpeg(video: str, a: float, b: float, out: str) -> None:
    n = max(1, int(round(b - a)))
    subprocess.run([FFMPEG, "-v", "error", "-y", "-ss", f"{a:.3f}", "-t", f"{b - a:.3f}", "-i", video,
                    "-vf", f"fps=1,scale=200:-2,tile={n}x1", "-frames:v", "1", "-update", "1", out],
                   check=False, capture_output=True)


# ---------------------------------------------------------------------------
# the model
# ---------------------------------------------------------------------------
ANALYST = """You are a dance movement analyst. Watch this clip of a dance number and describe ONLY what the body in it does, measured from the frames. Never guess who the performer is, never name a film, song, actor or dancer, and never describe anything not visible.

Return compact JSON (no markdown) with exactly these keys:
- "movesSentence": ONE present-tense sentence of at most 25 words describing the lead dancer's body movement, footwork and arm/hand gestures across the clip. No timestamps, no seconds, no camera words, no names.
- "blocks": a list of 2-4 strings, each "[a-bs: ...]" with a and b as whole seconds from the start of the clip, at most 30 words each, describing footwork, arms/hands/props, torso, head and face in that span.
- "framing": one of "wide", "medium", "close" (how much of the lead dancer's body is in frame: wide = full body small, medium = waist-up or full body large, close = head and shoulders).
- "cameraMove": one of "locked", "pan", "track", "crane", "handheld", "zoom".
- "subjectCount": integer, how many dancers are clearly in frame.
- "subjectVisible": true if the lead dancer's whole body (head to feet) is in frame for most of the clip, else false.
- "handsNearFace": true if a hand touches or covers the face at any point.
- "props": a short list of hand-held objects, e.g. ["lemon"], or [].
"""


def call_model(keys: list[str], model: str, parts: list[dict], cooldown: dict[str, float],
               max_output_tokens: int = 4096) -> dict:
    # max_output_tokens is SHARED with the model's thinking on gemini-3.5-flash: an
    # 80-word answer (~1.5k tokens) plus ~2.5k of thought hit the 4096 default and the
    # JSON came back cut off mid-row ("Expecting value: line 11") on every retry
    # (word_ends.py, 2026-09-05). A caller that expects a long answer raises it.
    body = {"contents": [{"parts": parts}],
            "generationConfig": {"temperature": 0.2, "maxOutputTokens": max_output_tokens,
                                 "responseMimeType": "application/json"}}
    data = json.dumps(body).encode()
    last = None
    for attempt in range(max(4, len(keys) + 2)):
        now = time.time()
        ready = [k for k in keys if cooldown.get(k, 0) <= now]
        if not ready:
            wait = min(cooldown.values()) - now + 0.5
            say(f"every key is cooling down — waiting {wait:.0f}s"); time.sleep(max(wait, 1))
            ready = keys
        key = ready[attempt % len(ready)]
        url = ENDPOINT.format(model=model) + "?key=" + key
        req = urllib.request.Request(url, data=data, headers={"Content-Type": "application/json"})
        try:
            with urllib.request.urlopen(req, timeout=240) as r:
                d = json.load(r)
            txt = d["candidates"][0]["content"]["parts"][0]["text"]
            txt = re.sub(r"^```(json)?\s*|\s*```$", "", txt.strip())
            return json.loads(txt)
        except urllib.error.HTTPError as e:
            msg = e.read()[:400].decode(errors="replace")
            last = f"HTTP {e.code}: {msg}"
            if e.code in (429, 503) or e.code >= 500:
                m = re.search(r"retry-after[^0-9]*(\d+)", msg, re.I)
                cooldown[key] = time.time() + (float(m.group(1)) if m else COOLDOWN_S)
                continue
            raise RuntimeError(last)
        except (KeyError, IndexError, ValueError) as e:
            last = f"unparseable answer: {e}"
            continue
    raise RuntimeError(last or "no attempts")


def analyse_shot(keys: list[str], model: str, video: str, a: float, b: float, cooldown: dict,
                 tmp: str) -> tuple[dict, str]:
    """(record, how) — how = 'video' (the clip inline) or 'frames' (the fallback)."""
    clip = os.path.join(tmp, f"shot-{a:.2f}.mp4")
    cut_preview(video, a, b, clip)
    b64 = base64.b64encode(open(clip, "rb").read()).decode()
    parts = [{"inlineData": {"mimeType": "video/mp4", "data": b64}}, {"text": ANALYST}]
    try:
        return call_model(keys, model, parts, cooldown), "video"
    except RuntimeError as exc:
        if "HTTP 4" not in str(exc):
            raise
        say(f"  inline video refused ({str(exc)[:80]}) — falling back to 8 frames")
        parts = [{"inlineData": {"mimeType": "image/jpeg", "data": f}} for f in frames_b64(video, a, b)]
        parts.append({"text": ANALYST.replace("Watch this clip", "These are 8 frames, in order, of a clip")})
        return call_model(keys, model, parts, cooldown), "frames"


# ---------------------------------------------------------------------------
# hygiene
# ---------------------------------------------------------------------------
def clean_sentence(s: str, nouns: list[str], deny: list[str]) -> tuple[str, list[str], list[str]]:
    """(sentence, stripped, flags). Stripped = removed for the render lane; flags = the
    operator must reword (the moderation nouns are not synonyms for one another)."""
    stripped, flags = [], []
    s = " ".join(str(s or "").split())
    for m in TECH_NOISE.finditer(s):
        stripped.append(m.group(0))
    s = TECH_NOISE.sub("", s)
    for word in deny:
        if word and re.search(rf"\b{re.escape(word)}\b", s, re.I):
            stripped.append(word)
            s = re.sub(rf"\b{re.escape(word)}\b", "", s, flags=re.I)
    if SECONDS_IN_PROSE.search(s):
        flags.append("seconds-in-sentence")
    if CAMERA_VERBS.search(s):
        flags.append("camera-verb")
    for n in nouns:
        if re.search(rf"\b{re.escape(n)}\b", s, re.I):
            flags.append(f"moderation-noun:{n}")
    # a capitalised word after a lowercase word reads as a name (a sentence or clause
    # opener does not: "The dancer…", ", She…", "Seated on the floor, …")
    for m in re.finditer(r"(?<=[a-z]\s)[A-Z][a-z]{2,}\b", s):
        if m.group(0) not in ("Kathak", "Rajasthani", "Indian", "Bollywood", "She", "He", "They",
                              "The", "Then", "While", "With", "Her", "His", "Their", "Seated", "Standing",
                              "Kneeling", "After", "Before", "As", "And", "But"):
            flags.append(f"proper-noun:{m.group(0)}")
    s = re.sub(r"\s{2,}", " ", s).strip(" ,;")
    words = s.split()
    if len(words) > 25:
        flags.append(f"too-long:{len(words)}w")
    return s, stripped, flags


# ---------------------------------------------------------------------------
# html
# ---------------------------------------------------------------------------
def write_html(path: str, doc: dict, strips_dir: str) -> None:
    sys.path.insert(0, os.path.join(REPO, "skills", "rap-avatar-mv", "scripts"))
    try:
        from lane_page import page as shell        # noqa: E402
    except Exception:                               # noqa: BLE001
        shell = None
    rows = []
    for s in doc["shots"]:
        strip = os.path.join(strips_dir, f"shot-{s['index']:02d}.jpg")
        img = f'<img src="file://{html.escape(strip)}" alt="shot {s["index"]}">' if os.path.isfile(strip) else ""
        blocks = "<br>".join(html.escape(b) for b in s.get("blocks") or [])
        flags = " ".join(f'<span class="b bad">{html.escape(f)}</span>' for f in s.get("flags") or [])
        meta = (f'{s["framing"]} · {s["cameraMove"]} · {s["subjectCount"]} in frame · '
                f'{"body in frame" if s.get("subjectVisible") else "body NOT fully in frame"}'
                f'{" · hands at face" if s.get("handsNearFace") else ""} · via {s.get("how")}')
        rows.append(f"""<tr><td><b>{s['index']}</b><br>{s['start']:.1f}–{s['end']:.1f}s<br>{s['seconds']:.1f}s</td>
<td>{img}</td><td><p><b>{html.escape(s.get('movesSentence') or '')}</b></p><p class="muted">{blocks}</p>
<p class="muted">{html.escape(meta)}</p>{flags}</td></tr>""")
    body = f"""<p>One row per measured shot: its 1-fps strip beside what the model saw the body do. The bold
sentence is what a SUNG window gets (no seconds — the audio is the only clock); the blocks are for
windows with no voice. Red chips are things the render lane refuses or the operator must reword.</p>
<table><tr><th>shot</th><th>frames</th><th>movement</th></tr>{''.join(rows)}</table>"""
    css = ("table{width:100%;border-collapse:collapse}td,th{vertical-align:top;padding:8px;border-top:1px solid var(--line,#333)}"
           "img{max-width:100%;height:auto}.muted{opacity:.75}.b{display:inline-block;padding:2px 8px;border-radius:12px;"
           "font-size:12px;margin:2px}.b.bad{background:#4a1d1d;color:#ff9a9a}")
    if shell:
        out = shell("Dance breakdown", f"{os.path.basename(doc['video'])} · {len(doc['shots'])} shots · {doc['model']}",
                    body, extra_css=css, eyebrow="dance-breakdown")
    else:
        out = f"<!doctype html><meta charset=utf-8><title>Dance breakdown</title><style>{css}</style>{body}"
    open(path, "w", encoding="utf-8").write(out)


# ---------------------------------------------------------------------------
def main() -> None:
    ap = argparse.ArgumentParser(description=__doc__, formatter_class=argparse.RawDescriptionHelpFormatter)
    ap.add_argument("--video", required=True)
    ap.add_argument("--out", required=True, help="dance-breakdown.json")
    ap.add_argument("--cuts", default="", help="an ffmpeg scene-detect log (pts_time lines)")
    ap.add_argument("--detect", action="store_true", help="run scene detection now (writes <out>.cuts.log)")
    ap.add_argument("--start", type=float, default=0.0)
    ap.add_argument("--end", type=float, default=0.0, help="default: the whole video")
    ap.add_argument("--max-shot", type=float, default=10.0, help="chunk longer shots (s)")
    ap.add_argument("--html", default="")
    ap.add_argument("--model", default=os.environ.get("DANCE_BREAKDOWN_MODEL", DEFAULT_MODEL))
    ap.add_argument("--deny", default="", help="comma-separated words to strip (titles, names)")
    ap.add_argument("--only", default="", help="comma-separated shot indexes to (re)run")
    ap.add_argument("--dry-run", action="store_true", help="print the shots and one request shape; no calls")
    args = ap.parse_args()

    if not os.path.isfile(args.video):
        say(f"no video at {args.video}"); raise SystemExit(2)
    end = args.end or probe_seconds(args.video)
    if args.cuts:
        cuts = parse_cuts(args.cuts); cuts_src = args.cuts
    else:
        cuts_src = args.out + ".cuts.log"
        if not os.path.isfile(cuts_src) or args.detect:
            say("detecting cuts…"); cuts = detect_cuts(args.video, cuts_src)
        else:
            cuts = parse_cuts(cuts_src)
    shots = shots_from_cuts(cuts, args.start, end, args.max_shot)
    say(f"{len(shots)} shots in {args.start:.0f}-{end:.0f}s (cuts: {os.path.basename(cuts_src)})")

    prior = {}
    if os.path.isfile(args.out):
        try:
            prior = {s["index"]: s for s in json.load(open(args.out, encoding="utf-8")).get("shots", [])}
        except (OSError, ValueError):
            prior = {}
    only = {int(x) for x in args.only.split(",") if x.strip()} if args.only else None

    if args.dry_run:
        for i, (a, b) in enumerate(shots):
            print(f"  shot {i:02d}  {a:6.2f}-{b:6.2f}  {b - a:5.2f}s")
        print(json.dumps({"model": args.model, "parts": [{"inlineData": {"mimeType": "video/mp4", "data": "<base64 of the shot's 640px muted preview>"}},
                                                          {"text": ANALYST[:160] + "…"}]}, indent=1))
        return

    keys = api_keys()
    if not keys:
        say("no Gemini key (GEMINI_API_KEYS / GOOGLE_API_KEY, env or the repo .env.local)"); raise SystemExit(3)
    nouns = load_moderation_nouns()
    deny = [d.strip() for d in args.deny.split(",") if d.strip()]
    strips_dir = os.path.splitext(args.out)[0] + "-strips"
    os.makedirs(strips_dir, exist_ok=True)
    tmp = tempfile.mkdtemp(prefix="dance-breakdown-")
    cooldown: dict[str, float] = {}
    records, failed = [], 0
    for i, (a, b) in enumerate(shots):
        if only is not None and i not in only and i in prior:
            records.append(prior[i]); continue
        if only is not None and i not in only:
            continue
        strip_jpeg(args.video, a, b, os.path.join(strips_dir, f"shot-{i:02d}.jpg"))
        try:
            raw, how = analyse_shot(keys, args.model, args.video, a, b, cooldown, tmp)
        except (RuntimeError, subprocess.CalledProcessError) as exc:
            say(f"  shot {i:02d} {a:.1f}-{b:.1f}s FAILED: {str(exc)[:160]}"); failed += 1
            records.append({"index": i, "start": a, "end": b, "seconds": round(b - a, 2), "error": str(exc)[:300]})
            continue
        sentence, stripped, flags = clean_sentence(raw.get("movesSentence", ""), nouns, deny)
        blocks = []
        for blk in raw.get("blocks") or []:
            t, s2, f2 = clean_sentence(blk, nouns, deny)
            blocks.append(t); stripped += s2; flags += [f for f in f2 if not f.startswith("seconds")]
        rec = {"index": i, "start": a, "end": b, "seconds": round(b - a, 2), "how": how,
               "movesSentence": sentence, "blocks": blocks,
               "framing": raw.get("framing"), "cameraMove": raw.get("cameraMove"),
               "subjectCount": raw.get("subjectCount"), "subjectVisible": bool(raw.get("subjectVisible")),
               "handsNearFace": bool(raw.get("handsNearFace")), "props": raw.get("props") or [],
               "stripped": sorted(set(stripped)), "flags": sorted(set(flags)), "raw": raw}
        records.append(rec)
        say(f"  shot {i:02d} {a:5.1f}-{b:5.1f}s {rec['framing']:6s} — {sentence[:80]}")
    records.sort(key=lambda r: r["index"])
    if records and failed == len([r for r in records if only is None or r["index"] in only]):
        raise SystemExit(4)
    doc = {"schemaVersion": 1, "video": os.path.abspath(args.video), "videoSha": sha256_file(args.video),
           "cuts": os.path.abspath(cuts_src), "cutsSha": sha256_file(cuts_src) if os.path.isfile(cuts_src) else None,
           "start": args.start, "end": end, "model": args.model, "sceneThreshold": SCENE_THRESHOLD,
           "shots": records}
    json.dump(doc, open(args.out, "w", encoding="utf-8"), indent=1, ensure_ascii=False)
    if args.html:
        write_html(args.html, doc, strips_dir)
    say(f"wrote {args.out} ({len(records)} shots, {failed} failed)" + (f" and {args.html}" if args.html else ""))


if __name__ == "__main__":
    main()