#!/usr/bin/env python3
# testHighRates.py - what the SwiftLink driver does at the line rates a C64 Ultimate's SwiftLink
# emulation might hand it, and what it does when the two ends do not agree on the rate.
#
#   python3 testHighRates.py <disk.d64> <rateA> <rateB> [windowSeconds] [cleanWindows] [tag]
#
# It reuses testRealtime.py wholesale - the socket null modem, the "warp only the loader" boot, the
# hot-key rate selection, the stopwatch check that the emulators really ran at 1 MHz, and the
# transmitter-deadlock probe.  What it adds is aimed at one question the earlier reports could not
# answer: at a rate the 20-byte receive ring cannot keep up with, does the link make progress,
# degrade and RECOVER, or WEDGE?
#
#   * rateA and rateB are selected independently, so 38400 against 300 - a device that ignores the
#     baud bits talking to a peer that does not - is the same script with different arguments;
#   * a window is measured three ways rather than once, because the instruments are not free:
#       w1..wN  clean, no tracepoints at all, so the byte counts are the honest throughput;
#       wSync   with tracepoints on beginByteSyncPhase ($E805) and on the instruction that
#               completes a sync ($E855, "sta ackPending", reached only when the peer's $FF has
#               arrived).  Those two count degradations and recoveries directly, and they are low
#               volume - at most a few per second - where a tracepoint on the error counter is not;
#       wErr    a short window with the error tracepoints from rate1200Report.md: $E73F (every
#               character the receiver threw away) and the conditional $E505 if x > $13 (the subset
#               thrown away because the 20-byte ring was already full).  At 38400 that can print
#               thousands of lines a second, which is why it is short and separate;
#   * uartRxCount ($E42D) is then sampled repeatedly with the emulator stopped, to see how full the
#     raw receive ring actually gets, alongside connectionPhase, isLinkActive, txCharActive and the
#     command shadow.  Those reads DO stop the emulator, so they are outside every measured window.
#
# The rig's own limits, which the report has to repeat: VICE's ACIA is byte level, so two ends at
# different rates do not garble each other the way real hardware would - the fast sender simply
# fills the slow receiver's socket buffer.
import os
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 (BITS_PER_BYTE, CONTROL_RATES, LOGS, NTSC_HZ, SHOTS, command,
                          decodeControl, isTxDeadlocked, measureWindowOnce, pauseAll,
                          probeTxDeadlock, readRange, resumeAll, selectRate, showState, snapshot,
                          waitForPhase)

RING_SIZE = 20                          # pushUartRxRing $E505 refuses a 21st character
SYNC_ENTER = 0xE805                     # beginByteSyncPhase: the link threw the connection away
SYNC_DONE = 0xE855                      # sta ackPending: the peer answered $FF, phase 3 is next
# The module's own jump table.  Sixteen bytes that cannot be anything else, which is the only honest
# "the opponent module has landed" test - see test1200AutoBaud.waitForModule and rate1200Report.md
# section 6.  $E5BE and $E013/$E014 both read plausible values out of the previous overlay's wreckage
# while the 1541 is still fetching track 34, and a run that trusts them measures a machine that never
# opened its link at all.
MODULE_JUMP_TABLE = [0x4C, 0x11, 0xE1, 0x4C, 0x9A, 0xE2, 0x4C, 0xA9, 0xE0, 0x4C, 0xF6, 0xE0, 0x4C,
                     0xCB, 0xE3, 0x60]


def countHits(text, address):
    return text.count(f"exec {address:04x}")


def measure(relay, sessions, seconds, label, traces, runTag, key):
    # One window, with whichever tracepoints this window is meant to carry, and a full state read of
    # both machines at each edge.  Nothing is retried and nothing is discarded: a window that
    # carried no bytes is the answer to the question, not a failed measurement.
    setTraces(sessions, traces)
    before = [snapshot(session, tag) for session, tag in sessions]
    for state in before:
        showState(state)
    result = measureWindowOnce(relay, sessions, seconds, label)
    after = [snapshot(session, tag) for session, tag in sessions]
    for state in after:
        showState(state)
    setTraces(sessions, [])
    result["statesBefore"] = before
    result["states"] = after
    result["healthy"] = all(value >= 0.9 for value in result["speed"].values())
    result["deadlocked"] = {state["tag"]: isTxDeadlocked(state) for state in after}
    result["traceByTag"] = {tag: {"syncEnter$E805": countHits(text, SYNC_ENTER),
                                  "syncDone$E855": countHits(text, SYNC_DONE),
                                  "linkError$E73F": countHits(text, 0xE73F),
                                  "rxRingFull$E505": countHits(text, 0xE505)}
                            for tag, text in result["tracedBy"].items()}
    print(f"  transmitter deadlock at this edge: {result['deadlocked']}", flush=True)
    print(f"  tracepoint hits by machine: {result['traceByTag']}", flush=True)
    open(f"{LOGS}/highRate.{runTag}.{key}.samples.txt", "w").write(
        "\n".join(f"{t}\t{c0}\t{c1}" for t, c0, c1 in result["samples"]) + "\n")
    if result["traced"].strip():
        open(f"{LOGS}/highRate.{runTag}.{key}.trace.txt", "w").write(result["traced"])
    return result


def probePendingCount(relay, sessions, seconds=20):
    # The direct test of the diagnosis, the way realtimeReport.md tested the 2400-baud transmitter
    # deadlock by writing $00 over txCharActive: on a machine that has gone silent with an empty
    # transmit ring and a non-zero uartPendingCount, write $00 over $E0A5 from the monitor - nothing
    # else - and see whether its bytes come back.  If they do, the byte the layers above wait on is
    # what wedged.  If they do not, the wedge is above the driver and no byte in the driver fixes it,
    # which is just as much of an answer.
    results = []
    for stage, addresses in (("uartPendingCount $E0A5", [0xE0A5]),
                             ("txCharActive $E5BB too", [0xE0A5, 0xE5BB])):
        before = [len(relay.log[0]), len(relay.log[1])]
        for session, tag in sessions:
            pauseAll([(session, tag)])
            for address in addresses:
                command(session, f"> {address:04x} 00")
            resumeAll([(session, tag)])
        time.sleep(seconds)
        moved = [len(relay.log[index]) - before[index] for index in (0, 1)]
        print(f"  {seconds} s after clearing {stage}: conn0 {moved[0]} bytes, conn1 {moved[1]} bytes",
              flush=True)
        results.append({"cleared": stage, "seconds": seconds, "bytes": moved})
    return results


def wedgeDump(session, tag, times=6):
    # Everything needed to say WHERE a silent machine is stuck rather than only that it is.  The
    # first 38400-baud run left machine B reading $DE01/$DE02/$DE03 as $FF/$00/$00 while machine A
    # read them correctly, which is what the monitor shows when the CPU it is reading through has the
    # I/O area banked out ($01 bit 2 clear) - so $01 and the program counter are the two bytes that
    # decide whether a wedge is in the driver or in the game above it.
    print(f"\n---- wedge dump [{tag}] ----", flush=True)
    for index in range(times):
        pauseAll([(session, tag)])
        command(session, "registers")
        command(session, "m 0000 0001")          # the bank register
        command(session, "m e039 e03a")          # the disk-load suspend handshake
        command(session, "m e42d e42f")          # uartRxCount and both ring indices
        command(session, "m e414 e416")          # uartTxCount and both transmit ring indices
        command(session, f"m {ACIA_BASE:04x} {ACIA_BASE + 3:04x}")
        command(session, "m 0ba5 0ba5")          # the game's solo/modem flag
        resumeAll([(session, tag)])
        time.sleep(1.0)


def moduleResident(session, tag, tries=60):
    # Wait for the opponent module itself, not for a byte inside it.  At true C64 speed the 1541
    # takes the best part of a minute to fetch track 34, and the fire press that starts the load does
    # not always land, so it is offered again every sixth poll while the options menu is still up.
    for attempt in range(tries):
        pauseAll([(session, tag)])
        got = readRange(session, 0xE000, 16)
        resumeAll([(session, tag)])
        if got == MODULE_JUMP_TABLE:
            print(f"[{tag}] the module is resident: the $E000 jump table matches", flush=True)
            return True
        print(f"[{tag}] waiting for the module: $E000 = "
              f"{' '.join('??' if value is None else f'{value:02X}' for value in got)}", flush=True)
        if attempt % 6 == 5:
            session.focus()
            session.hold("KP_0", 300)
        time.sleep(3)
    print(f"[{tag}] the module never became resident", flush=True)
    return False


def perSecond(samples):
    # The window's samples are cumulative relay byte counts; what matters for "did it keep making
    # progress" is the per-second delta and, above all, whether any second carried nothing.
    out = []
    for index in range(1, len(samples)):
        span = samples[index][0] - samples[index - 1][0]
        if span <= 0:
            continue
        out.append((round(span, 2),
                    (samples[index][1] - samples[index - 1][1]) / span,
                    (samples[index][2] - samples[index - 1][2]) / span))
    return out


def ringSample(session, tag):
    # How full the raw receive ring is right now.  This stops the emulator, so it never runs inside
    # a measured window.
    pauseAll([(session, tag)])
    ring = readRange(session, 0xE42D, 3)        # uartRxCount, read index, write index
    link = readRange(session, 0xE03B, 13)       # isLinkActive .. linkErrorCount
    acia = readRange(session, 0xE5BB, 4)        # txCharActive, statusSave, rxByte, command shadow
    txq = readRange(session, 0xE414, 1)         # uartTxCount
    resumeAll([(session, tag)])
    return {"tag": tag, "uartRxCount": ring[0], "uartTxCount": txq[0],
            "isLinkActive": link[0], "connectionPhase": link[5], "linkErrorCount": link[12],
            "txCharActive": acia[0], "aciaStatusSave": acia[1], "aciaCommandShadow": acia[3]}


def sampleRings(sessions, count, gap, runTag):
    # Repeated sampling, both machines, alternating.  The point is the distribution: the maximum the
    # ring ever reached, how often it was at the 20-byte limit, and whether connectionPhase moved
    # between 2 and 3 (degrade and recover) or sat still (wedge).
    rows = []
    for index in range(count):
        for session, tag in sessions:
            row = ringSample(session, tag)
            row["n"] = index
            rows.append(row)
            print(f"  [ring {index:02d}] [{tag}] uartRxCount={row['uartRxCount']} "
                  f"uartTxCount={row['uartTxCount']} phase={row['connectionPhase']} "
                  f"active=${(row['isLinkActive'] or 0):02X} err={row['linkErrorCount']} "
                  f"tx=${(row['txCharActive'] or 0):02X} cmd=${(row['aciaCommandShadow'] or 0):02X}",
                  flush=True)
        time.sleep(gap)
    with open(f"{LOGS}/highRate.{runTag}.ring.txt", "w") as handle:
        for row in rows:
            handle.write(f"{row['n']}\t{row['tag']}\t{row['uartRxCount']}\t{row['uartTxCount']}\t"
                         f"{row['connectionPhase']}\t{row['linkErrorCount']}\t"
                         f"{row['txCharActive']}\t{row['aciaCommandShadow']}\n")
    for _, tag in sessions:
        mine = [row for row in rows if row["tag"] == tag]
        depths = [row["uartRxCount"] for row in mine if row["uartRxCount"] is not None]
        phases = sorted({row["connectionPhase"] for row in mine})
        full = sum(1 for value in depths if value >= RING_SIZE)
        print(f"  [{tag}] uartRxCount over {len(depths)} samples: min {min(depths)}, "
              f"max {max(depths)}, mean {sum(depths) / len(depths):.1f}, at the {RING_SIZE}-byte "
              f"limit {full} times; connectionPhase values seen {phases}", flush=True)
    return rows


def setTraces(sessions, addresses):
    pauseAll(sessions)
    for session, tag in sessions:
        command(session, "del")
        for text in addresses:
            command(session, f"trace exec {text}")
        command(session, "break")
    resumeAll(sessions)


def summarise(name, run):
    print(f"\n--- {name} ---", flush=True)
    print(f"  {run['wall']:.2f} s wall, atFullSpeed={run['healthy']}", flush=True)
    for tag in run["cycles"]:
        print(f"   [{tag}] {run['cycles'][tag]} cycles = {run['cycles'][tag] / NTSC_HZ:.2f} "
              f"emulated s = {100.0 * run['speed'][tag]:.1f}% of real time", flush=True)
    for i in (0, 1):
        count = run["counts"][i]
        print(f"   conn{i} -> peer {count} bytes = {count / run['wall']:.2f} B/s = "
              f"{count / run['wall'] * BITS_PER_BYTE:.0f} bit/s of 8N1 line time; "
              f"wire {run['wire'][i]}", flush=True)
    buckets = perSecond(run["samples"])
    for i, name2 in ((1, "conn0"), (2, "conn1")):
        values = [row[i] for row in buckets]
        zeros = sum(1 for value in values if value < 0.5)
        print(f"   {name2} per second: min {min(values):.1f}, max {max(values):.1f}, "
              f"seconds carrying nothing: {zeros} of {len(values)}", flush=True)
    print(f"   tracepoints: {run['traceByTag']}", flush=True)
    print(f"   deadlock at the closing edge: {run['deadlocked']}", flush=True)
    for state in run["states"]:
        print(f"   [{state['tag']}] control {decodeControl(state['aciaControl'])}", flush=True)
        showState(state)


def main():
    disk = os.path.abspath(sys.argv[1])
    rateA = int(sys.argv[2])
    rateB = int(sys.argv[3]) if len(sys.argv) > 3 else int(sys.argv[2])
    seconds = float(sys.argv[4]) if len(sys.argv) > 4 else 75.0
    cleanWindows = int(sys.argv[5]) if len(sys.argv) > 5 else 1
    runTag = sys.argv[6] if len(sys.argv) > 6 else "h"
    os.makedirs(SHOTS, exist_ok=True)
    os.makedirs(LOGS, exist_ok=True)
    relay = NullModemRelay()
    print(f"relay on 127.0.0.1:{relay.port}; A wants {rateA} baud, B wants {rateB} baud", flush=True)
    # The host serial device's own baud setting is raised above every rate under test so that it can
    # never be the thing doing the limiting; the ACIA control register is what this test is about.
    args = aciaArgs(rsDevAddress=f"127.0.0.1:{relay.port}", baud=38400)
    a = ViceSession(disk, f"{SCRATCH}/hr.{runTag}.a.vice.log", args, label="A", warp=False)
    b = ViceSession(disk, f"{SCRATCH}/hr.{runTag}.b.vice.log", args, label="B", warp=False)
    sessions = [(a, "A"), (b, "B")]
    runs = []
    try:
        a.connect()
        b.connect()
        for session, tag in sessions:
            command(session, "warp on")
        a.bootPastLoader(waitSecs=900)
        b.bootPastLoader(waitSecs=900)
        pauseAll(sessions)
        for session, tag in sessions:
            command(session, "warp off")
            print(f"[{tag}] {command(session, 'warp').strip()}", flush=True)
        resumeAll(sessions)
        a.findWindow()
        b.findWindow()
        a.focus()
        b.focus()
        time.sleep(25)
        a.shot(f"{SHOTS}/hr{runTag}01aMenu.png")
        b.shot(f"{SHOTS}/hr{runTag}01bMenu.png")
        pickModemOpponent(a, "A", "a")
        pickModemOpponent(b, "B", "o")
        # Each machine gets its own rate, chosen at the "PRESS A OR O" prompt before the link is
        # ever opened, which is the order README.md tells the player to use.  baudIndex comes off
        # the disk as 0, so 300 needs no key - but pressing C= + 3 anyway proves the module is
        # resident and programs the chip at the same point in the sequence as every other rate.
        for (session, tag), rate in zip(sessions, (rateA, rateB)):
            if not moduleResident(session, tag):
                raise SystemExit(f"{tag}: the opponent module never loaded")
            selectRate(session, tag, rate)
        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)
        a.shot(f"{SHOTS}/hr{runTag}02aLinked.png")
        b.shot(f"{SHOTS}/hr{runTag}02bLinked.png")
        for state in [snapshot(session, tag) for session, tag in sessions]:
            print(f"  [{state['tag']}] control {decodeControl(state['aciaControl'])}", flush=True)
            showState(state)

        for index in range(cleanWindows):
            runs.append((f"clean w{index + 1}",
                         measure(relay, sessions, seconds,
                                 f"{rateA}/{rateB} baud, clean [{index + 1}/{cleanWindows}]",
                                 [], runTag, f"clean{index + 1}")))
        a.shot(f"{SHOTS}/hr{runTag}03aAfterClean.png")
        b.shot(f"{SHOTS}/hr{runTag}03bAfterClean.png")

        runs.append(("sync/recover", measure(relay, sessions, seconds,
                                             f"{rateA}/{rateB} baud, sync tracepoints",
                                             [f"${SYNC_ENTER:04x}", f"${SYNC_DONE:04x}"],
                                             runTag, "sync")))
        runs.append(("errors", measure(relay, sessions, min(seconds, 20.0),
                                       f"{rateA}/{rateB} baud, error tracepoints",
                                       ["$e73f", "$e505 if x > $13"], runTag, "err")))
        a.shot(f"{SHOTS}/hr{runTag}04aAfterTraced.png")
        b.shot(f"{SHOTS}/hr{runTag}04bAfterTraced.png")

        print("\n---- uartRxCount $E42D sampled repeatedly ----", flush=True)
        rows = sampleRings(sessions, 24, 0.5, runTag)
        probe = probeTxDeadlock(relay, sessions, runs[-1][1]["states"])
        print(f"deadlock probe: {probe}", flush=True)
        # A direction that carried nothing for a whole window is the wedge case, and it is worth more
        # than the byte count: dump where that machine actually is.
        # Both ends are dumped rather than the silent one, because which relay connection belongs to
        # which emulator is an inference (conn0 is whichever opened its ACIA first) and the machine
        # that is still talking is half the evidence.
        pending = None
        if any(0 in run["counts"] for _, run in runs):
            for session, tag in sessions:
                wedgeDump(session, tag)
            pending = probePendingCount(relay, sessions)
            for session, tag in sessions:
                wedgeDump(session, tag, times=2)
        a.shot(f"{SHOTS}/hr{runTag}05aEnd.png")
        b.shot(f"{SHOTS}/hr{runTag}05bEnd.png")

        print("\n================ SUMMARY ================", flush=True)
        print(f"A asked for {rateA} baud, B asked for {rateB} baud", flush=True)
        for name, run in runs:
            summarise(name, run)
        finalPhases = sorted({row["connectionPhase"] for row in rows})
        print(f"\nconnectionPhase values seen across the ring sampling: {finalPhases}", flush=True)
        print(f"deadlock probe: {probe}", flush=True)
        print(f"uartPendingCount probe: {pending}", flush=True)
    finally:
        a.close()
        b.close()
        time.sleep(1)
        relay.stop()
        for i in (0, 1):
            print(f"[relay] conn{i} total {len(relay.log[i])} bytes, "
                  f"tail {bytes(relay.log[i][-60:])!r}", flush=True)


if __name__ == "__main__":
    main()