#!/usr/bin/env python3
"""Build static-space/superbeam-tokens.js (window.SBT): per replay frame, the token stream the
in-weights model produces for that frame, for the page's "inside the model" panel.

Reads static-space/superbeam-data.js (run build_superbeam_page_data.py first): the ROM and the keys
held each frame. Frames are rebuilt the way beamboot/chip8_fast.FastSession builds them: a checkpoint
prompt of the machine state, then the frame's tokens (drafted by the reference emulator; on the live
run the model accepted the drafted tokens of every frame, which is what live.status = "live" records).
Each frame's tokens are fed through the token replay viewer and its display is checked against the
page's frame rows, so the page never shows tokens for a different frame.

  python scripts/build_token_trace.py
"""
import base64
import json
import sys
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
sys.path.insert(0, str(ROOT))

from beamboot import chip8_fast as F  # noqa: E402
from beamboot import chip8_lm as C    # noqa: E402
from beamboot.chip8 import Chip8      # noqa: E402

DATA = ROOT / "static-space" / "superbeam-data.js"
OUT = ROOT / "static-space" / "superbeam-tokens.js"
SHOW_RECORDS = 6


def fmt(t):
    k, v, a = C.untok(t)
    if k == "VW":
        return "V%X←%d" % (a, v)
    if k == "PCW":
        return "PC←%03X" % v
    if k == "IW":
        return "I←%03X" % v
    if k == "DTW":
        return "DT←%d" % v
    if k == "SPW":
        return "SP←%d" % v
    if k == "MEMW":
        return "M[I+%d]←%d" % (a, v)
    if k in ("CLKH", "CLKL", "ROW", "DL", "DR", "STKW"):
        return "%s %d" % (k, v)
    return k


def main():
    sb = json.loads(DATA.read_text().split("=", 1)[1].rstrip().rstrip(";"))
    rom = base64.b64decode(sb["play"]["rom_b64"])
    ipf = sb["play"]["ipf"]
    session = F.FastSession.__new__(F.FastSession)      # draft path only: no model is loaded
    session.rom, session.ipf = rom, ipf
    session.c = Chip8()
    session.c.load(rom)
    frames = []
    for f, st in enumerate(sb["state"]):
        for key, bit in ((4, 1), (5, 2), (6, 4)):
            session.c.set_key(key, bool(st[0] & bit))
        prompt = F.checkpoint_prompt(session.c, ipf)
        toks = session._draft([])
        replay = F._new_replay(session.c, ipf)
        recs, cur = [], []
        for t in toks:
            cur.append(t)
            if replay.feed(t) in ("instr", "tick"):
                recs.append(cur)
                cur = []
        session.c = replay.c
        if session.c.display_rows() != sb["frames"][f]:
            raise SystemExit("frame %d: token replay does not match the page's frame" % f)
        frames.append({"p": len(prompt), "n": len(toks), "r": len(recs),
                       "h": [[fmt(t) for t in r] for r in recs[:SHOW_RECORDS]]})
    OUT.write_text("window.SBT = " + json.dumps({"frames": frames}, separators=(",", ":")) + ";\n")
    print("wrote %s: %d frames, %d tokens in total" % (OUT, len(frames), sum(x["n"] for x in frames)))


if __name__ == "__main__":
    main()