#!/usr/bin/env python3
# testWedge.py - why the SwiftLink driver wedges at 38400 baud, and only there.
#
#   python3 testWedge.py <disk.d64> <rate> [tag] [watchSeconds]
#
# highRateReport.md characterised the wedge (four boots, four wedges, one CPU JAM at $0007) but did
# not explain it.  Three candidate causes were on the table and this script is built to tell them
# apart by measurement rather than by reading the source:
#
#   1. NMI re-entrancy.  commNmiHandler reads the ACIA status register early, which releases /IRQ,
#      so a character arriving before the handler finishes raises a fresh edge and re-enters it.
#      DIRECT TEST: trace the handler's entry ($E685) and both of its exits ($E6E9 RTI, $E67F the
#      chain JMP) and pair them with a stack.  An entry arriving while the depth is already non-zero
#      IS a re-entry - no inference needed.  Every VICE trace line carries SP and the free-running
#      cycle counter, so the same capture gives the nesting depth, the handler's duration in cycles
#      and the interrupt rate.
#   2. CPU starvation.  If the handler costs more cycles per second than the machine has, the main
#      line stops, the once-per-frame link service stops with it, and everything above the driver
#      waits for ever.  DIRECT TEST: the same capture measures the duty cycle exactly (sum of
#      exit-entry over the window), and a tracepoint on $E3CB - the $E00C jump-table entry the raster
#      IRQ calls once per frame - counts how often the game's per-frame link service actually runs.
#      60/s means the game is alive; 0/s means it is not.
#   3. Something above the driver.  Ruled in or out by the same $E3CB counter plus a PC/SP sampling
#      profiler and, after the wedge, a dump of the whole stack page.
#
# Nothing here changes the driver.  The only write to the machine is baudIndex $E055 before the link
# is opened, because the 38400 hot key was deliberately removed from baudEntryTable - the driver's
# own loadBaudParameters still programs the control register from that index, so the path under test
# is the driver's, not the monitor's.
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, CONTROL_RATES, LOGS, NTSC_HZ, SHOTS, command, decodeControl,
                          flush, pauseAll, readRange, resumeAll, snapshot, showState, stopwatch,
                          waitForPhase)
from testHighRates import moduleResident

NMI_ENTRY = 0xE685                      # commNmiHandler
NMI_EXIT_RTI = 0xE6E9                   # the normal exit
NMI_EXIT_CHAIN = 0xE67F                 # 'not mine' - jmp (nmiChainVector)
FRAME_SERVICE = 0xE3CB                  # pollCarrierState, the $E00C entry the raster IRQ calls
PENDING_DEC = 0xE6C9                    # dec uartPendingCount, inside the NMI transmit half
PENDING_INC = 0xE71E                    # inc uartPendingCount, in queueAndKickTx
PENDING_RESYNC = 0xE611                 # sta uartPendingCount - startNextTxChar's once-a-frame repair
TRACE_LINE = re.compile(r"^\.C:([0-9a-f]{4})\s+.*?SP:([0-9a-f]{2})\s+\S+\s+(\d+)\s*$", re.M)
REG_LINE = re.compile(r"^\.;([0-9a-f]{4}) ([0-9a-f]{2}) ([0-9a-f]{2}) ([0-9a-f]{2}) ([0-9a-f]{2})",
                      re.M)


def captureTrace(session, tag, addresses, seconds, maxBytes, path):
    # Tracepoints print without stopping the emulator, but printing thousands of lines a second over
    # the monitor socket does slow the host down.  That costs wall-clock speed, not emulated timing:
    # VICE's ACIA is clocked in emulated cycles, so the relationship between a character time and a
    # handler's cycle count - the only thing measured here - is unchanged.  Every reading below is
    # in emulated cycles for that reason, and the window's length is taken from the emulator's own
    # free-running counter, not from the wall clock.
    pauseAll([(session, tag)])
    command(session, "del")
    for address in addresses:
        command(session, f"trace exec {address:04x}")
    start = stopwatch(session)
    flush(session, 0.1, 1.0)
    session.sock.sendall(b"x\n")
    text = []
    total = 0
    session.sock.settimeout(0.5)
    deadline = time.time() + seconds
    while time.time() < deadline and total < maxBytes:
        try:
            chunk = session.sock.recv(1 << 16)
        except OSError:
            continue
        if not chunk:
            break
        text.append(chunk.decode("latin-1"))
        total += len(chunk)
    pauseAll([(session, tag)])
    stop = stopwatch(session)
    command(session, "del")
    resumeAll([(session, tag)])
    body = "".join(text)
    open(path, "w").write(body)
    print(f"[{tag}] captured {total} bytes of trace into {path}; "
          f"emulated cycles {start} -> {stop}", flush=True)
    return body, (start, stop)


def analyseTrace(body, span, tag, label):
    # Pair entries with exits on a stack.  An entry seen while the depth is already non-zero is a
    # re-entrant NMI, which is the whole question.
    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] == NMI_ENTRY]
    frames = [event for event in events if event[0] == FRAME_SERVICE]
    decs = [event for event in events if event[0] == PENDING_DEC]
    incs = [event for event in events if event[0] == PENDING_INC]
    resyncs = [event for event in events if event[0] == PENDING_RESYNC]
    depth = 0
    maxDepth = 0
    nested = 0
    stack = []
    durations = []
    unmatched = 0
    for pc, sp, cycle in events:
        if pc == NMI_ENTRY:
            if depth > 0:
                nested += 1
            depth += 1
            maxDepth = max(maxDepth, depth)
            stack.append(cycle)
        elif pc in (NMI_EXIT_RTI, NMI_EXIT_CHAIN):
            if stack:
                durations.append(cycle - stack.pop())
                depth -= 1
            else:
                unmatched += 1
    cycles = (span[1] - span[0]) if span[0] is not None and span[1] is not None else 0
    seconds = cycles / NTSC_HZ if cycles else 0.0
    busy = sum(durations)
    entrySps = [sp for pc, sp, _ in events if pc == NMI_ENTRY]
    gaps = [entries[i][2] - entries[i - 1][2] for i in range(1, len(entries))]
    result = {"tag": tag, "label": label, "cycles": cycles, "seconds": round(seconds, 4),
              "nmiEntries": len(entries), "nmiExits": len(durations), "unmatchedExits": unmatched,
              "reEntries": nested, "maxDepth": maxDepth,
              "frameServiceHits": len(frames),
              "pendingDec$E6C9": len(decs), "pendingInc$E71E": len(incs),
              "pendingResync$E611": len(resyncs),
              "frameServicePerSecond": round(len(frames) / seconds, 1) if seconds else None,
              "nmiPerSecond": round(len(entries) / seconds, 1) if seconds else None,
              "dutyCycle": round(busy / cycles, 4) if cycles else None,
              "durationMin": min(durations) if durations else None,
              "durationMax": max(durations) if durations else None,
              "durationMean": round(sum(durations) / len(durations), 1) if durations else None,
              "entrySpMin": min(entrySps) if entrySps else None,
              "entrySpMax": max(entrySps) if entrySps else None,
              "gapMin": min(gaps) if gaps else None,
              "gapMedian": sorted(gaps)[len(gaps) // 2] if gaps else None}
    print(f"\n  ---- NMI trace [{tag}] {label} ----", flush=True)
    for key, value in result.items():
        print(f"    {key}: {value}", flush=True)
    if durations:
        buckets = {}
        for value in durations:
            buckets[value // 32 * 32] = buckets.get(value // 32 * 32, 0) + 1
        print(f"    duration histogram (32-cycle buckets): "
              f"{dict(sorted(buckets.items()))}", flush=True)
    if entrySps:
        buckets = {}
        for value in entrySps:
            buckets[value // 16 * 16] = buckets.get(value // 16 * 16, 0) + 1
        print(f"    entry SP histogram (16-byte buckets): "
              f"{dict(sorted(buckets.items(), reverse=True))}", flush=True)
    return result


def registers(session, tag):
    for _ in range(4):
        out = command(session, "registers")
        match = REG_LINE.search(out)
        if match:
            return {"pc": int(match.group(1), 16), "a": int(match.group(2), 16),
                    "x": int(match.group(3), 16), "y": int(match.group(4), 16),
                    "sp": int(match.group(5), 16)}
    return None


def profile(sessions, count, label):
    # A sampling profiler made of monitor pauses.  Where the PC is when the machine is stopped, over
    # many stops, is the share of the CPU each region is getting; SP at the same moments says how
    # deep the interrupt nesting is.  Pausing does not change emulated timing - the emulator is
    # frozen, cartridge included - it only costs wall-clock time.
    out = {}
    for session, tag in sessions:
        rows = []
        for _ in range(count):
            pauseAll([(session, tag)])
            row = registers(session, tag)
            resumeAll([(session, tag)])
            if row:
                rows.append(row)
        buckets = {}
        for row in rows:
            buckets[region(row["pc"])] = buckets.get(region(row["pc"]), 0) + 1
        sps = sorted(row["sp"] for row in rows)
        print(f"\n  ---- PC/SP profile [{tag}] {label}, {len(rows)} samples ----", flush=True)
        for name, hits in sorted(buckets.items(), key=lambda item: -item[1]):
            print(f"    {name}: {hits} ({100.0 * hits / max(1, len(rows)):.1f}%)", flush=True)
        print(f"    SP: min ${min(sps):02X} max ${max(sps):02X} "
              f"median ${sps[len(sps) // 2]:02X}", flush=True)
        out[tag] = {"buckets": buckets, "samples": len(rows),
                    "spMin": min(sps), "spMax": max(sps), "spMedian": sps[len(sps) // 2],
                    "rows": rows}
    return out


def region(pc):
    # Which layer the CPU was in.  The helpers the NMI calls live outside $E685-$E732, so they are
    # named separately rather than lumped in with the rest of the module.
    if 0xE679 <= pc <= 0xE6F1:
        return "commNmiHandler $E679-$E6F1"
    for low, high, name in ((0xE4EA, 0xE51A, "UART ring helpers $E4EA-$E51A"),
                            (0xE551, 0xE56E, "ACIA accessors $E551-$E56E"),
                            (0xE408, 0xE40E, "aciaSetCommandIdle $E408"),
                            (0xE04C, 0xE054, "clearTxCharActive $E04F"),
                            (0xE73F, 0xE743, "countLinkError $E73F"),
                            (0xE3BC, 0xE3BE, "reportLinkError $E3BC")):
        if low <= pc <= high:
            return name
    if 0xE000 <= pc <= 0xEFFF:
        return "module, elsewhere"
    if pc < 0x0200:
        return "PAGE 0/1 - the CPU is off the rails"
    return "game code"


def bigDump(session, tag, title, times=4):
    print(f"\n======== {title} [{tag}] ========", flush=True)
    rows = []
    for index in range(times):
        pauseAll([(session, tag)])
        row = registers(session, tag)
        rows.append(row)
        command(session, "m 0000 0001")
        command(session, "m e039 e047")          # suspend handshake .. linkErrorCount
        command(session, "m e0a4 e0a8")          # frameInFlightFlag, uartPendingCount, packetInFlight
        command(session, "m e414 e416")          # uartTxCount + indices
        command(session, "m e42d e42f")          # uartRxCount + indices
        command(session, "m e5bb e5bf")          # txCharActive .. aciaPageLatch
        command(session, "m e055 e056")          # baudIndex, live control byte
        command(session, f"m {ACIA_BASE:04x} {ACIA_BASE + 3:04x}")
        command(session, "m fffa fffb")          # the NMI vector itself
        command(session, "m e031 e033")          # the chain vector and the CIA mask shadow
        command(session, "m 0ba5 0ba5")
        command(session, "stopwatch")
        if index == 0:
            command(session, "m 0100 017f", timeout=20)
            command(session, "m 0180 01ff", timeout=20)
        resumeAll([(session, tag)])
        time.sleep(0.8)
    print(f"  [{tag}] register samples: {rows}", flush=True)
    return rows


def watch(relay, sessions, seconds, gap=3.0):
    # Poll until both directions have gone quiet, or the time runs out.  The state read pauses the
    # emulator, which is fine here: this is a watch for a wedge, not a throughput measurement.
    start = time.time()
    marks = [len(relay.log[0]), len(relay.log[1])]
    quiet = 0
    history = []
    while time.time() - start < seconds:
        time.sleep(gap)
        now = [len(relay.log[0]), len(relay.log[1])]
        moved = [now[i] - marks[i] for i in (0, 1)]
        marks = now
        states = [snapshot(session, tag) for session, tag in sessions]
        for state in states:
            showState(state)
        elapsed = round(time.time() - start, 1)
        print(f"  [t+{elapsed}s] bytes this poll: conn0 {moved[0]}, conn1 {moved[1]}", flush=True)
        history.append({"t": elapsed, "moved": moved,
                        "states": [{k: v for k, v in state.items()} for state in states]})
        if moved[0] < 3 and moved[1] < 3:
            quiet += 1
            if quiet >= 3:
                print(f"  WEDGED: both directions carried almost nothing for "
                      f"{quiet * gap:.0f} s, at t+{elapsed}s", flush=True)
                return True, elapsed, history
        else:
            quiet = 0
    return False, round(time.time() - start, 1), history


def main():
    disk = os.path.abspath(sys.argv[1])
    rate = int(sys.argv[2])
    tag = sys.argv[3] if len(sys.argv) > 3 else "w"
    watchSeconds = float(sys.argv[4]) if len(sys.argv) > 4 else 150.0
    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}/wedge.{tag}.a.vice.log", args, label="A", warp=False)
    b = ViceSession(disk, f"{SCRATCH}/wedge.{tag}.b.vice.log", args, label="B", warp=False)
    sessions = [(a, "A"), (b, "B")]
    report = {"tag": tag, "rate": rate}
    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")
            # 38400's hot key was removed from baudEntryTable, so the rate is selected by writing
            # baudIndex and letting the driver's own loadBaudParameters program the chip when the
            # link opens.  Every other rate is written the same way, so the runs are comparable.
            pauseAll([(session, sessionTag)])
            command(session, f"> e055 {BAUD_INDEX[rate]:02x}")
            resumeAll([(session, sessionTag)])
        answerModemPrompts(a, "A", "a")
        answerModemPrompts(b, "B", "o")
        ok, phases = waitForPhase(sessions, 3, 180)
        print(f"both in the packet phase: {ok} ({phases})", flush=True)
        report["reachedPhase3"] = ok
        states = [snapshot(session, sessionTag) for session, sessionTag in sessions]
        for state in states:
            print(f"  [{state['tag']}] control {decodeControl(state['aciaControl'])}", flush=True)
            showState(state)
        report["control"] = {state["tag"]: state["aciaControl"] for state in states}
        a.shot(f"{SHOTS}/wg{tag}01linked.png")

        addresses = [NMI_ENTRY, NMI_EXIT_RTI, NMI_EXIT_CHAIN, FRAME_SERVICE,
                     PENDING_DEC, PENDING_INC, PENDING_RESYNC]
        body, span = captureTrace(a, "A", addresses, 12.0, 6 << 20,
                                  f"{LOGS}/wedge.{tag}.A.nmi.trace.txt")
        report["traceLive"] = analyseTrace(body, span, "A", f"{rate} baud, live link")
        report["profileLive"] = profile(sessions, 60, f"{rate} baud, live link")

        wedged, when, history = watch(relay, sessions, watchSeconds)
        report["wedged"] = wedged
        report["wedgeAt"] = when
        report["history"] = history
        a.shot(f"{SHOTS}/wg{tag}02end.png")
        b.shot(f"{SHOTS}/wg{tag}02endB.png")

        report["dump"] = {}
        for session, sessionTag in sessions:
            report["dump"][sessionTag] = bigDump(session, sessionTag,
                                                 "state after the watch window")
        report["profileAfter"] = profile(sessions, 60, "after the watch window")
        for session, sessionTag in sessions:
            body, span = captureTrace(session, sessionTag, addresses, 8.0, 4 << 20,
                                      f"{LOGS}/wedge.{tag}.{sessionTag}.after.trace.txt")
            report[f"traceAfter{sessionTag}"] = analyseTrace(body, span, sessionTag,
                                                            "after the watch window")
        print("\n================ SUMMARY ================", flush=True)
        for key in ("tag", "rate", "reachedPhase3", "control", "wedged", "wedgeAt"):
            print(f"  {key}: {report.get(key)}", flush=True)
        for key in ("traceLive", "traceAfterA", "traceAfterB"):
            print(f"  {key}: {report.get(key)}", flush=True)
        for key in ("profileLive", "profileAfter"):
            value = report.get(key) or {}
            for sessionTag, row in value.items():
                print(f"  {key}[{sessionTag}]: samples={row['samples']} "
                      f"spMin=${row['spMin']:02X} spMedian=${row['spMedian']:02X} "
                      f"buckets={row['buckets']}", 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, "
                  f"tail {bytes(relay.log[index][-60:])!r}", flush=True)


if __name__ == "__main__":
    main()