"""Run SUPERBEAM headless with a scripted input sequence.

  python scripts/play_chip8.py                 # replay roms/superbeam_inputs.json
  python scripts/play_chip8.py --generate      # (re)build the input script with the
                                               # emulator search bot, then replay it

Frame model: one frame = IPF instructions, then tick_timers().  Input events
[frame, key, "down"|"up"] are applied before that frame's instructions.
Writes demo/superbeam_frames.json and (with --generate) roms/superbeam_inputs.json.
"""

import argparse
import json
import os
import sys

ROOT = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))
sys.path.insert(0, ROOT)

from beamboot.chip8 import Chip8  # noqa: E402
from beamboot.chip8_asm import assemble, check_subset, longest_path  # noqa: E402

ASM = os.path.join(ROOT, "roms", "superbeam.asm")
ROM = os.path.join(ROOT, "roms", "superbeam.ch8")
INPUTS = os.path.join(ROOT, "roms", "superbeam_inputs.json")
FRAMES = os.path.join(ROOT, "demo", "superbeam_frames.json")

IPF = 300                      # instructions per frame (timer tick period)
KEY_LEFT, KEY_RIGHT, KEY_JUMP = 4, 6, 5
MAX_FRAMES = 600


def build():
    """Assemble roms/superbeam.asm and check that every reachable instruction is one the
    in-weights computer models (beamboot/chip8_asm.check_subset)."""
    with open(ASM) as f:
        rom, symbols = assemble(f.read())
    check_subset(rom, (0x200, symbols["halt"]))
    return rom, symbols


def run_game(rom, symbols, script, ipf=IPF, max_frames=MAX_FRAMES, capture=False):
    """Replay `script`; stop at the end of the frame in which RAM[WIN] becomes 1."""
    m = Chip8()
    m.load(rom)
    win, deaths = symbols["WIN"], symbols["DEATHS"]
    spin = {symbols["wait"], symbols["wait"] + 2, symbols["wait"] + 4}
    events = {}
    for fr, key, act in script:
        events.setdefault(fr, []).append((key, act == "down"))
    frames, work = [], []
    won = False
    f = 0
    while f < max_frames:
        for key, down in events.get(f, ()):
            m.set_key(key, down)
        busy = 0
        for _ in range(ipf):
            if m.step().pc not in spin:
                busy += 1
        work.append(busy)
        if capture:
            frames.append(m.display_rows())
        m.tick_timers()
        f += 1
        if m.mem[win]:
            won = True
            break
    return {"won": won, "frames_run": f, "instructions": m.instructions,
            "deaths": m.mem[deaths], "work": work, "frames": frames, "machine": m}


def generate(rom, symbols):
    """Winning, zero-death input script from the emulator search bot (beamboot/superbeam_levels.solve,
    coins=True: best-first over the 6 env actions at frame_skip=4 that also prefers coins)."""
    from beamboot import superbeam_levels as levels
    plan = levels.solve(rom, max_nodes=4000, coins=True)
    if not plan:
        raise RuntimeError("search bot found no winning plan")
    per_frame = []
    for a in plan:
        held, jump = levels._ACTION_KEYS[a]
        for i in range(levels.FRAME_SKIP):
            per_frame.append(set(held) | ({KEY_JUMP} if jump and i == 0 else set()))
    events, held = [], set()
    for f, keys in enumerate(per_frame):
        for k in sorted(keys - held):
            events.append([f, k, "down"])
        for k in sorted(held - keys):
            events.append([f, k, "up"])
        held = set(keys)
    return events


def ascii_frame(rows):
    return "\n".join("".join("#" if (int(r, 16) >> (63 - c)) & 1 else "." for c in range(64))
                     for r in rows)


def main(argv=None):
    ap = argparse.ArgumentParser()
    ap.add_argument("--generate", action="store_true")
    ap.add_argument("--every", type=int, default=20, help="print an ASCII frame every N frames")
    args = ap.parse_args(argv)

    rom, symbols = build()
    with open(ROM, "wb") as f:
        f.write(rom)
    if args.generate:
        script = generate(rom, symbols)
        with open(INPUTS, "w") as f:
            json.dump(script, f)
    else:
        with open(INPUTS) as f:
            script = json.load(f)

    res = run_game(rom, symbols, script, capture=True)
    for i, rows in enumerate(res["frames"]):
        if i % args.every == 0 or i == len(res["frames"]) - 1:
            print("--- frame %d ---" % i)
            print(ascii_frame(rows))
    work = res["work"]
    bound = longest_path(rom, symbols["wait"] + 6, {symbols["wait"]})
    print("won=%s frames=%d instructions=%d deaths=%d coins=%d ipf=%d work/frame max=%d mean=%.1f "
          "static worst-case=%d rom=%dB"
          % (res["won"], res["frames_run"], res["instructions"], res["deaths"],
             res["machine"].mem[symbols["COINS"]], IPF, max(work[:-1]), sum(work) / len(work), bound,
             len(rom)))
    os.makedirs(os.path.dirname(FRAMES), exist_ok=True)
    with open(FRAMES, "w") as f:
        json.dump({"ipf": IPF, "rom_size": len(rom), "total_instructions": res["instructions"],
                   "frames_run": res["frames_run"], "won": res["won"],
                   "inputs": script, "frames": res["frames"]}, f, separators=(",", ":"))
    return 0 if res["won"] else 1


if __name__ == "__main__":
    raise SystemExit(main())