#!/usr/bin/env python3
# c64trace.py -- ride a gate stream on the real game in VICE and print
# the trace stsim prints, line for line, so the two can be compared byte
# for byte (gate.sh c64).
#
#   c64trace.py <level letter> <ticks> <demobuf.json> [rawOut]
#   c64trace.py --prepare      make the cached snapshot (below) and exit
#
# demobuf.json is stsim's STSIM_DEMOBUF line: the whole 640-byte playback
# buffer ($0902) as the sim starts the ride with it. rawOut, if given,
# receives each tick's screen RAM, color RAM and charset in stsim's
# STSIM_RAW layout.
#
# How the ride is set up on the C64, all from the game's own attract demo:
#   * The original disk boots once (the crack's intro is skipped at its
#     key test, the trainer menu answered with START GAME, every cheat
#     off) to the point the first attract demo picks its screen ($50C6);
#     VICE's snapshot of that moment is cached in build/c64trace/.
#   * Each ride restores the snapshot, makes every demo slot ($50C1..
#     $50C4) the wanted screen, and when the screen and its recording are
#     loaded ($625F, the disk load is skipped) writes the sim's playback
#     buffer over $0902, zeroes the flame parity, wave index and walk
#     parity ($716C/$716E/$7161) that the title leaves behind -- the sim
#     starts its rides from a fresh game -- and sets the screens-played
#     count so that $6211's INC makes it the screen's place in the 24-hour
#     shift (A = 1 .. Y = 25), as stsim's rides have it.
#   * A finished wreck takes the real game's path ($6B52 -> $6B64: the cab
#     count drops and the cab respawns at $5F27) instead of the demo's
#     return to the title, as stsim's rides do. A ride ends when the cab
#     leaves the screen ($6BFB, end 1), the cabs run out ($6BD0, end 4),
#     the playback reaches the buffer's last pair (end "stream"), or after
#     <ticks> ticks.
#   * Every $5F40 (the main loop's top) the state is read and printed.
#
# Needs VICE's x64sc (VICE=path, else on PATH; VICE_DATA=its ROM folder
# when x64sc cannot find its own), Xvfb, and the original disk image
# (STAXI_D64, default stuff/SPACETAX.D64). Exits 2 when one is missing.
# The disk and everything derived from it stay out of the repository.

import json
import os
import shutil
import signal
import socket
import struct
import subprocess
import sys
import time

HERE   = os.path.dirname(os.path.abspath(__file__))
REPO   = os.path.normpath(os.path.join(HERE, '..', '..'))
CACHE  = os.path.join(REPO, 'build', 'c64trace')
D64    = os.environ.get('STAXI_D64', os.path.join(REPO, 'stuff', 'SPACETAX.D64'))
LIFE_S = 900          # an emulator outlives no ride by more than this

# VICE binary monitor protocol (src/monitor/monitor_binary.c).
STX           = 0x02
API_VERSION   = 0x02
CMD_MEM_GET   = 0x01
CMD_MEM_SET   = 0x02
CMD_CP_SET    = 0x12
CMD_CP_DELETE = 0x13
CMD_REGS_SET  = 0x32
CMD_REGS_GET  = 0x31
CMD_DUMP      = 0x41
CMD_UNDUMP    = 0x42
CMD_KEYS      = 0x72
CMD_PING      = 0x81
CMD_BANKS     = 0x82
CMD_EXIT      = 0xAA
CMD_QUIT      = 0xBB
EV_CHECKPOINT = 0x11
EV_JAM        = 0x61
EV_STOPPED    = 0x62
REG_PC        = 0x03
OP_EXEC       = 0x04

# The game's addresses.
INTRO_KEY_TEST = 0x0994   # the crack's "space pressed?" test
INTRO_GO       = 0x099E   # ...and where it goes when it is
DEMO_PICK      = 0x50C6   # $5010's demo branch: a demo picks its screen
DEMO_SLOTS     = 0x50C1   # "HWTX", the demo rotation
DEMO_LOAD      = 0x625F   # JSR $98C4: load the screen's recording
DEMO_LOADED    = 0x6262   # ...returned
DEMO_BUFFER    = 0x0902
CARRIED        = (0x716C, 0x716E, 0x7161)
SCREENS_PLAYED = 0x7215
TICK_TOP       = 0x5F40
CAB_LEFT       = 0x6BFB
WRECK_DONE     = 0x6B52
LOSE_CAB       = 0x6B64
PLAYER_OUT     = 0x6BD0
LAST_PAIR      = 638      # the playback buffer's last (mask, count) pair
END_EXIT       = 1        # StTickResultE: ST_TICK_LEVEL_EXIT
END_OUT        = 4        # ST_TICK_PLAYER_OUT
TRAINER_ROWS   = 12       # cursor-downs from the first option to START GAME
KEY_DOWN       = '\x11'
KEY_RETURN     = '\x0d'
KEY_WAIT_S     = 0.3
INTRO_WAIT_S   = 6
BOOT_WAIT_S    = 300
STOP_WAIT_S    = 120


class Monitor:
    """A VICE binary-monitor connection."""

    def __init__(self, port):
        self.sock = None
        for _ in range(200):
            try:
                self.sock = socket.create_connection(('127.0.0.1', port), timeout=2)
                break
            except OSError:
                time.sleep(0.1)
        if self.sock is None:
            raise RuntimeError('no VICE monitor on port %d' % port)
        self.sock.settimeout(None)
        self.reqId  = 1
        self.events = []
        self.buf    = b''

    def recvExact(self, n):
        while len(self.buf) < n:
            chunk = self.sock.recv(1 << 20)
            if not chunk:
                raise EOFError('VICE closed the monitor')
            self.buf += chunk
        out, self.buf = self.buf[:n], self.buf[n:]
        return out

    def packet(self):
        hdr    = self.recvExact(12)
        length = struct.unpack_from('<I', hdr, 2)[0]
        return hdr[6], hdr[7], struct.unpack_from('<I', hdr, 8)[0], self.recvExact(length)

    def cmd(self, ctype, body=b''):
        rid = self.reqId
        self.reqId += 1
        self.sock.sendall(struct.pack('<BBIIB', STX, API_VERSION, len(body), rid, ctype) + body)
        while True:
            rtype, err, r, b = self.packet()
            if r == rid:
                if err != 0:
                    raise RuntimeError('monitor command %02X failed (%02X)' % (ctype, err))
                return b
            self.events.append((rtype, b))

    def mem(self, start, end, bank):
        b = self.cmd(CMD_MEM_GET, struct.pack('<BHHBH', 0, start, end, 0, bank))
        return b[2:2 + struct.unpack_from('<H', b, 0)[0]]

    def poke(self, start, data, bank):
        self.cmd(CMD_MEM_SET, struct.pack('<BHHBH', 0, start, start + len(data) - 1, 0, bank) + bytes(data))

    def stopAt(self, addr):
        b = self.cmd(CMD_CP_SET, struct.pack('<HHBBBB', addr, addr, 1, 1, OP_EXEC, 0))
        return struct.unpack_from('<I', b, 0)[0]

    def unstop(self, num):
        self.cmd(CMD_CP_DELETE, struct.pack('<I', num))

    def pc(self):
        b = self.cmd(CMD_REGS_GET, b'\x00')
        off = 2
        for _ in range(struct.unpack_from('<H', b, 0)[0]):
            if b[off + 1] == REG_PC:
                return struct.unpack_from('<H', b, off + 2)[0]
            off += b[off] + 1
        raise RuntimeError('no PC in the register dump')

    def setPc(self, value):
        self.cmd(CMD_REGS_SET, struct.pack('<BH', 0, 1) + struct.pack('<BBH', 3, REG_PC, value))

    def banks(self):
        b = self.cmd(CMD_BANKS)
        out = {}
        off = 2
        for _ in range(struct.unpack_from('<H', b, 0)[0]):
            nameLen = b[off + 3]
            out[b[off + 4:off + 4 + nameLen].decode()] = struct.unpack_from('<H', b, off + 1)[0]
            off += b[off] + 1
        return out

    def go(self):
        self.events = []
        self.cmd(CMD_EXIT)

    def halt(self):
        self.cmd(CMD_PING)

    def waitStop(self, timeout):
        self.sock.settimeout(timeout)
        try:
            while True:
                if self.events:
                    rtype, b = self.events.pop(0)
                else:
                    rtype, _err, _rid, b = self.packet()
                if rtype == EV_STOPPED:
                    return struct.unpack_from('<H', b, 0)[0]
                if rtype == EV_JAM:
                    raise RuntimeError('the C64 jammed at $%04X' % struct.unpack_from('<H', b, 0)[0])
        finally:
            self.sock.settimeout(None)

    def keys(self, text):
        t = text.encode('latin-1')
        self.cmd(CMD_KEYS, bytes([len(t)]) + t)

    def snapshotSave(self, path):
        p = path.encode()
        self.cmd(CMD_DUMP, bytes([0, 0, len(p)]) + p)

    def snapshotLoad(self, path):
        p = path.encode()
        self.cmd(CMD_UNDUMP, bytes([len(p)]) + p)


class Vice:
    """x64sc on a private Xvfb display, driven through its monitor."""

    def __init__(self, autostart=None):
        exe = os.environ.get('VICE') or shutil.which('x64sc')
        if not exe or not shutil.which('Xvfb'):
            raise FileNotFoundError('VICE (x64sc) or Xvfb not found')
        self.procs = []
        display = self.freeDisplay()
        # Each in a session of its own, so close() can take down the
        # timeout guard and what it runs together.
        self.procs.append(subprocess.Popen(['timeout', '-s', 'KILL', str(LIFE_S), 'Xvfb', display, '-screen', '0', '640x480x24'],
                                           stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, start_new_session=True))
        os.makedirs(os.path.join(CACHE, 'home'), exist_ok=True)
        env = dict(os.environ, DISPLAY=display, XDG_CONFIG_HOME=os.path.join(CACHE, 'home'), XDG_CACHE_HOME=os.path.join(CACHE, 'home'))
        port = self.freePort()
        args = ['timeout', '-s', 'KILL', str(LIFE_S), exe, '-binarymonitor', '-binarymonitoraddress', 'ip4://127.0.0.1:%d' % port,
                '-sounddev', 'dummy', '-warp', '+confirmonexit']
        if os.environ.get('VICE_DATA'):
            args += ['-directory', os.environ['VICE_DATA']]
        if autostart:
            args += ['-autostart-warp', '-autostart', autostart]
        self.procs.append(subprocess.Popen(args, stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL, env=env, start_new_session=True))
        try:
            self.mon = Monitor(port)
        except BaseException:
            self.close()
            raise

    @staticmethod
    def freePort():
        s = socket.socket()
        s.bind(('127.0.0.1', 0))
        port = s.getsockname()[1]
        s.close()
        return port

    @staticmethod
    def freeDisplay():
        n = 100 + os.getpid() % 800
        while os.path.exists('/tmp/.X%d-lock' % n):
            n += 1
        return ':%d' % n

    def close(self):
        try:
            self.mon.cmd(CMD_QUIT)
        except Exception:
            pass
        for p in reversed(self.procs):
            try:
                os.killpg(p.pid, signal.SIGKILL)
            except ProcessLookupError:
                pass
            p.wait()


def fnv(data):
    h = 2166136261
    for b in data:
        h = ((h ^ b) * 16777619) & 0xFFFFFFFF
    return h


def demoOffset(mon, ram):
    zp = mon.mem(0x003F, 0x0040, ram)
    return ((zp[1] << 8) | zp[0]) - DEMO_BUFFER


def stateLine(mon, tick, ram, io, raw):
    """stsim's printState, from the C64's memory."""
    scr  = mon.mem(0x0400, 0x07E7, ram)
    chs  = mon.mem(0x2800, 0x2FFF, ram)
    vic  = mon.mem(0xD020, 0xD021, io)
    col  = bytes(b & 0x0F for b in mon.mem(0xD800, 0xDBE7, io))
    g    = mon.mem(0x7140, 0x71FF, ram)
    p    = mon.mem(0x7D5A, 0x7D9F, ram)
    G    = lambda a: g[a - 0x7140]
    P    = lambda a: p[a - 0x7D5A]
    en   = G(0x7196)
    msb  = G(0x718D)
    if raw is not None:
        raw.write(scr + col + chs)
    # A hidden sprite's registers print as dashes, as stsim's do.
    sprx = ''
    for k in range(8):
        x = G(0x71A7 + 2 * k) | (0x100 if (msb >> k) & 1 else 0)
        sprx += '%03X%02X' % (x, G(0x71A8 + 2 * k)) if (en >> k) & 1 else '-----'
    if en & 2:
        spr1 = '"ptr1":"%02X","s1":"%02X%02X.%02X",' % (G(0x7198), G(0x7186), G(0x7176), G(0x717E))
    else:
        spr1 = '"ptr1":"--","s1":"--",'
    s  = ('{"tick":%u,"x":"%02X%02X.%02X","y":"%02X.%02X","vx":"%04X","vy":"%04X",'
          '"pad":%u,"st":%u,"ph":%u,"in":"%02X","dir":"%02X","ptr0":"%02X",' % (
          tick, P(0x7D63), P(0x7D62), P(0x7D61), P(0x7D65), P(0x7D64),
          (G(0x714D) << 8) | G(0x714C), (G(0x714F) << 8) | G(0x714E),
          G(0x7150), G(0x7163), G(0x7164), G(0x7169), G(0x716A), G(0x7197)))
    s += spr1
    s += ('"en":"%02X","t4740":"%02X","off":%u,"rng":"%02X%02X","fuel":%u,"slots":%u,"asi":%u,"score":"' % (
          en, mon.mem(0x4740, 0x4740, ram)[0], demoOffset(mon, ram), G(0x7171), G(0x7172), G(0x71C9), G(0x715C), P(0x7D8E)))
    s += ''.join('%02X' % scr[0x3C2 + k] for k in range(7))
    s += '","fare":"' + ''.join('%02X' % scr[0x3E0 + k] for k in range(7))
    s += '","fuelc":"' + ''.join('%02X' % scr[0x3A6 + k] for k in range(12))
    s += '","sprx":"' + sprx
    s += '","scr":"%08X","col":"%08X","chr":"%08X","bd":"%02X","bg":"%02X"}' % (
         fnv(scr), fnv(col), fnv(chs), vic[0] & 0x0F, vic[1] & 0x0F)
    return s


def snapshot():
    """The cached VICE snapshot of the first attract demo picking its
    screen, made from the original disk on first use."""
    path = os.path.join(CACHE, 'attract.vsf')
    if os.path.exists(path):
        return path
    if not os.path.exists(D64):
        raise FileNotFoundError('no disk image at %s (STAXI_D64)' % D64)
    os.makedirs(CACHE, exist_ok=True)
    vice = Vice(autostart=D64)
    try:
        mon = vice.mon
        cp = mon.stopAt(INTRO_KEY_TEST)
        mon.go()
        mon.waitStop(BOOT_WAIT_S)
        mon.unstop(cp)
        mon.setPc(INTRO_GO)
        mon.go()
        time.sleep(INTRO_WAIT_S)
        mon.halt()
        for _ in range(TRAINER_ROWS):
            mon.keys(KEY_DOWN)
            mon.go()
            time.sleep(KEY_WAIT_S)
            mon.halt()
        mon.keys(KEY_RETURN)
        cp = mon.stopAt(DEMO_PICK)
        mon.go()
        mon.waitStop(BOOT_WAIT_S)
        mon.unstop(cp)
        mon.snapshotSave(path + '.part')
    finally:
        vice.close()
    os.replace(path + '.part', path)
    return path


def ride(letter, ticks, demobuf, out, raw):
    vice = Vice()
    try:
        mon = vice.mon
        mon.snapshotLoad(snapshot())
        banks = mon.banks()
        ram   = banks['ram']
        io    = banks['io']
        if mon.pc() != DEMO_PICK:
            raise RuntimeError('the snapshot is not at $%04X' % DEMO_PICK)
        mon.poke(DEMO_SLOTS, bytes([ord(letter)] * 4), ram)
        cp = mon.stopAt(DEMO_LOAD)
        mon.go()
        mon.waitStop(STOP_WAIT_S)
        mon.unstop(cp)
        mon.poke(DEMO_BUFFER, demobuf, ram)
        for a in CARRIED:
            mon.poke(a, b'\x00', ram)
        mon.poke(SCREENS_PLAYED, bytes([ord(letter) - ord('A')]), ram)
        mon.setPc(DEMO_LOADED)
        for a in (TICK_TOP, CAB_LEFT, WRECK_DONE, PLAYER_OUT):
            mon.stopAt(a)
        tick   = 0
        exited = False
        while tick < ticks:
            mon.go()
            pc = mon.waitStop(STOP_WAIT_S)
            if pc == TICK_TOP:
                if demoOffset(mon, ram) >= LAST_PAIR:
                    out.write('{"tick":%u,"end":"stream"}\n' % (tick - 1))
                    break
                out.write(stateLine(mon, tick, ram, io, raw) + '\n')
                tick += 1
            elif pc == CAB_LEFT:
                exited = True
            elif pc == WRECK_DONE:
                if exited:
                    out.write('{"tick":%u,"end":%d}\n' % (tick - 1, END_EXIT))
                    break
                mon.setPc(LOSE_CAB)
            elif pc == PLAYER_OUT:
                out.write('{"tick":%u,"end":%d}\n' % (tick - 1, END_OUT))
                break
            else:
                raise RuntimeError('unexpected stop at $%04X' % pc)
    finally:
        vice.close()


def main():
    if len(sys.argv) == 2 and sys.argv[1] == '--prepare':
        try:
            snapshot()
        except FileNotFoundError as e:
            print('c64trace: %s' % e, file=sys.stderr)
            return 2
        return 0
    if len(sys.argv) < 4:
        print('usage: c64trace.py <level letter> <ticks> <demobuf.json> [rawOut] | --prepare', file=sys.stderr)
        return 2
    letter  = sys.argv[1]
    ticks   = int(sys.argv[2], 0)
    demobuf = bytes.fromhex(json.load(open(sys.argv[3]))['demobuf'])
    raw     = open(sys.argv[4], 'wb') if len(sys.argv) > 4 else None
    try:
        ride(letter, ticks, demobuf, sys.stdout, raw)
    except FileNotFoundError as e:
        print('c64trace: %s' % e, file=sys.stderr)
        return 2
    finally:
        if raw is not None:
            raw.close()
    return 0


if __name__ == '__main__':
    sys.exit(main())