#!/usr/bin/env python3
"""
overlay_hook.py - burn a short-form note's `## hook` text onto its BufferMedia clip.

Usage:
  overlay_hook.py "<note.md>"                       # writes /Users/sean/BufferMedia/<uid>.mp4
  overlay_hook.py "<note.md>" --frame 3 --png /tmp/x.png   # single-frame preview, same filter
  overlay_hook.py "<note.md>" --dry-run             # print layout only
Options: --hook TEXT (override), --duration SEC (show only for first SEC seconds),
         --videotoolbox (h264_videotoolbox instead of libx264), --force (overwrite output),
         --keep-source (do not delete the pre-overlay BufferMedia source).

Layout spec (Sean, 2026-09-29):
  * ALL CAPS Impact, white fill, black outline (borderw 7 @1080 wide, scaled), no box.
  * Font size 7.5% of frame width; shrinks (never grows) until the hook fits in <= 4 lines.
  * Safe box: 8% from left, 18% from right (box width 74%), lines centered in that box.
  * Outline top edge of the block sits at 50% of frame height; whole block must end above 70%.
  * Text is pre-wrapped by measuring real glyph advances from the font file (pure-python
    TrueType hmtx/cmap reader, no PIL needed), balanced so lines are similar lengths.
  * Shown for the whole clip unless the note's ## ai notes says "overlay ... first N seconds".
  * Audio stream copied; video libx264 -crf 18 -preset medium, same resolution + fps.
Encodes to a verified temporary file before replacing the final BufferMedia clip.
"""
import argparse, itertools, json, os, re, shutil, struct, subprocess, sys, tempfile, urllib.parse

BUFFER_DIR = "/Users/sean/BufferMedia"

FONT_CANDIDATES = [
    "/System/Library/Fonts/Supplemental/Impact.ttf",
    "/Library/Fonts/Impact.ttf",
    "~/Library/Fonts/Impact.ttf",
    # fallbacks: closest condensed heavy sans
    "~/Library/Fonts/Anton-Regular.ttf", "/Library/Fonts/Anton-Regular.ttf",
    "~/Library/Fonts/LeagueGothic-Regular.ttf", "/Library/Fonts/LeagueGothic-Regular.ttf",
    "~/Library/Fonts/Oswald-Bold.ttf", "/Library/Fonts/Oswald-Bold.ttf",
]
FFMPEG_DIRS = ["/opt/homebrew/bin", "/usr/local/bin", "/Applications/Jellyfin.app/Contents/MacOS"]

# ---- spec constants (fractions of frame) ----
TOP_FRAC = 0.50          # outline top edge of text block
BOTTOM_LIMIT_FRAC = 0.70 # block must end above this
LEFT_FRAC = 0.08
RIGHT_FRAC = 0.18
FONT_FRAC = 0.075        # nominal font size as fraction of width (81px @1080)
MIN_FONT_FRAC = 0.040    # never shrink below this
BORDER_AT_1080 = 7       # outline width at 1080 px wide
LINE_PITCH_EM = 1.15     # baseline-to-baseline distance, in font-size units
MAX_LINES = 4
WIDTH_SAFETY = 0.98      # use 98% of the safe box for measured text (kerning/rounding slack)


# ---------------------------------------------------------------- font metrics
class TTFMetrics:
    """Minimal TrueType/OpenType reader: unitsPerEm, cap height, per-char advance widths."""
    def __init__(self, path):
        with open(path, "rb") as f:
            self.d = d = f.read()
        if d[:4] == b"ttcf":
            off0 = struct.unpack(">I", d[12:16])[0]
        else:
            off0 = 0
        num = struct.unpack(">H", d[off0 + 4:off0 + 6])[0]
        self.t = {}
        for i in range(num):
            rec = off0 + 12 + 16 * i
            tag = d[rec:rec + 4].decode("latin1")
            _, off, ln = struct.unpack(">III", d[rec + 4:rec + 16])
            self.t[tag] = (off, ln)
        self.upem = struct.unpack(">H", self._tbl("head", 18, 2))[0]
        nhm = struct.unpack(">H", self._tbl("hhea", 34, 2))[0]
        hoff, _ = self.t["hmtx"]
        self.adv = [struct.unpack(">H", d[hoff + 4 * i:hoff + 4 * i + 2])[0] for i in range(nhm)]
        self.cmap = self._read_cmap()
        self.cap = None
        if "OS/2" in self.t:
            o, ln = self.t["OS/2"]
            ver = struct.unpack(">H", d[o:o + 2])[0]
            if ver >= 2 and ln >= 90:
                self.cap = struct.unpack(">h", d[o + 88:o + 90])[0]
        if not self.cap:
            self.cap = self._glyph_ymax("H")  # e.g. Impact's OS/2 table is v1 (no sCapHeight)
        if not self.cap:
            self.cap = int(0.75 * self.upem)
        n = self.cmap.get(ord("N"))
        self.fallback_adv = self.adv[min(n, len(self.adv) - 1)] if n is not None else self.upem // 2

    def _glyph_ymax(self, ch):
        try:
            g = self.cmap[ord(ch)]
            d = self.d
            long_loca = struct.unpack(">h", self._tbl("head", 50, 2))[0] == 1
            lo, _ = self.t["loca"]
            go, _ = self.t["glyf"]
            if long_loca:
                off = struct.unpack(">I", d[lo + 4 * g:lo + 4 * g + 4])[0]
            else:
                off = 2 * struct.unpack(">H", d[lo + 2 * g:lo + 2 * g + 2])[0]
            return struct.unpack(">h", d[go + off + 8:go + off + 10])[0]
        except Exception:
            return None

    def _tbl(self, tag, rel, n):
        o, _ = self.t[tag]
        return self.d[o + rel:o + rel + n]

    def _read_cmap(self):
        d = self.d
        base, _ = self.t["cmap"]
        n = struct.unpack(">H", d[base + 2:base + 4])[0]
        subs = []
        for i in range(n):
            p, e, off = struct.unpack(">HHI", d[base + 4 + 8 * i:base + 12 + 8 * i])
            fmt = struct.unpack(">H", d[base + off:base + off + 2])[0]
            subs.append((p, e, fmt, base + off))
        # prefer format 12, then format 4 (unicode platforms)
        for want in (12, 4):
            for p, e, fmt, o in subs:
                if fmt == want and (p == 0 or (p == 3 and e in (1, 10))):
                    return self._fmt12(o) if fmt == 12 else self._fmt4(o)
        raise ValueError("no unicode cmap")

    def _fmt4(self, o):
        d = self.d
        segx2 = struct.unpack(">H", d[o + 6:o + 8])[0]
        seg = segx2 // 2
        ends = struct.unpack(">%dH" % seg, d[o + 14:o + 14 + segx2])
        so = o + 16 + segx2
        starts = struct.unpack(">%dH" % seg, d[so:so + segx2])
        deltas = struct.unpack(">%dh" % seg, d[so + segx2:so + 2 * segx2])
        ro = so + 2 * segx2
        ranges = struct.unpack(">%dH" % seg, d[ro:ro + segx2])
        m = {}
        for i in range(seg):
            for c in range(starts[i], ends[i] + 1):
                if c == 0xFFFF:
                    continue
                if ranges[i] == 0:
                    g = (c + deltas[i]) & 0xFFFF
                else:
                    a = ro + 2 * i + ranges[i] + 2 * (c - starts[i])
                    g = struct.unpack(">H", d[a:a + 2])[0]
                    if g:
                        g = (g + deltas[i]) & 0xFFFF
                if g:
                    m[c] = g
        return m

    def _fmt12(self, o):
        d = self.d
        ng = struct.unpack(">I", d[o + 12:o + 16])[0]
        m = {}
        for i in range(ng):
            s, e, g = struct.unpack(">III", d[o + 16 + 12 * i:o + 28 + 12 * i])
            if e - s > 0x3000:
                continue
            for c in range(s, e + 1):
                m[c] = g + (c - s)
        return m

    def has(self, ch):
        return ord(ch) in self.cmap

    def width(self, text, size):
        tot = 0
        for ch in text:
            g = self.cmap.get(ord(ch))
            tot += self.adv[min(g, len(self.adv) - 1)] if g is not None else self.fallback_adv
        return tot * size / self.upem


# ---------------------------------------------------------------- helpers
def find_font():
    for p in FONT_CANDIDATES:
        p = os.path.expanduser(p)
        if os.path.isfile(p):
            return p
    sys.exit("No Impact (or Anton/League Gothic/Oswald Bold) font found.")


def find_tool(name):
    w = shutil.which(name)
    if w:
        return w
    for dd in FFMPEG_DIRS:
        p = os.path.join(dd, name)
        if os.access(p, os.X_OK):
            return p
    sys.exit(f"{name} not found (looked in PATH and {FFMPEG_DIRS}).")


def section(text, name):
    m = re.search(r"^##\s+" + re.escape(name) + r"\s*\n(.*?)(?=^##\s|\Z)", text, re.S | re.M | re.I)
    return m.group(1).strip() if m else ""


def section_file_path(body):
    """Return a path from a note section containing either a file URI or path."""
    m = re.search(r"file://([^)\s>]+)", body)
    if m:
        return urllib.parse.unquote(m.group(1))
    for line in body.splitlines():
        line = line.strip()
        if line.startswith("/") or line.startswith("~"):
            return os.path.expanduser(line)
    return None


def parse_note(path):
    t = open(path, encoding="utf-8").read()
    fm = re.match(r"---\n(.*?)\n---", t, re.S)
    fm = fm.group(1) if fm else ""
    uid = re.search(r"^uid:\s*(\S+)", fm, re.M)
    if not uid:
        sys.exit("note has no uid in frontmatter")
    uid = uid.group(1).strip()
    hook = " ".join(section(t, "hook").split())
    clip = section_file_path(section(t, "buffer video"))
    if not clip or not os.path.isfile(clip):
        cands = [f for f in os.listdir(BUFFER_DIR)
                 if f.startswith(uid + ".") and f.lower().endswith((".mp4", ".mov", ".m4v"))]
        clip = os.path.join(BUFFER_DIR, cands[0]) if cands else clip
    source = section_file_path(section(t, "source file"))
    dur = None
    ai = section(t, "ai notes")
    m = re.search(r"overlay[^\n]*?first\s+(\d+(?:\.\d+)?)\s*(?:s\b|sec)", ai, re.I)
    if m:
        dur = float(m.group(1))
    return uid, hook, clip, source, dur


def probe(ffprobe, clip):
    out = subprocess.run([ffprobe, "-v", "error", "-print_format", "json", "-show_streams",
                          "-show_format", clip], capture_output=True, text=True, check=True).stdout
    j = json.loads(out)
    v = next(s for s in j["streams"] if s["codec_type"] == "video")
    w, h = int(v["width"]), int(v["height"])
    rot = 0
    for sd in v.get("side_data_list", []) or []:
        if "rotation" in sd:
            rot = int(sd["rotation"])
    if abs(rot) % 180 == 90:
        w, h = h, w
    return w, h, v, j


def normalize(hook, m):
    s = hook.upper()
    for a, b in {"\u2019": "'", "\u2018": "'", "\u201c": '"', "\u201d": '"', "\u2014": "-", "\u2013": "-"}.items():
        if a in s and not m.has(a):
            s = s.replace(a, b)
    return s


def best_split(words, n, widths_fn, maxw):
    """Split words into exactly n lines, all <= maxw, minimizing the widest line (balanced)."""
    best = None
    for cuts in itertools.combinations(range(1, len(words)), n - 1):
        b = (0,) + cuts + (len(words),)
        lines = [" ".join(words[b[i]:b[i + 1]]) for i in range(n)]
        ws = [widths_fn(l) for l in lines]
        if max(ws) > maxw:
            continue
        key = (round(max(ws)), -ws[-1])  # balanced; tie -> slightly heavier bottom line
        if best is None or key < best[0]:
            best = (key, lines, ws)
    return best


def layout(hook, W, H, m):
    text = normalize(hook, m)
    words = text.split()
    scale = W / 1080.0
    bw = max(2, round(BORDER_AT_1080 * scale))
    left = round(W * LEFT_FRAC)
    box_w = round(W * (1 - LEFT_FRAC - RIGHT_FRAC))
    top = round(H * TOP_FRAC)
    bottom_limit = round(H * BOTTOM_LIMIT_FRAC)
    maxw = box_w * WIDTH_SAFETY - 2 * bw
    size = round(W * FONT_FRAC)
    min_size = round(W * MIN_FONT_FRAC)
    while size >= min_size:
        wfn = lambda s: m.width(s, size)
        cap = m.cap * size / m.upem
        pitch = LINE_PITCH_EM * size
        if all(wfn(w) <= maxw for w in words):
            for n in range(1, min(MAX_LINES, len(words)) + 1):
                block_h = bw + cap + (n - 1) * pitch + bw
                if top + block_h > bottom_limit:
                    break
                r = best_split(words, n, wfn, maxw)
                if r:
                    _, lines, ws = r
                    base0 = top + bw + cap
                    baselines = [round(base0 + i * pitch) for i in range(n)]
                    return dict(text=text, lines=lines, line_widths=[round(x) for x in ws],
                                fontsize=size, borderw=bw, cap_height=round(cap),
                                line_pitch=round(pitch, 1), left=left, box_w=box_w,
                                right_margin=W - left - box_w, top=top,
                                block_bottom=round(top + block_h), bottom_limit=bottom_limit,
                                baselines=baselines, W=W, H=H)
        size -= 1
    sys.exit("hook too long to fit in %d lines even at minimum font size" % MAX_LINES)


def build_filter(L, fontfile, tmpdir, duration):
    parts = ["format=yuv420p"]
    enable = f":enable='between(t,0,{duration})'" if duration else ""
    for i, line in enumerate(L["lines"]):
        tf = os.path.join(tmpdir, f"line{i}.txt")
        with open(tf, "w", encoding="utf-8") as f:
            f.write(line)
        parts.append(
            f"drawtext=fontfile={fontfile}:textfile={tf}:expansion=none:fontsize={L['fontsize']}"
            f":fontcolor=white:borderw={L['borderw']}:bordercolor=black"
            f":x={L['left']}+({L['box_w']}-text_w)/2:y_align=baseline:y={L['baselines'][i]}{enable}")
    return ",".join(parts)


def media_duration(j):
    d = (j.get("format") or {}).get("duration")
    if d is not None:
        return float(d)
    for stream in j.get("streams", []):
        if stream.get("codec_type") == "video" and stream.get("duration") is not None:
            return float(stream["duration"])
    return None


def verify_encode(ffprobe, source_clip, output):
    sw, sh, sv, sj = probe(ffprobe, source_clip)
    ow, oh, ov, oj = probe(ffprobe, output)
    sd, od = media_duration(sj), media_duration(oj)
    source_audio = any(x.get("codec_type") == "audio" for x in sj.get("streams", []))
    output_audio = any(x.get("codec_type") == "audio" for x in oj.get("streams", []))
    problems = []
    if sd is None or od is None or abs(od - sd) > 0.2:
        problems.append("duration differs by more than 0.2 s")
    if (sw, sh) != (ow, oh):
        problems.append(f"resolution changed ({sw}x{sh} -> {ow}x{oh})")
    if source_audio and not output_audio:
        problems.append("source has audio but output does not")
    if problems:
        raise RuntimeError("encode verification failed: " + "; ".join(problems))
    return dict(source_duration=sd, output_duration=od, resolution=[ow, oh],
                source_audio=source_audio, output_audio=output_audio)


def update_buffer_link(path, output):
    with open(path, encoding="utf-8") as f:
        text = f.read()
    section_re = re.compile(r"^##\s+buffer video\s*\n(.*?)(?=^##\s|\Z)", re.S | re.M | re.I)
    sm = section_re.search(text)
    if not sm:
        raise RuntimeError("note has no ## buffer video section")
    body = sm.group(1)
    lm = re.search(r"file://[^)\s>]+", body)
    if not lm:
        raise RuntimeError("## buffer video section has no file link")
    new_uri = "file://" + urllib.parse.quote(os.path.abspath(output), safe="/:@")
    new_body = body[:lm.start()] + new_uri + body[lm.end():]
    updated = text[:sm.start(1)] + new_body + text[sm.end(1):]
    with open(path, "w", encoding="utf-8", newline="") as f:
        f.write(updated)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("note")
    ap.add_argument("--hook")
    ap.add_argument("--duration", type=float)
    ap.add_argument("--frame", type=float, help="render a single frame at this time (sec)")
    ap.add_argument("--png", help="output PNG for --frame")
    ap.add_argument("--out")
    ap.add_argument("--videotoolbox", action="store_true")
    ap.add_argument("--force", action="store_true")
    ap.add_argument("--keep-source", action="store_true",
                    help="do not delete the pre-overlay BufferMedia source copy")
    ap.add_argument("--dry-run", action="store_true")
    a = ap.parse_args()

    uid, hook, clip, source, note_dur = parse_note(a.note)
    hook = a.hook or hook
    duration = a.duration if a.duration is not None else note_dur
    if not hook:
        sys.exit("note has an empty ## hook")
    if not clip or not os.path.isfile(clip):
        sys.exit(f"clip not found for uid {uid}: {clip}")
    ffmpeg, ffprobe = find_tool("ffmpeg"), find_tool("ffprobe")
    fontfile = find_font()
    if re.search(r"[':,;\[\]\\ ]", fontfile):
        sys.exit("font path has characters that need filter escaping: " + fontfile)
    W, H, v, j = probe(ffprobe, clip)
    L = layout(hook, W, H, TTFMetrics(fontfile))
    info = dict(uid=uid, clip=clip, source=source, font=fontfile,
                show_seconds=duration or "whole clip", **L)
    print(json.dumps(info, indent=2))
    if a.dry_run:
        return

    tmpdir = tempfile.mkdtemp(prefix="hookovl_", dir="/tmp")
    part = None
    try:
        vf = build_filter(L, fontfile, tmpdir, duration)
        if a.frame is not None:
            png = a.png or f"/tmp/{uid}-overlay-frame.png"
            cmd = [ffmpeg, "-hide_banner", "-loglevel", "error", "-y", "-copyts", "-ss", str(a.frame),
                   "-i", clip, "-vf", vf, "-frames:v", "1", "-update", "1", png]
            subprocess.run(cmd, check=True)
            print("wrote", png, os.path.getsize(png), "bytes")
            return

        out = a.out or os.path.join(BUFFER_DIR, f"{uid}.mp4")
        same_target = os.path.realpath(out) == os.path.realpath(clip)
        if os.path.exists(out) and not a.force and not same_target:
            sys.exit(f"{out} exists (use --force)")
        need = int(os.path.getsize(clip) * 1.5)
        free = shutil.disk_usage(os.path.dirname(out)).free
        if free < need + 500 * 1024 * 1024:
            sys.exit(f"not enough free space on {os.path.dirname(out)}: {free/1e9:.2f} GB free, "
                     f"need ~{need/1e9:.2f} GB + 0.5 GB headroom")
        part = out + f".partial.{os.getpid()}.mp4"
        if a.videotoolbox:
            br = max(int(int(v.get("bit_rate", 0) or 0) * 1.2), 20_000_000)
            venc = ["-c:v", "h264_videotoolbox", "-b:v", str(br), "-profile:v", "high"]
        else:
            venc = ["-c:v", "libx264", "-crf", "18", "-preset", "medium", "-profile:v", "high"]
        color = []
        for key, flag in (("color_primaries", "-color_primaries"), ("color_transfer", "-color_trc"),
                          ("color_space", "-colorspace")):
            if v.get(key) and v[key] != "unknown":
                color += [flag, v[key]]
        cmd = [ffmpeg, "-hide_banner", "-loglevel", "error", "-stats", "-y", "-i", clip,
               "-map", "0:v:0", "-map", "0:a?", "-vf", vf, *venc, "-pix_fmt", "yuv420p", *color,
               "-fps_mode", "passthrough", "-c:a", "copy", "-map_metadata", "0",
               "-movflags", "+faststart", part]
        subprocess.run(cmd, check=True)
        verification = verify_encode(ffprobe, clip, part)
        print("verified", json.dumps(verification, sort_keys=True))

        original_match = False
        if source and os.path.isfile(source):
            original_match = os.path.getsize(source) == os.path.getsize(clip)
            if original_match:
                print("source check: original exists with matching byte size")
            else:
                print("source check: original byte size does not match; source kept")
        else:
            print("source check: original is missing; source kept")

        if same_target and not a.keep_source and not original_match:
            raise RuntimeError("refusing to replace the source in place without a matching Shorts/Used original")

        # The part has been verified before this replace.  When clip == out this
        # atomically replaces the pre-overlay source without ever reading from a
        # file that is being overwritten.
        os.replace(part, out)
        part = None
        if not a.keep_source and not same_target and os.path.lexists(clip):
            if original_match:
                os.remove(clip)
                print("deleted pre-overlay source copy", clip)
            else:
                print("kept pre-overlay source copy", clip)
        elif same_target:
            print("replaced pre-overlay source in place", out)
        elif a.keep_source:
            print("kept pre-overlay source copy (--keep-source)", clip)
        print("wrote", out, os.path.getsize(out), "bytes")
        update_buffer_link(a.note, out)
        print("updated note buffer link", a.note)
    finally:
        shutil.rmtree(tmpdir, ignore_errors=True)
        if part and os.path.exists(part):
            os.remove(part)


if __name__ == "__main__":
    main()