#!/usr/bin/env python3
"""STORYBOARD v0.2: grade 25s × 3 frames (BASE-CAM) → votos VLM → Viterbi de fases → storyboard.json.

Schema de saída INALTERADO (unified consome tc_start/tc_end/beat/importance/emotion).

Escape: STORYBOARD_VITERBI=0 usa caminho legado (1 frame, cam ativa, sem Viterbi).
Default: Viterbi ON.
"""
import json, os, glob
from collections import Counter

from vlm import ask_image, par_map
import judgelib as J
import cache_manifest as cm
import storyboard_viterbi as SV

WIN = 25 * J.FPS
OUT = "storyboard.json"
_SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
USE_VITERBI = os.environ.get("STORYBOARD_VITERBI", "1").strip() not in ("0", "false", "no")


def tcff(f):
    return f"{f//108000:02d}:{(f//1800)%60:02d}:{(f//30)%60:02d}:{f%30:02d}"


def core0_two_cams(placements):
    """E10: primeiro frame com ≥2 cams cobrindo (não starts[3])."""
    events = []
    for c in placements:
        a, b = int(c["startFrame"]), int(c["startFrame"]) + int(c["durationFrames"])
        events.append((a, +1))
        events.append((b, -1))
    events.sort()
    cov = 0
    for f, d in events:
        cov += d
        if cov >= 2:
            return f
    starts = sorted(int(c["startFrame"]) for c in placements)
    return starts[max(0, min(len(starts) - 1, 3))] if starts else 0


def load_base_cam_angle():
    """BASE-CAM = cam com menor mediana de gmot (tripé), como unified.py. Devolve angle 1..N ou None."""
    try:
        import numpy as np
        import vlib as V
    except Exception:
        return None
    GMOT = {}
    for c in J.P:
        # v0.4: footage_key(name, cam=) — não [:-4]
        n = V.footage_key(c["name"], cam=c.get("cam"))
        path = f"shakecache/{n}.json"
        if not os.path.exists(path):
            key = n
            hit = None
            for fp in glob.glob("shakecache/*.json"):
                if key and key in os.path.basename(fp):
                    hit = fp
                    break
            if not hit:
                continue
            path = hit
        try:
            sc = json.load(open(path))
        except Exception:
            continue
        cam = c["cam"]
        gm = [float(fr.get("gmot", 0) or 0) for fr in sc.get("frames", [])]
        if gm:
            GMOT.setdefault(cam, []).extend(gm)
    if not GMOT:
        return None
    base_cam = min(GMOT, key=lambda c: float(np.median(GMOT[c])))
    print(f"[storyboard] BASE_CAM={base_cam} (min median gmot)")
    return V.CAM_ANGLE.get(base_cam)


def sample_angle_fn(base_angle, active_fn, cov_fn=None):
    """BASE-CAM quando cobre; senão ativa. angle int."""
    import vlib as V

    def ang_at(f):
        if base_angle is not None:
            try:
                if V.covers(V.coverage(), base_angle, f, f + 1):
                    return base_angle
            except Exception:
                pass
        return active_fn(f)

    return ang_at


def classify_frame(img, prompt):
    if not img:
        return "officiant", "low", ""
    txt, _ = ask_image(img, prompt)
    if isinstance(txt, str) and txt.startswith("ERR:"):
        return "officiant", "low", ""
    v = J.jparse(txt)
    beat = SV.clamp_beat(v.get("beat") or "officiant")
    emo = (v.get("emotion") or "low").lower()
    if emo not in ("low", "med", "high"):
        emo = "low"
    return beat, emo, v.get("subject", "")


def main():
    # E5: cache SEM UNI/N_director (base-cam + placements + prompt + scripts)
    _sig = {
        "mtime": cm.input_signature([
            "placements.json",
            os.path.join(_SCRIPT_DIR, "prompts", "storyboard_v1.md"),
            os.path.join(_SCRIPT_DIR, "judgelib.py"),
            os.path.join(_SCRIPT_DIR, "storyboard_viterbi.py"),
            os.path.join(_SCRIPT_DIR, "storyboard_gen.py"),
        ]),
        "viterbi": USE_VITERBI,
    }
    # shakecache mtimes (base-cam dep)
    shake_files = sorted(glob.glob("shakecache/*.json"))
    if shake_files:
        _sig["shake"] = cm.input_signature(shake_files)

    if cm.is_fresh(".", "storyboard_gen", __file__, _sig) and os.path.exists(OUT):
        print(f"[cache] storyboard_gen: script+input inalterados, pulando VLM — reusando {OUT}")
        return

    active = J.active_fn()
    cov0 = min(int(c["startFrame"]) for c in J.P)
    cov1 = max(int(c["startFrame"]) + int(c["durationFrames"]) for c in J.P)
    core0 = core0_two_cams(J.P)
    prompt = J.prompt("storyboard_v1.md")

    if not USE_VITERBI:
        print("[storyboard] STORYBOARD_VITERBI=0 — caminho legado (1 frame, cam ativa)")
        _legacy(active, cov0, cov1, core0, prompt, _sig)
        return

    base_angle = load_base_cam_angle()
    print(f"[storyboard] v0.2 Viterbi ON  base_angle={base_angle}  core0={core0} ({J.tc(core0)})")
    ang_at = sample_angle_fn(base_angle, active)

    # grade de janelas
    f0 = core0 - (core0 % WIN)
    windows = []
    f = f0
    while f < cov1:
        windows.append(f)
        f += WIN

    # 3 samples por janela (25/50/75%)
    jobs = []  # (wi, si, frame)
    for wi, wf in enumerate(windows):
        for si, frac in enumerate((0.25, 0.50, 0.75)):
            fr = int(wf + frac * WIN)
            if fr >= cov1:
                fr = cov1 - 1
            jobs.append((wi, si, fr))

    def work(job):
        wi, si, fr = job
        a = ang_at(fr)
        img = J.frame_of(a, fr)
        beat, emo, subj = classify_frame(img, prompt)
        return wi, si, beat, emo, subj, fr, a

    print(f"[storyboard] classificando {len(jobs)} frames ({len(windows)} janelas × 3)…")
    results = par_map(work, jobs)
    # agrupa por janela
    votes = [[] for _ in windows]
    emos = [[] for _ in windows]
    subjects = [[] for _ in windows]
    for wi, si, beat, emo, subj, fr, a in results:
        votes[wi].append(beat)
        emos[wi].append(emo)
        subjects[wi].append(subj)

    decoded = SV.decode_windows(votes, emos)

    rows = []
    for wi, wf in enumerate(windows):
        d = decoded[wi] if wi < len(decoded) else {
            "beat": "officiant", "importance": 0.5, "emotion": "low", "phase": "liturgia"
        }
        subj = ""
        if subjects[wi]:
            # subject do voto que bate com beat final, senão o do meio
            subj = subjects[wi][min(1, len(subjects[wi]) - 1)] or ""
        rows.append({
            "f0": wf, "f1": wf + WIN,
            "beat": d["beat"],
            "subject": subj,
            "importance": float(d["importance"]),
            "emotion": d["emotion"],
            "phase": d.get("phase", ""),
            "votes": d.get("votes", votes[wi]),
        })
        if (wi + 1) % 20 == 0:
            print(f"  {wi+1} janelas... ({J.tc(wf)} {d['beat']} phase={d.get('phase')})")

    # refinamento de fronteira (+2 frames) onde a FASE muda
    rows = _refine_boundaries(rows, ang_at, prompt)

    out = SV.merge_same_beat(rows)
    if cov0 < core0:
        out.insert(0, {
            "f0": cov0, "f1": core0, "beat": "pre", "subject": "pre-ceremony",
            "importance": 0.2, "emotion": "low",
        })

    sb = [{
        "tc_start": tcff(r["f0"]), "tc_end": tcff(r["f1"]),
        "beat": r["beat"], "subject": r.get("subject", ""),
        "importance": round(float(r["importance"]), 2),
        "emotion": r.get("emotion", "low"), "note": "",
    } for r in out]

    json.dump(sb, open(OUT, "w"), indent=1)
    cm.mark_done(".", "storyboard_gen", __file__, _sig)

    # eval estrutural no próprio run
    metrics = SV.eval_sequence(out, cov0, cov1)
    print(f"storyboard.json: {len(sb)} batidas de {len(windows)} janelas — "
          f"{dict(Counter(b['beat'] for b in sb))}")
    print(f"eval: {json.dumps(metrics)}")


def _refine_boundaries(rows, ang_at, prompt):
    """Onde phase muda entre janelas, reclassifica 1/3 e 2/3 e ajusta f1/f0 (~8s)."""
    if len(rows) < 2:
        return rows
    out = [dict(r) for r in rows]
    for i in range(len(out) - 1):
        if out[i].get("phase") == out[i + 1].get("phase"):
            continue
        # só se phase presente (viterbi path)
        if not out[i].get("phase") or not out[i + 1].get("phase"):
            continue
        boundary = out[i]["f1"]
        f_a = boundary - WIN // 3
        f_b = boundary - WIN // 6
        for fr in (f_a, f_b):
            if fr <= out[i]["f0"] or fr >= out[i + 1]["f1"]:
                continue
            img = J.frame_of(ang_at(fr), fr)
            beat, emo, _ = classify_frame(img, prompt)
            # se o voto do frame de fronteira casa com a fase da janela SEGUINTE, move fronteira p/ trás
            ph_next = out[i + 1].get("phase")
            ph_prev = out[i].get("phase")
            if SV.BEAT_TO_PHASE.get(beat) == ph_next and beat != "empty_wait":
                out[i]["f1"] = fr
                out[i + 1]["f0"] = fr
                break
            if SV.BEAT_TO_PHASE.get(beat) == ph_prev:
                # fronteira fica / move um pouco à frente
                out[i]["f1"] = fr + (WIN // 12)
                out[i + 1]["f0"] = out[i]["f1"]
    return out


def _legacy(active, cov0, cov1, core0, prompt, _sig):
    rows = []
    f = core0 - (core0 % WIN)
    n = 0
    while f < cov1:
        img = J.frame_of(active(f + WIN // 2), f + WIN // 2)
        if img:
            beat, emo, subj = classify_frame(img, prompt)
            rows.append({
                "f0": f, "f1": f + WIN, "beat": beat, "subject": subj,
                "importance": SV.importance_for(beat, SV.BEAT_TO_PHASE.get(beat, "liturgia")),
                "emotion": emo,
            })
            n += 1
            if n % 20 == 0:
                print(f"  {n} janelas... ({J.tc(f)} {beat})")
        f += WIN
    out = SV.merge_same_beat(rows)
    if cov0 < core0:
        out.insert(0, {
            "f0": cov0, "f1": core0, "beat": "pre", "subject": "pre-ceremony",
            "importance": 0.2, "emotion": "low",
        })
    sb = [{
        "tc_start": tcff(r["f0"]), "tc_end": tcff(r["f1"]),
        "beat": r["beat"], "subject": r.get("subject", ""),
        "importance": round(float(r["importance"]), 2),
        "emotion": r.get("emotion", "low"), "note": "",
    } for r in out]
    json.dump(sb, open(OUT, "w"), indent=1)
    cm.mark_done(".", "storyboard_gen", __file__, _sig)
    print(f"storyboard.json: {len(sb)} batidas de {n} janelas — {dict(Counter(b['beat'] for b in sb))}")


if __name__ == "__main__":
    main()
