#!/usr/bin/env python3
"""style_render.py — render a clip in one of the 4 repost/commentary styles.

Pipeline: template.html?style=…&bg=none → playwright transparent overlay PNG → ffmpeg composite over the
cut footage. Two composite modes:
  - overlay (revolt, editorial): footage fills the frame, graphics/scrim/text composite on top.
  - stack  (culture, tweet):     footage sits in the lower region, an opaque top band carries the graphics.

CLI:  python3 style_render.py --source <video> --start S --end E --style revolt \
        --out out.mp4 --params '{"headline":"…","brand":"19 KEYS"}'
"""
import argparse
import json
import subprocess
import tempfile
import urllib.parse
from pathlib import Path

HERE = Path(__file__).resolve().parent
STUDIO = HERE.parent
TEMPLATE = STUDIO / "dashboard" / "styles" / "template.html"
FFMPEG = "/opt/homebrew/opt/ffmpeg-full/bin/ffmpeg"

# top opaque band height per stacked style (must match template.html)
BAND_H = {"culture": 560, "tweet": 820}
OVERLAY_STYLES = {"revolt", "editorial", "source"}
W, H = 1080, 1920
# xpost is the viral "tweet card over the FULL wide two-shot" — a SQUARE 1080x1080 where the whole
# clip is shown UNCROPPED (contained) in the bottom region and the black card carries the top.
XPOST_W, XPOST_H = 1080, 1080
XPOST_CARD = 472        # black card band (top); must match .xpost .card height in template.html
XPOST_VIDEO = XPOST_H - XPOST_CARD  # 608 — a full-width 16:9 clip fits here exactly, zero crop, zero bars
# conversation = HIGH LVL CONVERSATIONS card: gold stone header, clip window, gold headline+speaker, mono transcript
CONV_HEAD = 160         # must match .conversation .cvhead height
CONV_VIDEO = 608        # 16:9 clip, full width, no crop; window is y=CONV_HEAD .. CONV_HEAD+CONV_VIDEO (=768)
# quote = clean centered quote over the FULL wide frame. Aspect presets (w,h); default "wide" 16:9 = zero crop
QUOTE_ASPECTS = {"wide": (1920, 1080), "tall": (1080, 1920), "square": (1080, 1080)}
# docu = gold header + clip window + gold quote band + SCROLLING screenplay transcript (the bezos-gold style)
DOCU_HEAD = 218         # REAL "HIGH LVL CONVERSATIONS" header graphic (assets/hlc-header.png) — never recreated
DOCU_VIDEO = 608        # 16:9 clip window, y = 218 .. 826
DOCU_SCROLL_TOP = 1006  # scrolling transcript begins here (below the quote band 826..1006); window = H - this
DOCU_HEADER_IMG = HERE.parent / "dashboard" / "styles" / "assets" / "hlc-header.png"
_REACT = {"yeah", "right", "exactly", "no", "yes", "mm", "mmhmm", "mm-hmm", "true", "facts", "wow", "okay",
          "ok", "sure", "correct", "absolutely", "definitely", "word", "hmm", "yep", "nah"}

SCROLL_HTML = """<!doctype html><html><head><meta charset="utf-8">
<link rel="preconnect" href="https://fonts.gstatic.com" crossorigin>
<link href="https://fonts.googleapis.com/css2?family=Space+Mono:wght@400;700&display=swap" rel="stylesheet">
<style>*{margin:0;padding:0;box-sizing:border-box}html,body{width:1080px;background:transparent}
.wrap{width:1080px;padding:40px 62px 900px;color:#ECECEC;font-family:'Space Mono',monospace;font-weight:400;font-size:38px;line-height:1.6}
.wrap .r{margin:0 0 4px}.wrap .t{margin:24px 0 4px}.wrap b{color:#D2A93F;font-weight:700}
</style></head><body><div class="wrap">__BODY__</div></body></html>"""


def read_png_size(p):
    import struct
    with open(p, "rb") as f:
        head = f.read(24)
    return struct.unpack(">II", head[16:24])  # (w, h) from the IHDR chunk


def format_screenplay(words_path, start, end):
    """Verbatim transcript for the span as screenplay dialogue: reactions/short turns get a '– ' dash line."""
    try:
        words = json.loads(Path(words_path).read_text())
        words = words["words"] if isinstance(words, dict) else words
    except Exception:
        return "<div class='r'></div>"
    span = [w for w in words if w.get("type", "word") == "word" and start - 0.05 <= w.get("start", 0) <= end + 0.05]
    # split into sentences
    sents, buf = [], []
    for w in span:
        buf.append(w["text"])
        if w["text"].strip().endswith((".", "?", "!")) and len(buf) >= 2:
            sents.append(" ".join(buf).strip())
            buf = []
    if buf:
        sents.append(" ".join(buf).strip())
    rows = []
    for s in sents:
        first = s.lower().lstrip("\"'").split()[:1]
        is_turn = len(s.split()) <= 7 or (first and first[0].strip(".,!?") in _REACT)
        esc = s.replace("&", "&amp;").replace("<", "&lt;").replace(">", "&gt;")
        rows.append(f'<div class="t">&ndash; {esc}</div>' if is_turn else f'<div class="r">{esc}</div>')
    return "\n".join(rows) or "<div class='r'></div>"


def render_scroll_png(script_html, out_png):
    """Render the screenplay transcript as a tall transparent PNG; return its height."""
    from playwright.sync_api import sync_playwright
    html = SCROLL_HTML.replace("__BODY__", script_html)
    with sync_playwright() as pw:
        b = pw.chromium.launch()
        pg = b.new_page(viewport={"width": 1080, "height": 1400}, device_scale_factor=1)
        pg.set_content(html, wait_until="networkidle")
        pg.wait_for_timeout(500)
        pg.screenshot(path=str(out_png), omit_background=True, full_page=True)
        b.close()
    return read_png_size(out_png)[1]


def _ass_ts(t):
    t = max(0.0, t)
    h = int(t // 3600); m = int((t % 3600) // 60); s = t % 60
    return f"{h}:{m:02d}:{s:05.2f}"


def build_ass(words_path, start, end, out_ass, max_words=5, max_chars=30):
    """Word-synced captions (ASS) from the transcript span, clip-relative. Returns True if written."""
    try:
        words = json.loads(Path(words_path).read_text()).get("words", [])
    except Exception:
        return False
    span = [w for w in words if w.get("type", "word") == "word"
            and start - 0.05 <= w.get("start", 0) <= end + 0.05]
    if not span:
        return False
    cues, buf, t0 = [], [], None
    for w in span:
        if t0 is None:
            t0 = w["start"]
        buf.append(w)
        txt = " ".join(x["text"] for x in buf)
        if len(buf) >= max_words or len(txt) >= max_chars or w["text"].strip().endswith((".", "?", "!", ",")):
            cues.append((t0 - start, w["end"] - start, txt.strip()))
            buf, t0 = [], None
    if buf:
        cues.append((t0 - start, span[-1]["end"] - start, " ".join(x["text"] for x in buf).strip()))
    header = (
        "[Script Info]\nScriptType: v4.00+\nPlayResX: 1080\nPlayResY: 1920\nWrapStyle: 2\nScaledBorderAndShadow: yes\n\n"
        "[V4+ Styles]\nFormat: Name, Fontname, Fontsize, PrimaryColour, SecondaryColour, OutlineColour, BackColour, "
        "Bold, Italic, Underline, StrikeOut, ScaleX, ScaleY, Spacing, Angle, BorderStyle, Outline, Shadow, "
        "Alignment, MarginL, MarginR, MarginV, Encoding\n"
        "Style: Cap,Arial,60,&H00FFFFFF,&H00FFFFFF,&H00000000,&H64000000,-1,0,0,0,100,100,0,0,1,5,2,2,90,90,175,1\n\n"
        "[Events]\nFormat: Layer, Start, End, Style, Name, MarginL, MarginR, MarginV, Effect, Text\n"
    )
    rows = []
    for a, b, txt in cues:
        if b <= a:
            b = a + 0.4
        txt = txt.replace("\n", " ").replace("{", "(").replace("}", ")")
        rows.append(f"Dialogue: 0,{_ass_ts(a)},{_ass_ts(b)},Cap,,0,0,0,,{txt}")
    Path(out_ass).write_text(header + "\n".join(rows) + "\n")
    return True


def render_overlay_png(style, params, out_png, w=W, h=H):
    """Render the transparent graphics overlay (w x h) for a style."""
    from playwright.sync_api import sync_playwright
    q = {"style": style, "bg": "none", "w": w, "h": h, **{k: v for k, v in params.items() if v is not None}}
    url = TEMPLATE.as_uri() + "?" + urllib.parse.urlencode(q)
    with sync_playwright() as pw:
        b = pw.chromium.launch()
        pg = b.new_page(viewport={"width": w, "height": h}, device_scale_factor=1)
        pg.goto(url, wait_until="networkidle")
        pg.wait_for_timeout(600)
        pg.screenshot(path=str(out_png), omit_background=True,
                      clip={"x": 0, "y": 0, "width": w, "height": h})
        b.close()


def _encode(ff, ss, to, source, overlay, fc, vmap, out):
    """Composite footage (span ss..to) with the overlay PNG under filter chain fc, encode, return path."""
    cmd = [ff, "-y", "-ss", ss, "-to", to, "-i", str(source), "-i", str(overlay),
           "-filter_complex", fc, "-map", vmap, "-map", "0:a?",
           "-c:v", "libx264", "-pix_fmt", "yuv420p", "-preset", "medium", "-crf", "18",
           "-c:a", "aac", "-b:a", "160k", "-movflags", "+faststart", "-shortest", str(out)]
    r = subprocess.run(cmd, capture_output=True, text=True)
    if not Path(out).exists():
        raise RuntimeError("ffmpeg failed: " + (r.stderr or "")[-400:])
    return str(out)


def render(source, start, end, style, out, params=None, words=None, captions="auto"):
    params = params or {}
    out = Path(out)
    out.parent.mkdir(parents=True, exist_ok=True)
    # card styles bake their own text — no burned captions
    NO_CAP = {"xpost", "conversation", "quote", "docu"}
    cap_on = captions == "on" or (captions == "auto" and style not in OVERLAY_STYLES and style not in NO_CAP)
    with tempfile.TemporaryDirectory() as td:
        overlay = Path(td) / "overlay.png"
        ss, to = f"{float(start):.3f}", f"{float(end):.3f}"
        if style == "docu":
            # gold header + quote FRAME (overlay), clip in the window, screenplay transcript SCROLLING below
            render_overlay_png("docu", params, overlay, w=W, h=H)
            scroll_png = Path(td) / "scroll.png"
            script_html = params.get("script") or (
                format_screenplay(words, float(start), float(end)) if words and Path(words).exists()
                else "<div class='r'>(transcript)</div>")
            Hs = render_scroll_png(script_html, scroll_png)
            dur = max(0.1, float(end) - float(start))
            win = H - DOCU_SCROLL_TOP           # 914
            maxy = max(0, Hs - win)
            # inputs: 0 source · 1 frame(quote) · 2 scroll · 3 REAL header graphic (overlaid last, on top)
            fc = (f"[0:v]scale={W}:{DOCU_VIDEO}:force_original_aspect_ratio=decrease,"
                  f"pad={W}:{DOCU_VIDEO}:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1,"
                  f"pad={W}:{H}:0:{DOCU_HEAD}:color=black[base];"
                  f"[2:v]crop={W}:{win}:0:'min(t/{dur:.3f}\\,1)*{maxy}',setsar=1[sc];"
                  f"[base][sc]overlay=0:{DOCU_SCROLL_TOP}:shortest=1[b2];"
                  f"[b2][1:v]overlay=0:0[b3];"
                  f"[b3][3:v]overlay=0:0[v]")
            cmd = [FFMPEG, "-y", "-ss", ss, "-to", to, "-i", str(source), "-i", str(overlay),
                   "-loop", "1", "-i", str(scroll_png), "-i", str(DOCU_HEADER_IMG),
                   "-filter_complex", fc, "-map", "[v]", "-map", "0:a?",
                   "-c:v", "libx264", "-pix_fmt", "yuv420p", "-preset", "medium", "-crf", "18",
                   "-c:a", "aac", "-b:a", "160k", "-movflags", "+faststart", "-shortest", str(out)]
            r = subprocess.run(cmd, capture_output=True, text=True)
            if not out.exists():
                raise RuntimeError("ffmpeg failed: " + (r.stderr or "")[-500:])
            return str(out)
        if style == "conversation":
            # 9:16. Contain the whole clip (no crop) into the window under the gold header; card overlays the rest.
            render_overlay_png(style, params, overlay, w=W, h=H)
            fc = (f"[0:v]scale={W}:{CONV_VIDEO}:force_original_aspect_ratio=decrease,"
                  f"pad={W}:{CONV_VIDEO}:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1,"
                  f"pad={W}:{H}:0:{CONV_HEAD}:color=black[bg];[bg][1:v]overlay=0:0[v]")
            return _encode(FFMPEG, ss, to, source, overlay, fc, "[v]", out)
        if style == "quote":
            # clean quote over the FULL wide frame; footage fills the chosen aspect (default 16:9 = zero crop)
            aspect = str(params.get("aspect", "wide")).lower()
            qw, qh = QUOTE_ASPECTS.get(aspect, QUOTE_ASPECTS["wide"])
            render_overlay_png(style, params, overlay, w=qw, h=qh)
            fc = (f"[0:v]scale={qw}:{qh}:force_original_aspect_ratio=increase,"
                  f"crop={qw}:{qh},setsar=1[bg];[bg][1:v]overlay=0:0[v]")
            return _encode(FFMPEG, ss, to, source, overlay, fc, "[v]", out)
        if style == "xpost":
            # SQUARE 1080x1080. Contain the WHOLE clip (no crop) into the bottom 608px, card overlays the top.
            render_overlay_png(style, params, overlay, w=XPOST_W, h=XPOST_H)
            fc = (f"[0:v]scale={XPOST_W}:{XPOST_VIDEO}:force_original_aspect_ratio=decrease,"
                  f"pad={XPOST_W}:{XPOST_VIDEO}:(ow-iw)/2:(oh-ih)/2:color=black,setsar=1,"
                  f"pad={XPOST_W}:{XPOST_H}:0:{XPOST_CARD}:color=black[bg];"
                  f"[bg][1:v]overlay=0:0[v]")
            return _encode(FFMPEG, ss, to, source, overlay, fc, "[v]", out)
        render_overlay_png(style, params, overlay)
        if style in OVERLAY_STYLES:
            fc = (f"[0:v]scale={W}:{H}:force_original_aspect_ratio=increase,"
                  f"crop={W}:{H},setsar=1[bg];[bg][1:v]overlay=0:0[v]")
        else:  # stacked: footage in the lower region, band on top
            band = BAND_H[style]
            lower = H - band
            fc = (f"[0:v]scale={W}:{lower}:force_original_aspect_ratio=increase,"
                  f"crop={W}:{lower},setsar=1,pad={W}:{H}:0:{band}:color=black[bg];"
                  f"[bg][1:v]overlay=0:0[v]")
        vmap = "[v]"
        if cap_on and words and Path(words).exists():
            assf = Path(td) / "cap.ass"
            if build_ass(words, float(start), float(end), assf):
                esc = str(assf).replace("\\", "\\\\").replace(":", "\\:")
                fc += f";[v]ass={esc}[vc]"
                vmap = "[vc]"
        cmd = [FFMPEG, "-y", "-ss", ss, "-to", to, "-i", str(source),
               "-i", str(overlay), "-filter_complex", fc,
               "-map", vmap, "-map", "0:a?",
               "-c:v", "libx264", "-pix_fmt", "yuv420p", "-preset", "medium", "-crf", "18",
               "-c:a", "aac", "-b:a", "160k", "-movflags", "+faststart", "-shortest", str(out)]
        r = subprocess.run(cmd, capture_output=True, text=True)
        if not out.exists():
            raise RuntimeError("ffmpeg failed: " + (r.stderr or "")[-400:])
    return str(out)


if __name__ == "__main__":
    ap = argparse.ArgumentParser()
    ap.add_argument("--source", required=True)
    ap.add_argument("--start", type=float, required=True)
    ap.add_argument("--end", type=float, required=True)
    ap.add_argument("--style", required=True,
                    choices=["revolt", "editorial", "culture", "tweet", "xpost", "source", "conversation", "quote", "docu"])
    ap.add_argument("--out", required=True)
    ap.add_argument("--params", default="{}")
    ap.add_argument("--words", default=None, help="transcript words.json for word-synced captions")
    ap.add_argument("--captions", default="auto", choices=["auto", "on", "off"])
    a = ap.parse_args()
    print(render(a.source, a.start, a.end, a.style, a.out, json.loads(a.params),
                 words=a.words, captions=a.captions))
