#!/usr/bin/env python3
# screenshots.py - boot the game in VICE (NTSC, warp off after load) and take screenshots at intervals.
import socket, subprocess, sys, time, os
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])
    times = [float(t) for t in sys.argv[2].split(",")]
    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}/shot_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"); recvUntilPrompt(sock, 10)
    sock.sendall(b"x\n"); recvUntilPrompt(sock, 5)
    print("waiting for $0800 ...", recvUntilPrompt(sock, 300)[-120:])
    # leave warp mode so timing is realistic from here on: toggle via monitor? use 'warp' resource
    sock.sendall(b"resourceset WarpMode 0\n"); print(recvUntilPrompt(sock, 5)[-80:])
    t0 = time.time()
    for i, t in enumerate(times):
        sock.sendall(b"x\n"); recvUntilPrompt(sock, 2)
        time.sleep(max(0, t - (time.time() - t0)))
        sock.sendall(b"\n")                         # any input re-enters the monitor
        out = recvUntilPrompt(sock, 10)
        fn = f"{SCRATCH}/shot_{i:02d}_{int(t)}s.png"
        sock.sendall(f'screenshot "{fn}" 2\n'.encode()); print(recvUntilPrompt(sock, 10)[-100:])
        sock.sendall(b"r\n"); print(recvUntilPrompt(sock, 5)[-200:])
    sock.sendall(b"quit\n")
    time.sleep(1)
    proc.kill(); proc.wait()

main()
