#!/usr/bin/env python3
"""Build static-space/superbeam-data.js (window.SB) for index.html.  Re-runnable.

Sources (priority order for the player replay):
  1. results/beam_live_weights_run.json: Beam played LIVE against the in-weights computer.  Its
     best attempt's key_schedule is replayed on the current emulator; if every frame equals
     weights.frame_rows the run is for the current ROM -> live.status = "live".  A partial prefix
     -> "partial"; no match (file missing / old ROM) -> "pending".
  2. otherwise the best run in results/levels/beam_level<seed>_run.json (held-out test level,
     emulator backend): shown as a replay, in-weights frames = 0.
Leaderboard: every results/levels/beam_level*_run.json; baselines: results/gym_baselines_levels.json.
"""
import glob
import json
import re
import sys
from pathlib import Path

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

from beamboot.superbeam_gym import (SuperBeamEnv, FLAGCOL, START_X, WIN_X,  # noqa: E402
                                    frames_from_events)

OUT = ROOT / "static-space" / "superbeam-data.js"


def load(path):
    p = ROOT / path
    return json.load(open(p)) if p.exists() else None


def replay(events, n, level="classic"):
    """Emulator replay of a key schedule.  Returns per-frame display rows, state
    [keys mask (1 left, 2 jump, 4 right), VA x, VB y, VD camera], walkers [[alive, col, off], x2],
    coins, plus the initial level bytes, level changes [[frame, col, byte]] and walkers."""
    env = SuperBeamEnv(backend="emulator", level=level)
    env.reset()
    m = env.backend.m
    et = env.sym["ETAB"]
    lv0 = list(env.level_cols)      # TMAP in RAM is only filled once the ROM has started
    cur = list(lv0)
    rows, state, walk, coins, changes = [], [], [], [], []
    for f, keys in enumerate(frames_from_events(events, n)):
        env.backend.advance(keys)
        rows.append(m.display_rows())
        mask = (1 if 4 in keys else 0) | (2 if 5 in keys else 0) | (4 if 6 in keys else 0)
        state.append([mask, m.v[0xA], m.v[0xB], m.v[0xD]])
        walk.append([[m.mem[et + 4 * e + i] for i in range(3)] for e in range(2)])
        coins.append(env.coins())
        lv = env.live_level()
        for c in range(len(lv)):
            if lv[c] != cur[c]:
                changes.append([f, c, lv[c]])
        cur = lv
    return {"rows": rows, "state": state, "walkers": walk, "coins": coins, "level0": lv0,
            "changes": changes, "level_cols": list(env.level_cols), "level_seed": env.level_seed,
            "start_walkers": [list(w) for w in env.walkers]}


def play_data(level):
    """What the page's in-browser player needs: the exact ROM of the shown level (the same bytes the
    in-weights model runs), the RAM addresses it reads and the tile-code -> level-byte table."""
    import base64
    from beamboot import superbeam_levels as L
    if level == "classic":
        rom = (ROOT / "roms" / "superbeam.ch8").read_bytes()
    else:
        env = SuperBeamEnv(backend="emulator", level=level)
        env.reset()
        rom = bytes(env.backend.rom)
    _, sym = L._symbols()
    return {"rom_b64": base64.b64encode(rom).decode(), "ipf": 300,
            "sym": {k: sym[k] for k in ("TMAP", "WIN", "DEATHS", "COINS", "ETAB")},
            "level_len": L.LEVEL_LEN, "code_of_tile": {str(k): v for k, v in L.CODE_OF_TILE.items()},
            "keys": {"left": 4, "jump": 5, "right": 6}}


def excerpt(text, limit=200):
    text = " ".join((text or "").split())
    if not text:
        return ""
    sents = [s.strip() for s in re.split(r"(?<=[.!?])\s+", text) if s.strip()]
    pick = sents[-1] if sents else text
    i = len(sents) - 1
    while i > 0 and len(sents) - i < 3 and len(sents[i - 1]) + 1 + len(pick) <= limit:
        i -= 1
        pick = sents[i] + " " + pick
    if len(pick) > limit:
        pick = pick[: limit - 1].rstrip() + "…"
    return pick


def reasoning(req):
    try:
        msg = req["response"]["choices"][0]["message"]
    except (KeyError, IndexError, TypeError):
        return ""
    return msg.get("reasoning_content") or msg.get("reasoning") or msg.get("content") or ""


def attempts_summary(run):
    return [{"attempt": a["attempt"], "status": a["status"], "won": a["won"],
             "progress": round(a["progress"], 4), "frames": a["frames"],
             "requests": a["requests_used"]} for a in run["attempts"]]


def best_attempt(run):
    return run["attempts"][run["best_attempt"] - 1]


def actions_of(run, att):
    out = []
    for r in att["requests"]:
        end = r["frames"]
        out.append({"request": r["request"], "action": r["action"],
                    "frame_start": max(0, end - run["frame_skip"]), "frame_end": end - 1,
                    "progress": round(r["progress"], 4), "excerpt": excerpt(reasoning(r))})
    return out


def main():
    notes = []
    live_run = load("results/beam_live_weights_run.json")
    live = {"status": "pending", "reason": "no live run file", "frames_equal": 0, "attempts": []}
    player = None          # (run, attempt, level, weights dict or None)
    if live_run:
        att = best_attempt(live_run)
        lvl = live_run.get("level", "classic")
        w = att.get("weights") or {}
        n = att["frames"]
        rep = replay(att["key_schedule"], n, lvl)
        fr = w.get("frame_rows", [])
        k = 0
        while k < len(fr) and k < n and fr[k] == rep["rows"][k]:
            k += 1
        status = "live" if (k == n and len(fr) == n) else ("partial" if k > 0 else "pending")
        live = {"status": status, "frames_equal": k, "frames_model": len(fr), "attempts": attempts_summary(live_run),
                "reason": "" if status == "live" else
                ("model frames equal the emulator for the first %d of %d frames only" % (k, n) if status == "partial" else
                 "live run's model frames do not reproduce on the current ROM (run is for an older game version)")}
        if status != "pending":
            player = (live_run, att, lvl, w)
            player_rep = rep
        else:
            notes.append("live run file is for an older ROM; not rendered")
    if player is None:
        cands = []
        for p in sorted(glob.glob(str(ROOT / "results/levels/beam_level*_run.json"))):
            r = json.load(open(p))
            if r.get("attempts"):
                cands.append(r)
        wins = [r for r in cands if best_attempt(r)["won"]]
        pick = wins[0] if wins else (cands[0] if cands else None)
        if pick is None:
            sys.exit("no usable player run (live run stale and no results/levels runs)")
        att = best_attempt(pick)
        player = (pick, att, pick["level"], None)
        player_rep = replay(att["key_schedule"], att["frames"], pick["level"])
    run, att, lvl, w = player
    n = att["frames"]
    rep = player_rep

    # in-weights info
    inw = {"n": 0, "verified": False}
    model = {"cls": "BeamForCausalLM", "layers": 7, "hidden": 640, "dtype": "float32",
             "tok_per_s": None, "frames_all_draft_accepted": None, "compute_seconds": None}
    if w and live["status"] in ("live", "partial"):
        pf = w.get("per_frame", [])
        toks, secs = sum(p["tokens"] for p in pf), sum(p["s"] for p in pf)
        inw = {"n": live["frames_equal"], "verified": live["status"] == "live"}
        model.update(tok_per_s=round(toks / secs, 1) if secs else None, n_tokens=toks,
                     frames_all_draft_accepted=w.get("frames_all_draft_accepted"),
                     compute_seconds=w.get("compute_seconds"),
                     s_per_frame=round(secs / len(pf), 2) if pf else None)

    # leaderboard: Beam on held-out test levels
    board = []
    for p in sorted(glob.glob(str(ROOT / "results/levels/beam_level*_run.json"))):
        r = json.load(open(p))
        if not r.get("attempts"):
            continue
        a = best_attempt(r)
        rr = replay(a["key_schedule"], a["frames"], r["level"])
        board.append({"seed": r["level"], "won": bool(a["won"]), "progress": round(a["progress"], 4),
                      "coins": rr["coins"][-1], "coins_total": sum(1 for b in rr["level_cols"] if b & 0x0C),
                      "requests": r["total_requests"], "attempts": len(r["attempts"]),
                      "frames": a["frames"], "model": r["model"]})
    base = load("results/gym_baselines_levels.json")
    baselines = []
    if base:
        baselines = [{"agent": b["agent"], "levels": b["episodes"], "levels_won": b["levels_won"],
                      "mean_progress": round(b["mean_progress"], 3), "mean_coins": round(b["mean_coins"], 2)}
                     for b in base["results"]]
    bw = load("results/beam_weights.json")             # network whose weights Beam wrote (scripts/beam_weights.py)
    if bw and bw.get("test"):
        t = bw["test"]
        baselines.insert(0, {"agent": "beam_weights", "levels": t["levels"], "levels_won": t["won"],
                             "mean_progress": t["mean_progress"],
                             "mean_coins": round(sum(r["coins"] for r in t["detail"]) / t["levels"], 2),
                             "rounds": bw.get("best_round")})
    beam_row = None
    if board:
        beam_row = {"model": board[0]["model"], "levels": len(board),
                    "levels_won": sum(b["won"] for b in board),
                    "mean_progress": round(sum(b["progress"] for b in board) / len(board), 3),
                    "mean_coins": round(sum(b["coins"] for b in board) / len(board), 2)}

    data = {"n_frames": n, "frames": rep["rows"], "state": rep["state"], "walkers": rep["walkers"],
            "coins": rep["coins"], "level": rep["level0"], "level_changes": rep["changes"],
            "level_seed": rep["level_seed"], "level_kind": "classic" if lvl == "classic" else "test seed",
            "start_walkers": rep["start_walkers"],
            "flag_col": FLAGCOL, "start_x": START_X, "win_x": WIN_X,
            "actions": actions_of(run, att), "inweights": inw, "live": live, "model": model,
            "player": {"model": run["model"], "backend": att.get("backend"), "won": att["won"],
                       "attempt": att["attempt"], "attempts": attempts_summary(run),
                       "requests": run["total_requests"], "frames": n,
                       "progress": round(att["progress"], 4), "coins": rep["coins"][-1],
                       "coins_total": sum(1 for b in rep["level_cols"] if b & 0x0C),
                       "frame_skip": run["frame_skip"], "ipf": run["ipf"],
                       "source": "live in-weights run" if live["status"] in ("live", "partial")
                       else "held-out level run (emulator backend)"},
            "leaderboard": {"beam": beam_row, "beam_levels": board, "baselines": baselines,
                            "baseline_levels": (base or {}).get("levels"),
                            "difficulty": (base or {}).get("difficulty")},
            "notes": notes,
            "play": play_data(lvl)}
    OUT.write_text("window.SB = " + json.dumps(data, separators=(",", ":")) + ";\n")
    print("wrote %s: %d frames, %d actions, live=%s (equal %d), player=%s, board=%d levels"
          % (OUT, n, len(data["actions"]), live["status"], live["frames_equal"], data["player"]["source"], len(board)))


if __name__ == "__main__":
    main()