#!/usr/bin/env python3
# viceDriver.py - launch x64sc headless with the remote text monitor, run a list of monitor
# command "stages" (each stage = commands to run when the monitor becomes active), capture output.
import socket, subprocess, sys, time, os, shlex

SCRATCH = os.path.dirname(os.path.abspath(__file__))
def pickFreePort():
    s = socket.socket()
    s.bind(("127.0.0.1", 0))
    port = s.getsockname()[1]
    s.close()
    return port

PORT = pickFreePort()

def recvUntilPrompt(sock, timeout):
    # The monitor prompt looks like "(C:$xxxx) " or "(8:$xxxx) "; wait for it.
    sock.settimeout(0.5)
    buf = b""
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            chunk = sock.recv(65536)
            if not chunk:
                break
            buf += chunk
            tail = buf[-40:]
            if b") " in tail and (b"(C:$" in tail or b"(8:$" in tail):
                break
        except socket.timeout:
            pass
    return buf.decode("latin-1")

def runStages(disk, stages, extraArgs, totalTimeout):
    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] + extraArgs + ["-autostart", disk]
    print("launch:", " ".join(shlex.quote(c) for c in cmd), flush=True)
    proc = subprocess.Popen(cmd, stdout=open(f"{SCRATCH}/vice_stdout.txt", "w"), stderr=subprocess.STDOUT)
    sock = None
    for _ in range(100):
        if os.path.exists(f"{SCRATCH}/vice_stdout.txt") and "bind() failed" in open(f"{SCRATCH}/vice_stdout.txt").read():
            proc.kill()
            raise SystemExit("monitor port collision")
        try:
            sock = socket.create_connection(("127.0.0.1", PORT), timeout=1)
            break
        except OSError:
            time.sleep(0.3)
    if sock is None:
        proc.kill()
        raise SystemExit("could not connect to monitor")
    log = []
    print(recvUntilPrompt(sock, 10), flush=True)
    for stageName, waitSecs, cmds in stages:
        print(f"=== stage: {stageName}", flush=True)
        for c in cmds:
            sock.sendall((c + "\n").encode())
            out = recvUntilPrompt(sock, 30)
            print(f">>> {c}\n{out}", flush=True)
            log.append((c, out))
        # last command of each stage is expected to resume emulation ("x" / "g"); wait for next break
        if waitSecs:
            out = recvUntilPrompt(sock, waitSecs)
            print(f"--- resumed; monitor re-entered with:\n{out}", flush=True)
    try:
        sock.sendall(b"quit\n")
    except OSError:
        pass
    time.sleep(1)
    proc.kill()
    proc.wait()
    return log

if __name__ == "__main__":
    disk = os.path.abspath(sys.argv[1])
    stages = [
        ("setup", 240, ["break 0800", "x"]),
        ("dump at 0800", 0, [
            "r",
            "bank ram",
            f"bsave \"{SCRATCH}/ram_at_0800.bin\" 0 0000 ffff",
            "bank cpu",
            "m 0000 00ff",
            "dev 8:",
            "r",
            f"bsave \"{SCRATCH}/drive_at_0800.bin\" 0 0000 07ff",
            "dev c:", "x"
        ]),
    ]
    runStages(disk, stages, [], 300)