#!/usr/bin/env python3
# testStack.py - catch the 38400-baud wedge while it is happening, rather than after.
#
#   python3 testStack.py <disk.d64> <rate> [tag] [watchSeconds]
#
# testWedge.py established what a wedged machine looks like afterwards: the game running with the I
# flag set so its raster IRQ never runs again, and the whole of page 1 covered in commNmiHandler
# frames.  That is wreckage, not a mechanism.  This script watches for the mechanism.
#
# The 6502 does not mask NMI, so a fresh /NMI edge re-enters commNmiHandler at the next instruction
# boundary, and every re-entry costs 6 bytes of stack: the CPU's own PCH/PCL/P plus the handler's A,
# X and saved $01.  So the stack pointer AT THE HANDLER'S FIRST INSTRUCTION is a direct read-out of
# how many NMIs are nested at that moment - and VICE prints SP and the free-running cycle counter on
# every tracepoint hit.
#
#   trace exec e685 if sp < $a8    - prints nothing while nesting is shallow (the main line runs with
#                                    SP around $BD-$E5, so an entry below $A8 is at least three deep)
#                                    and turns into a staircase the moment the handler stops keeping
#                                    up.  Cheap: no output at all until it matters.
#   break exec e685 if sp < $20    - stops the emulator once the stack has nearly wrapped, so the
#                                    stack page, the registers and the ACIA can be read with the
#                                    runaway still in progress instead of after it has destroyed
#                                    itself.
#
# Both machines carry both checkpoints, and the raw monitor output of each is saved, so the staircase
# is timestamped in the emulator's own cycles.
import os
import re
import sys
import time

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from viceHarness import ViceSession, aciaArgs, ACIA_BASE, SCRATCH
from testTwoMachines import NullModemRelay, pickModemOpponent, answerModemPrompts
from testRealtime import (BAUD_INDEX, LOGS, NTSC_HZ, SHOTS, command, decodeControl, flush,
                          pauseAll, readRange, resumeAll, snapshot, showState, stopwatch,
                          waitForPhase)
from testHighRates import moduleResident
from testWedge import TRACE_LINE, registers

DEEP_TRACE_SP = 0xB4                    # one NMI frame below the deepest main line ever sampled
BREAK_SP = 0x20                         # nearly a whole page of stack consumed
# A healthy machine never pushes below about $01A0: the deepest stack byte either emulator touched in
# a 200-second 38400-baud run that did not wedge was $01AE.  So a STORE anywhere in $0100-$0180 means
# the stack has descended thirteen NMI frames further than anything healthy does, and it is the
# cheapest possible trigger - it costs nothing at all until the runaway starts.
DEEP_TRACE_STORE = (0x0100, 0x0180)
BREAK_STORE = (0x0100, 0x0140)


def collect(sessions, seconds, maxBytes, runTag):
    # Read both monitor sockets raw, at the same time, and stop early if either machine's breakpoint
    # fires.  A resumed emulator prints nothing, so any prompt in the stream is the break.
    buffers = {tag: [] for _, tag in sessions}
    totals = {tag: 0 for _, tag in sessions}
    broke = None
    for session, _ in sessions:
        session.sock.settimeout(0.2)
    deadline = time.time() + seconds
    while time.time() < deadline and broke is None:
        for session, tag in sessions:
            if totals[tag] >= maxBytes:
                continue
            try:
                chunk = session.sock.recv(1 << 16)
            except OSError:
                continue
            if not chunk:
                continue
            text = chunk.decode("latin-1")
            buffers[tag].append(text)
            totals[tag] += len(chunk)
            if "Stop on" in text:
                broke = tag
                print(f"\n*** [{tag}] BREAKPOINT: the handler was entered with SP below "
                      f"${BREAK_SP:02X} ***", flush=True)
    out = {}
    for _, tag in sessions:
        body = "".join(buffers[tag])
        path = f"{LOGS}/stack.{runTag}.{tag}.deep.txt"
        open(path, "w").write(body)
        out[tag] = body
        print(f"[{tag}] {totals[tag]} bytes of deep-nesting trace -> {path}", flush=True)
    return out, broke


def staircase(body, tag):
    # Every line is one entry into commNmiHandler with the stack already at least three frames deep.
    # Consecutive lines whose SP falls by exactly 6 are consecutive nesting levels; the cycle gap
    # between them is how long one pass of the handler managed before the next character arrived.
    events = [(int(pc, 16), int(sp, 16), int(cycle)) for pc, sp, cycle in TRACE_LINE.findall(body)]
    entries = [event for event in events if event[0] == 0xE685]
    if not entries:
        print(f"  [{tag}] no entry to commNmiHandler was ever seen with SP < ${DEEP_TRACE_SP:02X}",
              flush=True)
        return {"tag": tag, "deepEntries": 0}
    sps = [sp for _, sp, _ in entries]
    gaps = [entries[i][2] - entries[i - 1][2] for i in range(1, len(entries))]
    descending = sum(1 for i in range(1, len(entries))
                     if (entries[i - 1][1] - entries[i][1]) % 256 == 6)
    print(f"\n  ---- deep-nesting staircase [{tag}] ----", flush=True)
    print(f"    entries below SP ${DEEP_TRACE_SP:02X}: {len(entries)}", flush=True)
    print(f"    SP first ${sps[0]:02X}, last ${sps[-1]:02X}, lowest ${min(sps):02X}", flush=True)
    print(f"    steps where SP fell by exactly 6 (one more NMI frame): {descending} of "
          f"{len(entries) - 1}", flush=True)
    if gaps:
        print(f"    cycles between consecutive entries: min {min(gaps)}, median "
              f"{sorted(gaps)[len(gaps) // 2]}, max {max(gaps)}", flush=True)
    span = entries[-1][2] - entries[0][2]
    print(f"    the whole staircase took {span} cycles = {span / NTSC_HZ * 1000:.2f} ms", flush=True)
    print(f"    first 24: {[(f'${sp:02X}', cycle) for _, sp, cycle in entries[:24]]}", flush=True)
    print(f"    last 8:   {[(f'${sp:02X}', cycle) for _, sp, cycle in entries[-8:]]}", flush=True)
    return {"tag": tag, "deepEntries": len(entries), "spFirst": sps[0], "spLast": sps[-1],
            "spLowest": min(sps), "stepsOfSix": descending, "spanCycles": span,
            "gapMedian": sorted(gaps)[len(gaps) // 2] if gaps else None,
            "gapMin": min(gaps) if gaps else None}


def dumpAtBreak(session, tag):
    print(f"\n======== the machine, stopped inside the runaway [{tag}] ========", flush=True)
    flush(session, 0.3, 3)
    row = registers(session, tag)
    print(f"  [{tag}] registers: {row}", flush=True)
    command(session, "d e685 e6f1", timeout=20)
    command(session, "m 0100 017f", timeout=20)
    command(session, "m 0180 01ff", timeout=20)
    command(session, "m 0000 0001")
    command(session, f"m {ACIA_BASE:04x} {ACIA_BASE + 3:04x}")
    command(session, "m e5bb e5bf")
    command(session, "m e42d e42f")
    command(session, "m e414 e416")
    command(session, "m e0a4 e0a6")
    command(session, "m 0038 0038")              # irqInProgress, the game's raster-IRQ guard
    command(session, "m fffa fffb")
    command(session, "stopwatch")
    return row


def main():
    disk = os.path.abspath(sys.argv[1])
    rate = int(sys.argv[2])
    tag = sys.argv[3] if len(sys.argv) > 3 else "s"
    watchSeconds = float(sys.argv[4]) if len(sys.argv) > 4 else 200.0
    # traceOnly: arm the tracepoints only, never a breakpoint, and start reading the moment the link
    # is open.  st384b caught a machine with SP already at $12, but it caught it with a BREAKPOINT,
    # which stops the emulator - and once stopped it stays stopped, so the descent itself was never
    # printed.  Tracepoints print and let the machine run, so this mode records the staircase as it
    # happens and never interferes.
    traceOnly = len(sys.argv) > 5 and sys.argv[5] == "traceonly"
    os.makedirs(SHOTS, exist_ok=True)
    os.makedirs(LOGS, exist_ok=True)
    relay = NullModemRelay()
    print(f"relay on 127.0.0.1:{relay.port}; both machines want {rate} baud", flush=True)
    args = aciaArgs(rsDevAddress=f"127.0.0.1:{relay.port}", baud=38400)
    a = ViceSession(disk, f"{SCRATCH}/stack.{tag}.a.vice.log", args, label="A", warp=False)
    b = ViceSession(disk, f"{SCRATCH}/stack.{tag}.b.vice.log", args, label="B", warp=False)
    sessions = [(a, "A"), (b, "B")]
    try:
        a.connect()
        b.connect()
        for session, _ in sessions:
            command(session, "warp on")
        a.bootPastLoader(waitSecs=900)
        b.bootPastLoader(waitSecs=900)
        pauseAll(sessions)
        for session, sessionTag in sessions:
            command(session, "warp off")
            print(f"[{sessionTag}] {command(session, 'warp').strip()}", flush=True)
        resumeAll(sessions)
        a.findWindow()
        b.findWindow()
        a.focus()
        b.focus()
        time.sleep(25)
        pickModemOpponent(a, "A", "a")
        pickModemOpponent(b, "B", "o")
        for session, sessionTag in sessions:
            if not moduleResident(session, sessionTag):
                raise SystemExit(f"{sessionTag}: the opponent module never loaded")
            pauseAll([(session, sessionTag)])
            command(session, f"> e055 {BAUD_INDEX[rate]:02x}")
            # Arm both checkpoints before the link is ever opened, so nothing that happens on the
            # way up can be missed.
            command(session, "del")
            command(session, f"trace exec e685 if sp < ${DEEP_TRACE_SP:02x}")
            if not traceOnly:
                command(session, f"trace store {DEEP_TRACE_STORE[0]:04x} "
                                 f"{DEEP_TRACE_STORE[1]:04x}")
                command(session, f"break exec e685 if sp < ${BREAK_SP:02x}")
                command(session, f"break store {BREAK_STORE[0]:04x} {BREAK_STORE[1]:04x}")
            resumeAll([(session, sessionTag)])
        answerModemPrompts(a, "A", "a")
        answerModemPrompts(b, "B", "o")
        if not traceOnly:
            ok, phases = waitForPhase(sessions, 3, 180)
            print(f"both in the packet phase: {ok} ({phases})", flush=True)
            for state in [snapshot(session, sessionTag) for session, sessionTag in sessions]:
                print(f"  [{state['tag']}] control {decodeControl(state['aciaControl'])}",
                      flush=True)
                showState(state)
            # waitForPhase and snapshot both stop and restart the emulators, so re-arm and hand the
            # sockets over to the raw collector with nothing else talking on them.
            for session, sessionTag in sessions:
                pauseAll([(session, sessionTag)])
                command(session, "del")
                command(session, f"trace exec e685 if sp < ${DEEP_TRACE_SP:02x}")
                command(session, f"trace store {DEEP_TRACE_STORE[0]:04x} "
                                 f"{DEEP_TRACE_STORE[1]:04x}")
                command(session, f"break exec e685 if sp < ${BREAK_SP:02x}")
                command(session, f"break store {BREAK_STORE[0]:04x} {BREAK_STORE[1]:04x}")
                resumeAll([(session, sessionTag)])
        # Bracket the watch with both machines' own cycle counters, so the run says for itself
        # whether it was at true C64 speed.
        pauseAll(sessions)
        startCycles = {sessionTag: stopwatch(session) for session, sessionTag in sessions}
        resumeAll(sessions)
        t0 = time.time()
        bodies, broke = collect(sessions, watchSeconds, 8 << 20, tag)
        wall = time.time() - t0
        if broke is None:
            pauseAll(sessions)
            endCycles = {sessionTag: stopwatch(session) for session, sessionTag in sessions}
            resumeAll(sessions)
            for sessionTag in startCycles:
                span = endCycles[sessionTag] - startCycles[sessionTag]
                print(f"[{sessionTag}] watch {wall:.2f} s wall, {span} emulated cycles = "
                      f"{span / NTSC_HZ:.2f} emulated s = "
                      f"{100.0 * span / NTSC_HZ / wall:.1f}% of real time", flush=True)
        summary = {}
        for _, sessionTag in sessions:
            summary[sessionTag] = staircase(bodies[sessionTag], sessionTag)
        if broke:
            session = dict((sessionTag, s) for s, sessionTag in sessions)[broke]
            dumpAtBreak(session, broke)
            print(f"\n  letting [{broke}] run on from the break", flush=True)
            command(session, "del")
            resumeAll([(session, broke)])
            time.sleep(10)
            pauseAll([(session, broke)])
            print(f"  [{broke}] ten seconds later: {registers(session, broke)}", flush=True)
            command(session, "m 0038 0038")
            command(session, "m e42d e42f")
            command(session, "m e0a5 e0a5")
            resumeAll([(session, broke)])
        for session, sessionTag in sessions:
            pauseAll([(session, sessionTag)])
            command(session, "del")
            row = registers(session, sessionTag)
            print(f"  [{sessionTag}] final registers: {row}", flush=True)
            command(session, "m 0038 0038")
            command(session, "m 0100 017f", timeout=20)
            command(session, "m 0180 01ff", timeout=20)
            resumeAll([(session, sessionTag)])
        a.shot(f"{SHOTS}/st{tag}end.png")
        b.shot(f"{SHOTS}/st{tag}endB.png")
        print("\n================ SUMMARY ================", flush=True)
        print(f"  rate {rate}, tag {tag}, broke on {broke}", flush=True)
        for sessionTag, row in summary.items():
            print(f"  {sessionTag}: {row}", flush=True)
    finally:
        a.close()
        b.close()
        time.sleep(1)
        relay.stop()
        for index in (0, 1):
            print(f"[relay] conn{index} total {len(relay.log[index])} bytes", flush=True)


if __name__ == "__main__":
    main()