#!/usr/bin/env python3
# traceRun.py - run the game under VICE (NTSC) to $0800, then enable an exec tracepoint over the
# whole RAM range and collect every executed PC for a wall-clock budget; write a coverage file.
import socket, subprocess, sys, time, os, re
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
from viceDriver import pickFreePort, recvUntilPrompt, SCRATCH

def main():
    disk = os.path.abspath(sys.argv[1])
    budget = int(sys.argv[2])
    outName = sys.argv[3]
    port = pickFreePort()
    cmd = ["xvfb-run", "-a", "x64sc", "-remotemonitor", "-remotemonitoraddress", f"127.0.0.1:{port}",
           "-drive8truedrive", "-drive8type", "1541", "-sounddev", "dummy", "-warp", "-jamaction", "2",
           "-autostart-warp", "-model", "ntsc", "-chdir", SCRATCH, "-autostart", disk]
    proc = subprocess.Popen(cmd, stdout=open(f"{SCRATCH}/trace_stdout.txt", "w"), stderr=subprocess.STDOUT)
    sock = None
    for _ in range(100):
        try:
            sock = socket.create_connection(("127.0.0.1", port), timeout=1); break
        except OSError:
            time.sleep(0.3)
    print(recvUntilPrompt(sock, 10))
    sock.sendall(b"break 0800\n"); print(recvUntilPrompt(sock, 10))
    sock.sendall(b"x\n"); print(recvUntilPrompt(sock, 300))
    sock.sendall(b"del\n"); print(recvUntilPrompt(sock, 10))
    sock.sendall(b"tr exec 0000 ffff\n"); print(recvUntilPrompt(sock, 10))
    sock.sendall(b"x\n")
    pcs = {}
    buf = b""
    deadline = time.time() + budget
    sock.settimeout(1.0)
    total = 0
    rx = re.compile(rb"\.C:([0-9a-f]{4})")
    while time.time() < deadline:
        try:
            chunk = sock.recv(1 << 20)
            if not chunk: break
            total += len(chunk)
            buf += chunk
            lines = buf.split(b"\n")
            buf = lines[-1]
            for ln in lines[:-1]:
                m = rx.search(ln)
                if m:
                    pc = int(m.group(1), 16)
                    pcs[pc] = pcs.get(pc, 0) + 1
        except socket.timeout:
            pass
    with open(f"{SCRATCH}/{outName}", "w") as f:
        for pc in sorted(pcs):
            f.write(f"{pc:04X} {pcs[pc]}\n")
    print(f"bytes received {total}, distinct PCs {len(pcs)}")
    try:
        sock.sendall(b"\n")
        time.sleep(0.5)
        sock.sendall(b"quit\n")
    except OSError:
        pass
    time.sleep(1)
    proc.kill(); proc.wait()

main()