#!/usr/bin/env python3
# testHarden.py - verify the two hardenings the driver has just been given, on either strap:
#
#   1. the page latch (aciaPageLatch $E5BF).  aciaDetect probes $DE00 once, on a cold start, and
#      never again: every later call - every link open and every baud hot key - sees a non-zero latch
#      and returns in five bytes.  So on a $DF00 machine nothing is ever written into $DE00-$DE03
#      after the single cold-start probe, and on either strap the only thing a rate change writes to
#      the control register is the rate itself, never a probe pattern.
#   2. the 38400 retirement.  baudEntryTable's last entry keeps its control byte $1F but its hot-key
#      code is now $00, which the search at $E344 can never match because it only runs on codes with
#      bit 7 set.  C= + 8 must therefore change nothing and be swallowed like any other unrecognised
#      Commodore-key combination.
#
#   python3 testHarden.py <disk.d64> <tag>
#   SWIFTLINK_ACIA_BASE=0xDF00 python3 testHarden.py <disk.d64> <tag>
#
# Two machines are booted and cross connected through the socket null modem, so every measurement
# below is taken on a link that has reached connection phase 3 - the game's own packet protocol -
# rather than on a machine talking to itself.  Both machines run at true C64 speed; only the loader
# is warped.
#
# Everything is measured with VICE *tracepoints*, not watchpoints.  A watchpoint stops the emulator
# at every hit, which is fine for a cold-start log (strapReport.md counted 800 of them that way) but
# useless while keys are being pressed at true speed.  A tracepoint prints the accessing instruction
# and its A register and lets the machine run on, so an ordered access log costs nothing and the hot
# keys land normally.
import os
import re
import subprocess
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
from testRealtime import (command, decodeControl, flush, hotKey, pauseAll, readRange, resumeAll,
                          showState, snapshot, waitForPhase, BAUD_HOTKEY, BAUD_INDEX,
                          CONTROL_RATES)
from testHighRates import MODULE_JUMP_TABLE

SHOTS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "shots")
LOGS = os.path.join(os.path.dirname(os.path.abspath(__file__)), "testLogs")

OTHER_BASE = 0xDF00 if ACIA_BASE == 0xDE00 else 0xDE00
PROBE_PATTERNS = (0x1E, 0x15)           # aciaProbe's two control-register patterns

# The six one-instruction ACIA accessors; the operand high byte is the driver's whole record of where
# the hardware is, and aciaPageLatch $E5BF is the readable seventh copy of it.
ACCESSORS = [(0xE554, "aciaCmdWrite   sta $xx02  command"),
             (0xE558, "aciaStatusRead lda $xx01  status"),
             (0xE55F, "aciaDataRead   lda $xx00  receive"),
             (0xE563, "aciaDataWrite  sta $xx00  transmit"),
             (0xE567, "aciaCtrlWrite  sta $xx03  control"),
             (0xE56B, "aciaCtrlRead   lda $xx03  control read-back / the probe")]

TRACE_HEAD = re.compile(r"#(\d+)\s+\(Trace\s+(\w+)\s+([0-9a-f]{4})\)")
DISASM = re.compile(r"\.C:([0-9a-f]{4})\s+((?:[0-9A-F]{2} )+)\s*(\S+)\s+(\S*)\s*-\s*A:([0-9A-Fa-f]{2})")

# The order the hot keys are pressed in.  38400 is in the middle rather than at the end so that the
# rates either side of it prove the keyboard was still working when it was pressed.
RATE_SEQUENCE = [1200, 2400, 4800, 38400, 9600, 19200, 300]


def armTraces(sessions, ranges):
    # ranges: list of (first, last).  Both directions on each, deleted first so a phase's log
    # contains only what that phase armed.
    pauseAll(sessions)
    for session, tag in sessions:
        command(session, "del")
        for first, last in ranges:
            command(session, f"trace store ${first:04x} ${last:04x}")
            command(session, f"trace load ${first:04x} ${last:04x}")
    resumeAll(sessions)


def collect(sessions, seconds, sink):
    # Read whatever the tracepoints printed, without stopping anything.  sink is a dict of tag ->
    # list of strings.
    deadline = time.time() + seconds
    while time.time() < deadline:
        for session, tag in sessions:
            sink[tag].append(flush(session, 0.25, 0.4))


def describe(blob):
    # Ordered "who touched what, with what" out of the tracepoint stream.
    events = []
    lines = blob.splitlines()
    for index, line in enumerate(lines):
        head = TRACE_HEAD.search(line)
        if not head:
            continue
        for follow in lines[index:index + 3]:
            hit = DISASM.search(follow)
            if hit:
                events.append({"kind": head.group(2), "addr": int(head.group(3), 16),
                               "pc": int(hit.group(1), 16), "op": hit.group(3),
                               "operand": hit.group(4), "a": int(hit.group(5), 16)})
                break
    return events


def pageOf(addr):
    return addr & 0xFF00


def readAccessors(session, tag):
    pauseAll([(session, tag)])
    values = {}
    for addr, name in ACCESSORS:
        values[addr] = (name, readRange(session, addr + 2, 1)[0])
    latch = readRange(session, 0xE5BF, 1)[0]
    resumeAll([(session, tag)])
    return values, latch


def showAccessors(values, latch, tag, note):
    print(f"--- [{tag}] the six ACIA accessors ({note}) ---", flush=True)
    pages = set()
    for addr, name in ACCESSORS:
        _, value = values[addr]
        pages.add(value)
        print(f"    ${addr:04X}  {name:42s} operand high byte at ${addr + 2:04X} = "
              f"${'??' if value is None else '%02X' % value}", flush=True)
    print(f"    aciaPageLatch $E5BF = ${'??' if latch is None else '%02X' % latch}   "
          f"all six agree: {len(pages) == 1}   they point at page "
          f"${(pages.pop() if len(pages) == 1 else 0):02X}00", flush=True)
    return latch


def summariseEvents(events, label):
    counts = {}
    for event in events:
        key = (event["addr"], event["kind"])
        counts[key] = counts.get(key, 0) + 1
    byPage = {}
    for (addr, kind), n in counts.items():
        byPage[pageOf(addr)] = byPage.get(pageOf(addr), 0) + n
    print(f"=== {label}: {len(events)} accesses ===", flush=True)
    for (addr, kind) in sorted(counts):
        print(f"    ${addr:04X} {kind:5s} x {counts[(addr, kind)]}", flush=True)
    for page in sorted(byPage):
        print(f"    page ${page:04X}: {byPage[page]} accesses", flush=True)
    return counts, byPage


def moduleResident(session, tag, tries=60):
    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
        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 openLink(sessions, keys, sink, tries=24):
    # Answer "PRESS A OR O" and then "PRESS SPACE", with the cold-detect traces already armed and
    # nothing sent to either monitor, so the log of the probe itself is clean.
    for attempt in range(tries):
        for (session, tag), key in zip(sessions, keys):
            session.focus()
            session.hold(key, 250)
            session.hold("space", 250)
        collect(sessions, 3, sink)
        if all(describe("".join(sink[tag])) for _, tag in sessions):
            print(f"both machines touched an ACIA register after {attempt + 1} attempt(s)",
                  flush=True)
            collect(sessions, 4, sink)
            return True
    return False


def pressRate(sessions, rate, sink, tries=4):
    # Press one speed hot key on both machines and report what each driver did with it.  38400 is in
    # the table by index but has no key code any more, so C= + 8 is expected to do nothing at all -
    # which is why "landed" is reported rather than asserted.
    key = BAUD_HOTKEY[rate]
    want = [code for code, value in CONTROL_RATES.items() if value == rate][0]
    before = [snapshot(session, tag) for session, tag in sessions]
    for state in before:
        showState(state)
    for attempt in range(tries):
        for session, tag in sessions:
            hotKey(session, key)
        collect(sessions, 2, sink)
        after = [snapshot(session, tag) for session, tag in sessions]
        landed = all(state["aciaControl"] == want and state["baudIndex"] == BAUD_INDEX[rate]
                     for state in after)
        if landed or rate == 38400:
            break
        print(f"  C= + {key} did not land on both machines - retrying", flush=True)
    for state in after:
        showState(state)
        print(f"  [{state['tag']}] control {decodeControl(state['aciaControl'])}", flush=True)
    return {"rate": rate, "key": key, "wantControl": want, "before": before, "after": after,
            "landed": landed}


def swallowProof(session, tag):
    # $E38A is the 'lda #$FF' that swallows an unrecognised Commodore-key combination, and it is the
    # only way out of the hot-key layer for C= + 8 now that the baud walk cannot match it.  Arm the
    # breakpoint before the key goes down: entering the monitor clears VICE's emulated key matrix.
    pauseAll([(session, tag)])
    command(session, "del")
    command(session, "break e38a")
    resumeAll([(session, tag)])
    session.focus()
    subprocess.run(["xdotool", "keydown", "Tab"], env=session.env, check=False)
    time.sleep(0.2)
    subprocess.run(["xdotool", "keydown", "8"], env=session.env, check=False)
    hit = flush(session, 0.6, 12)
    print(f"[{tag}] break at $E38A while C= + 8 is held:\n{hit[-600:]}", flush=True)
    code = readRange(session, 0xE046, 1)[0]
    baud = readRange(session, 0xE055, 2)
    control = readRange(session, ACIA_BASE + 3, 1)[0]
    print(f"[{tag}] at the swallow: lastHotkeyCode $E046 = "
          f"{'??' if code is None else '$%02X' % code}, baudIndex $E055 = {baud[0]}, "
          f"bitPeriodLo $E056 = {'??' if baud[1] is None else '$%02X' % baud[1]}, "
          f"${ACIA_BASE + 3:04X} = {'??' if control is None else '$%02X' % control}", flush=True)
    command(session, "del")
    resumeAll([(session, tag)])
    subprocess.run(["xdotool", "keyup", "8"], env=session.env, check=False)
    subprocess.run(["xdotool", "keyup", "Tab"], env=session.env, check=False)
    time.sleep(0.4)
    return {"hit": hit, "reachedSwallow": "e38a" in hit.lower(), "lastHotkeyCode": code,
            "baudIndex": baud[0], "bitPeriodLo": baud[1], "control": control}


def main():
    disk = os.path.abspath(sys.argv[1])
    tag = sys.argv[2] if len(sys.argv) > 2 else "harden"
    os.makedirs(SHOTS, exist_ok=True)
    os.makedirs(LOGS, exist_ok=True)
    relay = NullModemRelay()
    print(f"relay on 127.0.0.1:{relay.port}", flush=True)
    print(f"acia base for this run: ${ACIA_BASE:04X}; the other page is ${OTHER_BASE:04X}",
          flush=True)
    args = aciaArgs(rsDevAddress=f"127.0.0.1:{relay.port}", baud=2400)
    a = ViceSession(disk, f"{SCRATCH}/hd.{tag}.a.vice.log", args, label="A", warp=False)
    b = ViceSession(disk, f"{SCRATCH}/hd.{tag}.b.vice.log", args, label="B", warp=False)
    sessions = [(a, "A"), (b, "B")]
    sink = {"A": [], "B": []}
    results = {}
    try:
        a.connect()
        b.connect()
        for session, name in sessions:
            command(session, "warp on")
            command(session, "sidefx off")
        a.bootPastLoader(waitSecs=900)
        b.bootPastLoader(waitSecs=900)
        pauseAll(sessions)
        for session, name in sessions:
            command(session, "warp off")
            print(f"[{name}] {command(session, 'warp').strip()}", flush=True)
        resumeAll(sessions)
        a.findWindow()
        b.findWindow()
        a.focus()
        b.focus()
        time.sleep(25)
        a.shot(f"{SHOTS}/hd{tag}01aMenu.png")
        b.shot(f"{SHOTS}/hd{tag}01bMenu.png")
        pickModemOpponent(a, "A", "a")
        pickModemOpponent(b, "B", "o")
        for session, name in sessions:
            if not moduleResident(session, name):
                raise SystemExit(f"{name}: the opponent module never loaded")

        # ---- phase 1: the cold detect, with both pages watched -------------------------------
        armTraces(sessions, [(0xDE00, 0xDE03), (0xDF00, 0xDF03)])
        ok = openLink(sessions, ["a", "o"], sink)
        print(f"link opened on both machines: {ok}", flush=True)
        coldEvents = {name: describe("".join(sink[name])) for _, name in sessions}
        results["cold"] = {}
        for _, name in sessions:
            print(f"=== [{name}] the first 12 cartridge-page accesses, in order ===", flush=True)
            for event in coldEvents[name][:12]:
                print(f"    {event['kind']:5s} ${event['addr']:04X}  {event['op']} "
                      f"{event['operand']} at ${event['pc']:04X}  A=${event['a']:02X}", flush=True)
            counts, byPage = summariseEvents(coldEvents[name], f"[{name}] cold detect + link open")
            results["cold"][name] = {"events": len(coldEvents[name]), "byPage": byPage,
                                     "counts": {f"{k[0]:04X}:{k[1]}": v for k, v in counts.items()}}
            open(f"{LOGS}/harden.{tag}.{name}.cold.txt", "w").write("".join(sink[name]))

        okPhase, phases = waitForPhase(sessions, 3, 240)
        print(f"both machines in the packet phase: {okPhase} ({phases})", flush=True)
        results["reachedPhase3"] = okPhase
        a.shot(f"{SHOTS}/hd{tag}02aLinked.png")
        b.shot(f"{SHOTS}/hd{tag}02bLinked.png")
        for session, name in sessions:
            values, latch = readAccessors(session, name)
            showAccessors(values, latch, name, "after the cold detect")
            results.setdefault("accessors", {})[name] = {
                "operands": {f"{addr:04X}": values[addr][1] for addr, _ in ACCESSORS},
                "latch": latch}

        # ---- phase 2: the hot keys, with the other page and the control register watched ------
        sink = {"A": [], "B": []}
        armTraces(sessions, [(OTHER_BASE, OTHER_BASE + 3), (ACIA_BASE + 3, ACIA_BASE + 3)])
        results["rates"] = []
        for rate in RATE_SEQUENCE:
            print(f"\n---- C= + {BAUD_HOTKEY[rate]} ({rate} baud) on a live link ----", flush=True)
            results["rates"].append(pressRate(sessions, rate, sink))
        collect(sessions, 3, sink)
        hotEvents = {name: describe("".join(sink[name])) for _, name in sessions}
        results["hotkeyPhase"] = {}
        for _, name in sessions:
            print(f"=== [{name}] every access to ${OTHER_BASE:04X}-${OTHER_BASE + 3:04X} and to the "
                  f"control register ${ACIA_BASE + 3:04X} during the whole hot-key sweep ===",
                  flush=True)
            for event in hotEvents[name]:
                print(f"    {event['kind']:5s} ${event['addr']:04X}  {event['op']} "
                      f"{event['operand']} at ${event['pc']:04X}  A=${event['a']:02X}", flush=True)
            counts, byPage = summariseEvents(hotEvents[name], f"[{name}] hot-key sweep")
            writes = [e for e in hotEvents[name]
                      if e["kind"] == "store" and e["addr"] == ACIA_BASE + 3]
            # $1E and $15 are both probe patterns AND real rate codes (19200 and 300), so a single
            # value proves nothing.  What a re-probe would look like is the pair $1E, $15 written
            # back to back in front of the rate the player asked for, i.e. three writes per hot key
            # instead of one.  So the honest measurement is the ordered list of everything written
            # to the control register during the sweep, against the rates that were asked for.
            written = " ".join(f"${e['a']:02X}" for e in writes)
            pairs = sum(1 for i in range(len(writes) - 1)
                        if writes[i]["a"] == PROBE_PATTERNS[0]
                        and writes[i + 1]["a"] == PROBE_PATTERNS[1])
            print(f"    control-register writes, in order: {written}", flush=True)
            print(f"    back-to-back $1E,$15 pairs (the probe's signature): {pairs}", flush=True)
            results["hotkeyPhase"][name] = {
                "otherPageAccesses": byPage.get(OTHER_BASE, 0),
                "controlWrites": [e["a"] for e in writes],
                "probePairs": pairs,
                "events": len(hotEvents[name])}
            open(f"{LOGS}/harden.{tag}.{name}.hotkeys.txt", "w").write("".join(sink[name]))

        # ---- phase 3: a census of every register access, both pages, 20 s --------------------
        sink = {"A": [], "B": []}
        armTraces(sessions, [(0xDE00, 0xDE03), (0xDF00, 0xDF03)])
        collect(sessions, 20, sink)
        results["census"] = {}
        for _, name in sessions:
            events = describe("".join(sink[name]))
            counts, byPage = summariseEvents(events, f"[{name}] 20 s census on a live link")
            results["census"][name] = {"byPage": byPage, "total": len(events)}
            open(f"{LOGS}/harden.{tag}.{name}.census.txt", "w").write("".join(sink[name]))
        pauseAll(sessions)
        for session, name in sessions:
            command(session, "del")
        resumeAll(sessions)

        # ---- phase 4: C= + 8 really does reach the swallow ------------------------------------
        a.shot(f"{SHOTS}/hd{tag}03aBeforeC8.png")
        results["swallow"] = swallowProof(a, "A")
        time.sleep(3)
        a.shot(f"{SHOTS}/hd{tag}04aAfterC8.png")
        b.shot(f"{SHOTS}/hd{tag}04bAfterC8.png")
        for session, name in sessions:
            values, latch = readAccessors(session, name)
            showAccessors(values, latch, name, "at the end of the run")

        print("\n================ SUMMARY ================", flush=True)
        print(f"strap: ${ACIA_BASE:04X}   other page: ${OTHER_BASE:04X}", flush=True)
        print(f"both reached phase 3: {results.get('reachedPhase3')}", flush=True)
        for _, name in sessions:
            print(f"[{name}] cold: {results['cold'][name]['byPage']}", flush=True)
            print(f"[{name}] latch/accessors: {results['accessors'][name]}", flush=True)
            print(f"[{name}] hot-key sweep: {results['hotkeyPhase'][name]}", flush=True)
            print(f"[{name}] census: {results['census'][name]}", flush=True)
        for entry in results["rates"]:
            print(f"  C= + {entry['key']} -> {entry['rate']} baud: landed={entry['landed']} "
                  f"control={[('%02X' % s['aciaControl']) if s['aciaControl'] is not None else '??' for s in entry['after']]} "
                  f"baudIndex={[s['baudIndex'] for s in entry['after']]}", flush=True)
        print(f"C= + 8 swallow: {results['swallow']['reachedSwallow']}, "
              f"lastHotkeyCode={results['swallow']['lastHotkeyCode']}", 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()