"""Record deterministic Space replays from real checkpoint forward passes."""
import json
from pathlib import Path
from beamboot.causal_lm import CausalMachine
from beamboot.os import OperatingMonitor, firmware


def main():
    base = CausalMachine.from_pretrained('artifacts/beamboot-monitor')
    dataset = {}
    cases = [('boot', None), ('help', 'help'), ('version', 'version'),
             ('add_7_8', (7, 8)), ('add_15_3', (15, 3)), ('add_0_0', (0, 0))]
    for name, command in cases:
        runtime = CausalMachine.__new__(CausalMachine)
        runtime.model = base.model
        runtime._reset(base.maximum)
        monitor = OperatingMonitor.__new__(OperatingMonitor)
        monitor.labels, monitor.runtime = firmware()[1], runtime
        monitor.started, monitor.waiting = False, None
        trace = []
        original = runtime.step
        def step():
            result = original()
            trace.append({'pc': result.pc, 'acc': result.acc, 'token': runtime.token,
                          'writes': runtime.writes, 'write': result.write, 'output': result.output})
            return result
        runtime.step = step
        boot = monitor.boot()
        if command is None:
            response = boot
        else:
            trace.clear()
            if isinstance(command, tuple): response = str(monitor.add(*command))
            else: response = getattr(monitor, command)()
        dataset[name] = {'boot': boot, 'response': response, 'trace': trace}
        print(name, len(trace), repr(response), flush=True)
    Path('static-space/traces.json').write_text(json.dumps(dataset, separators=(',', ':')))


if __name__ == '__main__': main()