"""Render two recorded runs (scripts/race_record.py) as a side-by-side race video, replayed on a shared clock.

  python3 scripts/race_render.py race/base.json race/dory-beam.json --out race/race.mp4 [--speed auto] [--hw "an M4 Max"]

Every chunk appears at the moment it was actually streamed (scaled by --speed, which is shown on screen).
Left panel: the first recording (gray); right: the second (accent). Needs ffmpeg.
"""

import argparse
import json
import math
import re
import subprocess
import sys
import textwrap
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont

ROOT = Path(__file__).resolve().parent.parent
W, H, FPS, HOLD = 1600, 900, 20, 5
SURFACE, INK, INK2, MUTED, GRID = "#fcfcfb", "#0b0b0b", "#52514e", "#898781", "#e1e0d9"
GRAY, ACCENT, GOOD, BAD = "#898781", "#2a78d6", "#006300", "#d03b3b"
SANS, MONO = "/System/Library/Fonts/HelveticaNeue.ttc", "/System/Library/Fonts/SFNSMono.ttf"


def font(path, size, index=0):
    return ImageFont.truetype(path, size, index=index)


F_TITLE, F_SUB, F_NAME = font(SANS, 34, 1), font(SANS, 19), font(SANS, 25, 1)
F_STAT, F_MONO, F_CLOCK = font(SANS, 20), font(MONO, 15), font(SANS, 18)
F_BADGE = font("/System/Library/Fonts/SFNS.ttf", 21)  # SF Pro has the ✓/✗ glyphs Helvetica Neue lacks
try:
    F_BADGE.set_variation_by_name("Semibold")
except Exception:
    pass


def last_boxed(text):
    i = text.rfind("\\boxed{")
    if i < 0:
        return None
    depth, j = 0, i + 7
    for k in range(j, len(text)):
        depth += text[k] == "{"
        if text[k] == "}":
            if depth == 0:
                return text[j:k]
            depth -= 1
    return None


def grade(run):
    answer = run["text"].split("</think>")[-1] if "</think>" in run["text"] else ""
    try:
        sys.path.insert(0, str(ROOT / "scripts"))
        from verify import check
        ok, _ = check(answer, run["gold"], "math")
    except Exception:
        ok = (last_boxed(answer) or "").replace(" ", "") == run["gold"].replace(" ", "")
    return bool(ok), last_boxed(answer)


def wrap(text, width):
    out = []
    for para in text.split("\n"):
        out += textwrap.wrap(para, width, replace_whitespace=False, drop_whitespace=True) or [""]
    return out


def fmt_t(s):
    return f"{int(s // 60)}:{int(s % 60):02d}"


def panel(d, x, y, w, h, run, t, color, max_tokens):
    done = t >= run["end"]
    n_ev = sum(1 for e in run["events"] if e[0] <= t)
    text = "".join(e[1] for e in run["events"][:n_ev]).replace("<think>", "").lstrip()
    tokens = round(run["completion_tokens"] * n_ev / max(len(run["events"]), 1))
    # header: identity dot + name; live stats on the right
    d.ellipse((x, y + 9, x + 14, y + 23), fill=color)
    d.text((x + 24, y), run["name"], font=F_NAME, fill=INK)
    stats = f"{tokens:,} tokens  ·  {fmt_t(min(t, run['end']))}"
    d.text((x + w - d.textlength(stats, font=F_STAT), y + 4), stats, font=F_STAT, fill=INK2)
    # progress bar: tokens so far against the longer run's total
    by = y + 44
    d.rounded_rectangle((x, by, x + w, by + 8), radius=4, fill=GRID)
    if tokens:
        d.rounded_rectangle((x, by, x + max(8, w * tokens / max_tokens), by + 8), radius=4, fill=color)
    # streamed text: thinking in secondary ink, the final answer in primary ink
    ty, line_h = by + 26, 20
    cols = int(w / d.textlength("M", font=F_MONO))
    think, sep, answer = text.partition("</think>")
    lines = [(l, INK2) for l in wrap(think, cols)] + ([("", INK2)] + [(l, INK) for l in wrap(answer.strip(), cols)] if sep else [])
    rows = int((h - (ty - y) - (64 if done else 8)) / line_h)
    for i, (l, c) in enumerate(lines[-rows:]):
        d.text((x, ty + i * line_h), l, font=F_MONO, fill=c)
    if done:  # finish badge: status icon + label, never color alone
        ok, boxed = run["grade"]
        label = f"{'✓' if ok else '✗'} {'Correct' if ok else 'Wrong'}: {boxed or '(no answer)'}  ·  finished in {fmt_t(run['end'])}"
        bx, bh = y + h - 48, 44
        d.rounded_rectangle((x, bx, x + w, bx + bh), radius=8, fill="#f0efec")
        d.text((x + 16, bx + 10), label, font=F_BADGE, fill=GOOD if ok else BAD)


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("runs", nargs="+", help="2 runs side by side, or 4 in a 2x2 grid; the first is the baseline")
    ap.add_argument("--title", default="Same problem, same laptop: base Dory vs Dory-Beam-4B")
    ap.add_argument("--out", required=True)
    ap.add_argument("--speed", default="auto", help="playback speed multiplier, or auto (fits in ~45 s)")
    ap.add_argument("--hw", default="an Apple M4 Max")
    args = ap.parse_args()
    runs = [json.loads(Path(p).read_text()) for p in args.runs]
    assert len(runs) in (2, 4)
    for r in runs:
        r["end"] = r["events"][-1][0] if r["events"] else 0
        r["grade"] = grade(r)
    span = max(r["end"] for r in runs)
    speed = max(1, math.ceil(span / 45)) if args.speed == "auto" else float(args.speed)
    max_tokens = max(r["completion_tokens"] for r in runs)
    frames = int(span / speed * FPS) + HOLD * FPS
    prompt = re.sub(r"\s+", " ", runs[0]["prompt"].replace("Put your final answer within \\boxed{}.", "")).strip()

    ff = subprocess.Popen(["ffmpeg", "-y", "-loglevel", "error", "-f", "rawvideo", "-pix_fmt", "rgb24", "-s", f"{W}x{H}",
                           "-r", str(FPS), "-i", "-", "-c:v", "libx264", "-pix_fmt", "yuv420p", "-crf", "23",
                           "-movflags", "+faststart", args.out], stdin=subprocess.PIPE)
    for f in range(frames):
        t = min(f / FPS * speed, span)
        img = Image.new("RGB", (W, H), SURFACE)
        d = ImageDraw.Draw(img)
        d.text((48, 32), args.title, font=F_TITLE, fill=INK)
        sp = "real time" if speed == 1 else f"played at {speed:g}× speed"
        d.text((48, 82), f"llama.cpp Q5_K_M on {args.hw}, {runs[0]['loops']} loops, one run each (seed "
               f"{runs[0]['seed']}). Recorded token timings, {sp}.", font=F_SUB, fill=INK2)
        q = textwrap.shorten(prompt, 150, placeholder=" …")
        d.text((48, 116), f"{runs[0]['problem'].split('-')[0].upper().replace('MATH500', 'MATH-500')}: {q}",
               font=F_SUB, fill=INK)
        d.line((48, 154, W - 48, 154), fill=GRID, width=1)
        pw = (W - 48 * 2 - 40) / 2
        colors = [GRAY] + [ACCENT] * (len(runs) - 1)  # identity is in each panel's name; color marks base vs Beam
        if len(runs) == 2:
            for i, (r, c) in enumerate(zip(runs, colors)):
                panel(d, 48 + i * (pw + 40), 176, pw, H - 176 - 56, r, t, c, max_tokens)
        else:
            ph = (H - 176 - 56 - 24) / 2
            for i, (r, c) in enumerate(zip(runs, colors)):
                panel(d, 48 + (i % 2) * (pw + 40), 176 + (i // 2) * (ph + 24), pw, ph, r, t, c, max_tokens)
        if t >= span:  # all finished: one-line result, in text ink
            a = runs[0]
            if len(runs) == 2:
                b = runs[1]
                summ = (f"{b['name']}: {a['completion_tokens'] / b['completion_tokens']:.1f}× fewer tokens, "
                        f"finished {a['end'] / b['end']:.1f}× sooner than {a['name']}.")
            else:
                summ = "Fewer tokens than base: " + "  ·  ".join(
                    f"{r.get('effort') or 'default'} {a['completion_tokens'] / r['completion_tokens']:.1f}×" for r in runs[1:])
            d.text((48, H - 42), summ, font=F_BADGE, fill=INK)
        clock = f"elapsed {fmt_t(t)}"
        d.text((W - 48 - d.textlength(clock, font=F_CLOCK), H - 40), clock, font=F_CLOCK, fill=MUTED)
        ff.stdin.write(img.tobytes())
    ff.stdin.close()
    ff.wait()
    print(f"wrote {args.out}: {frames / FPS:.0f} s at {speed:g}x; "
          + "; ".join(f"{r['name']} {r['completion_tokens']} tok in {r['end']:.0f}s ok={r['grade'][0]}" for r in runs))


if __name__ == "__main__":
    main()