#!/usr/bin/env python3
"""assemble_answer.py — the rendered takes of an answer → one finished clip (captions, cutaways).

Proven on `ad-answer-rundown` (2026-09-06, six Veo-fast takes → a 40 s 9:16 answer to a
50 s competitor ad). Per beat: whisper primed with the scripted line → trim at the last
SCRIPTED word + TAIL (never past the take; whisper's hallucinated repeat of the last two
words is never captioned) → 1080x1920/30 fps encode. A beat with a cutaway shows the still
full-frame for CUTAWAY_S from the moment the take's words reach the beat's noun (video only;
the take's audio keeps running). Then: concat → -14 LUFS with the limiter AFTER gain →
phrase captions from the heard words (never two on screen at once) burned through
ffmpeg-full `subtitles=` (big Inter, centred low) → optional end card → QC json on the
ENCODED file, one row per beat with heard-vs-scripted words. A take that covers less than
MIN_COVERAGE of its line is REFUSED with the beat named; nothing is invented for it.

Whisper is cached by the take's sha + the line, so a re-assembly only re-encodes and a
re-rendered take invalidates its own cache.

usage: assemble_answer.py --project <p> --script <plan.json|script.json> --out final/answer.mp4
                          [--cutaway <beatId>=<image>]... [--nouns tip1=air,tip2=clips,...]
                          [--fonts <dir>] [--font-name Inter] [--no-subs]
                          [--end-card "<text>"] [--end-card-seconds 3]
Exit 0 written · 2 bad input · 4 a beat refused (the clip is still written from the beats that passed).
Takes are read from <project>/outputs/scene-<i>.mp4 in beat order (vclaw's `produce --scene i`).
"""
from __future__ import annotations

import argparse
import hashlib
import json
import os
import re
import shutil
import subprocess
from pathlib import Path

FFMPEG = (os.environ.get("VCLAW_FFMPEG")
          or ("/opt/homebrew/opt/ffmpeg-full/bin/ffmpeg" if os.path.exists("/opt/homebrew/opt/ffmpeg-full/bin/ffmpeg") else "ffmpeg"))
FFPROBE = os.environ.get("FFPROBE_BIN") or "ffprobe"
# Inter-*.ttf for the burned captions and the end card: $VCLAW_FONTS_DIR, else <workspace>/fonts when
# VCLAW_WORKSPACE is set, else none (libass then uses the system's Inter if installed; the end card
# falls back to PIL's default face). Nothing in this file assumes where the operator keeps fonts.
DEFAULT_FONTS = (os.environ.get("VCLAW_FONTS_DIR")
                 or (os.path.join(os.environ["VCLAW_WORKSPACE"], "fonts") if os.environ.get("VCLAW_WORKSPACE") else ""))
TAIL = 0.40
MIN_COVERAGE = 0.80
CUTAWAY_S = 2.0
W, H, FPS = 1080, 1920, 30


def sh(cmd, **kw):
    r = subprocess.run(cmd, capture_output=True, text=True, **kw)
    if r.returncode != 0:
        raise SystemExit(f"command failed ({r.returncode}): {' '.join(map(str, cmd))}\n{r.stderr[-800:]}")
    return r


def tok(t: str) -> list[str]:
    return [x for x in re.split(r"[^a-z0-9']+", t.lower()) if x]


def probe_seconds(p: str) -> float:
    r = subprocess.run([FFPROBE, "-v", "error", "-show_entries", "format=duration", "-of", "csv=p=0", p], capture_output=True, text=True)
    return float(r.stdout.strip() or 0)


def whisper_words(take: str, line: str, build: Path, tag: str) -> list[dict]:
    """Cached by the take's sha + the line (a re-rendered take invalidates its cache)."""
    key = hashlib.sha1(open(take, "rb").read() + line.encode()).hexdigest()[:12]
    cached = build / f"{tag}-{key}.json"
    if not cached.is_file():
        wav = build / f"{tag}.wav"
        sh([FFMPEG, "-v", "error", "-y", "-i", take, "-vn", "-ac", "1", str(wav)])
        sh(["whisper", str(wav), "--model", "small", "--language", "en", "--output_format", "json",
            "--output_dir", str(build), "--word_timestamps", "True", "--initial_prompt", line], cwd=str(build))
        (build / f"{tag}.json").rename(cached)
    d = json.load(open(cached))
    return [{"w": w["word"].strip(), "s": float(w["start"]), "e": float(w["end"])}
            for seg in d.get("segments", []) for w in seg.get("words", [])]


def align(words: list[dict], line: str) -> tuple[float, float, list[str], int]:
    """Coverage of the scripted line by the heard words, the end of the last SCRIPTED word heard,
    the scripted words that were missed, and the index of the last heard word that was scripted."""
    L = tok(line)
    heard = [(tok(w["w"]), w) for w in words]
    hit = 0; last_end = 0.0; i = 0; missed = []; last_idx = -1
    for t in L:
        j = i
        found = None
        while j < len(heard) and j < i + 6:
            if t in heard[j][0]:
                found = j; break
            j += 1
        if found is None:
            missed.append(t); continue
        hit += 1; last_end = heard[found][1]["e"]; i = found + 1; last_idx = found
    return hit / max(1, len(L)), last_end, missed, last_idx


def noun_time(words: list[dict], noun: str, fallback: float) -> float:
    for w in words:
        if noun and noun in tok(w["w"]):
            return w["s"]
    return fallback


def encode_take(src: str, dst: Path, end: float, cutaway: str | None, cut_at: float) -> None:
    vf = f"scale={W}:{H}:force_original_aspect_ratio=increase,crop={W}:{H},fps={FPS},format=yuv420p"
    if cutaway:
        # the cutaway is a still shown full-frame for CUTAWAY_S from cut_at; audio untouched
        c_end = min(end, cut_at + CUTAWAY_S)
        sh([FFMPEG, "-v", "error", "-y", "-i", src, "-loop", "1", "-i", cutaway,
            "-filter_complex",
            f"[0:v]{vf}[base];[1:v]scale={W}:{H}:force_original_aspect_ratio=decrease,pad={W}:{H}:(ow-iw)/2:(oh-ih)/2:color=#101418,fps={FPS},format=yuv420p[cut];"
            f"[base][cut]overlay=enable='between(t,{cut_at:.2f},{c_end:.2f})':eof_action=pass[v]",
            "-map", "[v]", "-map", "0:a", "-t", f"{end:.2f}", "-c:v", "libx264", "-crf", "18", "-preset", "medium",
            "-c:a", "aac", "-b:a", "160k", "-ar", "48000", str(dst)])
    else:
        sh([FFMPEG, "-v", "error", "-y", "-i", src, "-t", f"{end:.2f}", "-vf", vf,
            "-c:v", "libx264", "-crf", "18", "-preset", "medium", "-c:a", "aac", "-b:a", "160k", "-ar", "48000", str(dst)])


def end_card(text: str, seconds: float, build: Path, fonts: str) -> Path:
    """A card after the CTA: dark ground, the text in Inter, silence under it (so the concat keeps an audio stream)."""
    from PIL import Image, ImageDraw, ImageFont
    im = Image.new("RGB", (W, H), (16, 20, 24)); dr = ImageDraw.Draw(im)
    font_path = os.path.join(fonts, "Inter-ExtraBold.ttf") if fonts else ""
    font = ImageFont.truetype(font_path, 84) if os.path.isfile(font_path) else ImageFont.load_default()
    lines = text.split("\\n")
    y = H // 2 - 60 * len(lines)
    for ln in lines:
        w = dr.textlength(ln, font=font); dr.text(((W - w) / 2, y), ln, font=font, fill=(255, 255, 255)); y += 120
    png = build / "end-card.png"; im.save(png)
    out = build / "end-card.mp4"
    sh([FFMPEG, "-v", "error", "-y", "-loop", "1", "-i", str(png), "-f", "lavfi", "-i", "anullsrc=r=48000:cl=stereo", "-t", f"{seconds:.2f}",
        "-r", str(FPS), "-c:v", "libx264", "-crf", "18", "-pix_fmt", "yuv420p", "-c:a", "aac", "-b:a", "160k", "-shortest", str(out)])
    return out


def ts(t: float) -> str:
    h = int(t // 3600); m = int(t % 3600 // 60); s = t % 60
    return f"{h}:{m:02d}:{s:05.2f}"


def ass_header(font: str) -> str:
    return ("[Script Info]\nScriptType: v4.00+\nPlayResX: 1080\nPlayResY: 1920\nWrapStyle: 0\n\n[V4+ Styles]\n"
            "Format: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, Alignment, MarginL, MarginR, MarginV, Encoding\n"
            f"Style: Big,{font},78,&H00FFFFFF,&H000000FF,&H00000000,&H80000000,1,0,0,0,100,100,0,0,1,6,2,2,70,70,430,1\n\n"
            "[Events]\nFormat: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n")


def phrase_cues(words: list[dict], offset: float, end_cap: float, size: int = 4) -> list[tuple[float, float, str]]:
    """3–5 words per cue from the heard words; a cue ends where the next one starts."""
    cues = []
    for k in range(0, len(words), size):
        grp = [w for w in words[k:k + size] if w["s"] < end_cap]
        if not grp:
            break
        s = grp[0]["s"] + offset; e = min(grp[-1]["e"] + 0.15, end_cap) + offset
        cues.append((s, max(e, s + 0.3), " ".join(w["w"] for w in grp).strip()))
    for k in range(len(cues) - 1):
        s, e, t = cues[k]; n = cues[k + 1][0]
        if e > n - 0.05:
            cues[k] = (s, max(s + 0.2, n - 0.05), t)
    return cues


def lufs(path: str) -> float:
    m = subprocess.run([FFMPEG, "-i", path, "-af", "ebur128=peak=true", "-f", "null", "-"], capture_output=True, text=True).stderr
    mm = re.findall(r"I:\s+(-?[\d.]+) LUFS", m)
    return float(mm[-1]) if mm else 0.0


def parse_pairs(items: list[str] | None, what: str) -> dict[str, str]:
    out: dict[str, str] = {}
    for item in items or []:
        for part in item.split(","):
            if not part.strip():
                continue
            if "=" not in part:
                raise SystemExit(f"--{what} wants <beatId>=<value>[,<beatId>=<value>...], got {part!r}")
            k, v = part.split("=", 1); out[k.strip()] = v.strip()
    return out


def load_beats(path: str) -> list[dict]:
    d = json.load(open(path))
    beats = d.get("beats") or []
    if not beats or not all(isinstance(b, dict) and b.get("id") and b.get("line") for b in beats):
        raise SystemExit(f"{path}: `beats` must be a list of {{id, line[, cutaway]}}")
    return beats


def main() -> int:
    ap = argparse.ArgumentParser(description=__doc__.split("\n\n")[0])
    ap.add_argument("--project", required=True)
    ap.add_argument("--script", required=True, help="a plan.json from answer_plan.py, or any {beats:[{id,line}]} file")
    ap.add_argument("--out", required=True)
    ap.add_argument("--cutaway", action="append", default=[], help="<beatId>=<image path> (repeatable)")
    ap.add_argument("--nouns", action="append", default=[], help="<beatId>=<word the cutaway waits for>[,...]; default: a beat's `noun` in the plan, else 2 s in")
    ap.add_argument("--fonts", default=DEFAULT_FONTS, help="directory holding Inter-*.ttf ($VCLAW_FONTS_DIR, else $VCLAW_WORKSPACE/fonts)")
    ap.add_argument("--font-name", default="Inter")
    ap.add_argument("--no-subs", action="store_true")
    ap.add_argument("--end-card", help="text for an end card after the CTA (\\n for a line break)")
    ap.add_argument("--end-card-seconds", type=float, default=3.0)
    a = ap.parse_args()
    try:
        P = Path(a.project); build = P / "build" / "answer"; build.mkdir(parents=True, exist_ok=True)
        beats = load_beats(a.script)
        cutaways = parse_pairs(a.cutaway, "cutaway")
        nouns = {b["id"]: b["noun"] for b in beats if b.get("noun")} | parse_pairs(a.nouns, "nouns")
    except SystemExit as e:
        print(f"assemble_answer: {e}", file=__import__("sys").stderr); return 2
    for k, v in cutaways.items():
        if not os.path.isfile(v):
            print(f"assemble_answer: --cutaway {k}: {v} is not a file", file=__import__("sys").stderr); return 2

    parts = []; qc = {"schemaVersion": 1, "beats": [], "refused": []}; t_offset = 0.0; all_cues = []
    for i, b in enumerate(beats):
        take = str(P / "outputs" / f"scene-{i}.mp4")
        if not os.path.isfile(take):
            qc["refused"].append({"beat": b["id"], "why": "no take on disk"})
            qc["beats"].append({"beat": b["id"], "line": b["line"], "refused": True, "why": "no take on disk"}); continue
        words = whisper_words(take, b["line"], build, f"scene-{i}")
        cov, last_end, missed, last_idx = align(words, b["line"])
        words = words[:last_idx + 1] if last_idx >= 0 else words   # never caption a hallucinated repeat
        take_len = probe_seconds(take)
        end = min(take_len, (last_end or take_len) + TAIL)
        row = {"beat": b["id"], "line": b["line"], "coverage": round(cov, 2), "heard": " ".join(w["w"] for w in words),
               "missed": missed, "trimTo": round(end, 2), "takeSeconds": take_len}
        if cov < MIN_COVERAGE:
            row["refused"] = True; qc["refused"].append({"beat": b["id"], "coverage": round(cov, 2), "missed": missed}); qc["beats"].append(row); continue
        cut = cutaways.get(b["id"]); cut_at = noun_time(words, nouns.get(b["id"], ""), 2.0) if cut else 0.0
        if cut:
            cut_at = min(cut_at, max(0.0, end - CUTAWAY_S - 0.3))
            row["cutawayAt"] = round(cut_at, 2)
        dst = build / f"part-{i}.mp4"
        encode_take(take, dst, end, cut, cut_at)
        parts.append(dst)
        all_cues += phrase_cues(words, t_offset, end)
        t_offset += probe_seconds(str(dst))
        qc["beats"].append(row)
    if not parts:
        print(json.dumps(qc)); return 4
    if a.end_card:
        parts.append(end_card(a.end_card, a.end_card_seconds, build, a.fonts))
    lst = build / "concat.txt"; lst.write_text("".join(f"file '{p}'\n" for p in parts))
    raw = build / "concat-raw.mp4"
    sh([FFMPEG, "-v", "error", "-y", "-f", "concat", "-safe", "0", "-i", str(lst), "-c", "copy", str(raw)])
    gain = -14.0 - lufs(str(raw))
    lev = build / "concat-lev.mp4"
    sh([FFMPEG, "-v", "error", "-y", "-i", str(raw), "-af", f"volume={gain:.2f}dB,alimiter=limit=0.79:level=disabled:attack=5:release=50",
        "-c:v", "copy", "-c:a", "aac", "-b:a", "160k", str(lev)])
    out = Path(a.out); out.parent.mkdir(parents=True, exist_ok=True)
    if a.no_subs:
        shutil.copyfile(lev, out)
    else:
        ass = build / "captions.ass"
        ass.write_text(ass_header(a.font_name) + "".join(f"Dialogue: 0,{ts(s)},{ts(e)},Big,,0,0,0,,{t}\n" for s, e, t in all_cues))
        sub = f"subtitles={ass}" + (f":fontsdir={a.fonts}" if a.fonts else "")
        sh([FFMPEG, "-v", "error", "-y", "-i", str(lev), "-vf", sub, "-c:v", "libx264", "-crf", "18",
            "-pix_fmt", "yuv420p", "-c:a", "copy", str(out)])
    qc["output"] = {"path": str(out), "seconds": probe_seconds(str(out)), "lufs": round(lufs(str(out)), 2), "cues": len(all_cues),
                    "gainDb": round(gain, 2), "endCard": bool(a.end_card)}
    sh([FFMPEG, "-v", "error", "-y", "-i", str(out), "-vf", "fps=1,scale=120:-1,tile=12x5", str(build / "final-strip.jpg")])
    sh([FFMPEG, "-v", "error", "-y", "-ss", str(qc["output"]["seconds"] / 2), "-i", str(out), "-frames:v", "1", str(build / "final-mid.jpg")])
    qc_path = out.with_name(out.stem + "-qc.json")
    json.dump(qc, open(qc_path, "w"), indent=1)
    print(json.dumps(qc["output"] | {"qc": str(qc_path), "refused": qc["refused"]}))
    return 0 if not qc["refused"] else 4


if __name__ == "__main__":
    raise SystemExit(main())