#!/usr/bin/env python3
"""Styled burn-in captions: rounded cards, custom font, per-word emphasis.

`render_reel.py --srt` burns plain libass text. libass BorderStyle=3 draws a
hard rectangle with no corner radius, so the rounded-box look modern IG/TikTok
captions use cannot be produced that way. This renders each phrase as an image
card and overlays the cards on the video.

Two stages, deliberately. Whisper mishears proper nouns and numbers ("70.3" ->
"17.3", "borrowed" -> "Boren"), so the text is surfaced for correction BEFORE
anything is burned in:

  1. review   words.json  -> phrases.json   (edit text, timings, emphasis)
  2. burn     phrases.json + video -> captioned video

Emphasis is explicit, not guessed: set "emphasis": [i, ...] on a phrase to
italicise the i-th word of that phrase.

Input is WORD-LEVEL timestamps, not a segment SRT — phrase chunking needs word
boundaries. Produce them with:
  mlx_whisper take.wav --word-timestamps True --output-format json

Usage:
  gen_captions.py review words.json --out work/captions/phrases.json
  gen_captions.py burn work/captions/phrases.json --video work/master.mp4 \\
      --out work/master_captioned.mp4
"""
import argparse, json, re, shutil, subprocess, sys
from pathlib import Path

MAX_WORDS = 4          # phrase length cap; 2-4 words is the readable range
MAX_GAP = 0.6          # a pause longer than this starts a new phrase
BOX_RADIUS = 18
PAD_X, PAD_Y = 34, 18
SPACE_RATIO = 0.28     # inter-word gap as a fraction of point size


def need(exe, why):
    path = shutil.which(exe)
    if not path:
        sys.exit(f"{exe} not found — {why}")
    return path


def run(cmd, label):
    r = subprocess.run(cmd, capture_output=True, text=True)
    if r.returncode != 0:
        sys.exit(f"{label} failed:\n{' '.join(str(c) for c in cmd)[:300]}\n"
                 f"{r.stderr[-700:]}")
    return r


# ---------- stage 1: words -> phrases ----------

def load_words(path):
    """Accept whisper's nested JSON or a flat [{word,start,end}] list."""
    data = json.loads(Path(path).read_text())
    if isinstance(data, list):
        words = data
    elif "words" in data:
        words = data["words"]
    elif "segments" in data:
        words = [w for s in data["segments"] for w in s.get("words", [])]
    else:
        sys.exit("Unrecognised words JSON. Expected a flat list, {'words': …}, "
                 "or whisper's {'segments': [{'words': …}]}.")
    out = []
    for w in words:
        text = (w.get("word") or w.get("text") or "").strip()
        if not text:
            continue
        try:
            out.append({"text": text, "start": float(w["start"]),
                        "end": float(w["end"])})
        except (KeyError, TypeError, ValueError):
            sys.exit(f"Word entry missing usable start/end: {w}. Transcribe "
                     f"with --word-timestamps True.")
    if not out:
        sys.exit("No words found — was this transcribed with word timestamps?")
    return out


def chunk(words, max_words, max_gap):
    phrases, cur = [], []
    for w in words:
        if cur:
            gap = w["start"] - cur[-1]["end"]
            ends_clause = re.search(r"[.,!?;:]$", cur[-1]["text"])
            if len(cur) >= max_words or gap > max_gap or ends_clause:
                phrases.append(cur)
                cur = []
        cur.append(w)
    if cur:
        phrases.append(cur)
    return [{"text": " ".join(w["text"] for w in p),
             "words": [w["text"] for w in p],
             "start": round(p[0]["start"], 3),
             "end": round(p[-1]["end"], 3),
             "emphasis": []} for p in phrases]


# ---------- stage 2: phrases -> cards -> burned video ----------

def render_card(phrase, opts, dst):
    """One card per phrase.

    Built from per-word `label:` renders appended with transparent spacers,
    because ImageMagick's pango delegate errors with "no decode delegate" in
    the common Homebrew IM7 build even though `-list format` advertises it.
    That rules out inline markup, so emphasis needs a separate font file.
    """
    magick = need("magick", "install ImageMagick (brew install imagemagick)")
    s = opts.get("_scale", 1.0)
    pad_x, pad_y = round(PAD_X * s), round(PAD_Y * s)
    radius = round(BOX_RADIUS * s)
    parts, tmps = [], []
    spacer = dst.parent / f".sp_{dst.stem}.png"
    run([magick, "-size", f"{int(opts['size'] * SPACE_RATIO)}x1",
         "xc:none", str(spacer)], "spacer")
    tmps.append(spacer)

    for i, word in enumerate(phrase["words"]):
        font = opts["font_italic"] if i in phrase.get("emphasis", []) else opts["font"]
        colour = opts["accent"] if i in phrase.get("emphasis", []) else opts["fg"]
        wp = dst.parent / f".w{i}_{dst.stem}.png"
        run([magick, "-background", "none", "-fill", colour, "-font", font,
             "-pointsize", str(opts["size"]), f"label:{word}", str(wp)],
            f"word {word!r}")
        tmps.append(wp)
        if parts:
            parts.append(str(spacer))
        parts.append(str(wp))

    text_img = dst.parent / f".txt_{dst.stem}.png"
    run([magick] + parts + ["-background", "none", "-gravity", "center",
                            "+append", str(text_img)], "append words")
    tmps.append(text_img)

    r = run([magick, str(text_img), "-format", "%w %h", "info:"], "measure")
    tw, th = (int(n) for n in r.stdout.split())
    bw, bh = tw + pad_x * 2, th + pad_y * 2
    run([magick, "-size", f"{bw}x{bh}", "xc:none", "-fill", opts["box"],
         "-draw", f"roundrectangle 0,0 {bw - 1},{bh - 1} {radius},{radius}",
         str(text_img), "-gravity", "center", "-composite", str(dst)],
        "compose card")
    for t in tmps:
        t.unlink(missing_ok=True)
    return dst


def video_width(path):
    r = subprocess.run(["ffprobe", "-v", "error", "-select_streams", "v:0",
                        "-show_entries", "stream=width", "-of", "csv=p=0",
                        str(path)], capture_output=True, text=True)
    if r.returncode != 0 or not r.stdout.strip():
        sys.exit(f"Cannot read width from {path}")
    return int(r.stdout.strip().split(",")[0])


def burn(phrases, video, out, opts):
    need("ffmpeg", "install ffmpeg (brew install ffmpeg)")
    # Card geometry is authored at 1080 wide and scaled to the real frame, so
    # a 4K master gets proportionally sized captions instead of half-size ones.
    scale = video_width(video) / 1080.0
    opts = dict(opts, size=max(8, round(opts["size"] * scale)))
    opts["_scale"] = scale
    print(f"  frame is {round(1080 * scale)}px wide -> caption scale "
          f"{scale:.2f}x (point size {opts['size']})", flush=True)
    cards_dir = out.parent / f"{out.stem}_cards"
    cards_dir.mkdir(parents=True, exist_ok=True)

    cards = []
    for n, p in enumerate(phrases):
        dst = cards_dir / f"card_{n:04d}.png"
        render_card(p, opts, dst)
        cards.append((dst, p))
        print(f"  card {n + 1}/{len(phrases)}: {p['text'][:48]}", flush=True)

    # Build one filtergraph. Every card is an input; overlay is gated on
    # enable=between(t,start,end) so only the current phrase is visible.
    cmd = ["ffmpeg", "-y", "-v", "error", "-i", str(video)]
    for dst, _ in cards:
        cmd += ["-i", str(dst)]
    chain, prev = [], "[0:v]"
    for i, (_, p) in enumerate(cards, start=1):
        label = f"[v{i}]"
        chain.append(
            f"{prev}[{i}:v]overlay=x=(W-w)/2:y=H*{opts['position']}-h/2:"
            f"enable='between(t,{p['start']:.3f},{p['end']:.3f})'{label}")
        prev = label
    cmd += ["-filter_complex", ";".join(chain), "-map", prev,
            "-map", "0:a?", "-c:a", "copy",
            "-c:v", "libx264", "-preset", "slow", "-crf", str(opts["crf"]),
            "-pix_fmt", "yuv420p", "-movflags", "+faststart", str(out)]
    run(cmd, "burn captions")
    return cards_dir


def main():
    ap = argparse.ArgumentParser()
    sub = ap.add_subparsers(dest="cmd", required=True)

    r = sub.add_parser("review", help="words.json -> editable phrases.json")
    r.add_argument("words_json")
    r.add_argument("--out", default="work/captions/phrases.json")
    r.add_argument("--max-words", type=int, default=MAX_WORDS)
    r.add_argument("--max-gap", type=float, default=MAX_GAP)

    b = sub.add_parser("burn", help="phrases.json + video -> captioned video")
    b.add_argument("phrases_json")
    b.add_argument("--video", required=True)
    b.add_argument("--out", required=True)
    b.add_argument("--font", default="Helvetica-Bold")
    b.add_argument("--font-italic", default="Helvetica-BoldOblique",
                   help="Used for emphasised words; pango markup is not "
                        "available, so emphasis needs its own font file/name")
    b.add_argument("--size", type=int, default=54,
                   help="Point size authored at 1080 wide; scaled to the "
                        "actual frame, so 4K gets proportional captions")
    b.add_argument("--fg", default="white")
    b.add_argument("--accent", default="#F26B21")
    b.add_argument("--box", default="#0F1D2BE6", help="Card fill, RGBA hex")
    b.add_argument("--crf", type=int, default=16,
                   help="x264 quality for the re-encode, lower is better "
                        "(default 16). Match the master you are burning onto "
                        "— burning at a higher CRF silently discards its "
                        "quality.")
    b.add_argument("--position", type=float, default=0.72,
                   help="Card centre as a fraction of frame height. Default "
                        "0.72 is lower-third — deliberately off the face.")
    args = ap.parse_args()

    if args.cmd == "review":
        phrases = chunk(load_words(args.words_json), args.max_words, args.max_gap)
        out = Path(args.out)
        out.parent.mkdir(parents=True, exist_ok=True)
        out.write_text(json.dumps(phrases, indent=2))
        print(f"{len(phrases)} phrases -> {out}\n")
        for p in phrases[:8]:
            print(f"  {p['start']:6.2f}  {p['text']}")
        if len(phrases) > 8:
            print(f"  … {len(phrases) - 8} more")
        print("\nREAD THIS BEFORE BURNING. Whisper mishears proper nouns and "
              "numbers — fix the text now, and set \"emphasis\": [i] on a "
              "phrase to italicise its i-th word.", file=sys.stderr)
        return

    phrases = json.loads(Path(args.phrases_json).read_text())
    if not phrases:
        sys.exit("No phrases in that file — nothing to burn.")
    out = Path(args.out)
    out.parent.mkdir(parents=True, exist_ok=True)
    opts = {"font": args.font, "font_italic": args.font_italic,
            "size": args.size, "fg": args.fg, "accent": args.accent,
            "box": args.box, "position": args.position, "crf": args.crf}
    cards_dir = burn(phrases, args.video, out, opts)
    print(json.dumps({"phrases": len(phrases), "cards": str(cards_dir),
                      "output": str(out), "position": args.position,
                      "font": args.font, "font_italic": args.font_italic},
                     indent=2))


if __name__ == "__main__":
    main()