#!/usr/bin/env python3
"""m2drive.py — pty driver for scripted Turbo Modula-2 (cpmulator) sessions.
Usage: m2drive.py <drive-dir> <script.txt> [--timeout N]
script.txt lines:  EXPECT <regex> | SEND <text-with-\\r-\\n> | SLEEP <sec> | LOG <note>
Matching is done on a cleaned screen buffer with an advancing cursor, so
repeated EXPECTs (e.g. several "SYM produced") each wait for FRESH output.
Exits 0 iff every EXPECT matched in order.
"""
import os, pty, re, select, sys, time

ANSI = re.compile(rb'\x1b\[[0-9;?]*[a-zA-Z]|\x1b[()][AB0]')


def main():
    drv, script = sys.argv[1], sys.argv[2]
    timeout = float(sys.argv[4]) if len(sys.argv) > 4 else 10.0
    steps = []
    for line in open(script):
        line = line.rstrip('\n')
        if not line or line.startswith('#'):
            continue
        cmd, _, arg = line.partition(' ')
        steps.append((cmd, arg))
    pid, fd = pty.fork()
    if pid == 0:
        os.execvp('/home/eric/bin/CPM/Cpmulator/cpmulator',
                  ['cpmulator', '-drive-d', drv])
    clean = b''   # ANSI-stripped screen, consumed as matched
    try:
        for cmd, arg in steps:
            if cmd == 'SEND':
                data = arg.encode().decode('unicode_escape').encode('latin1')
                os.write(fd, data)
                time.sleep(1.0)
                clean += drain(fd, 1.0)
            elif cmd == 'SLEEP':
                clean += drain(fd, float(arg))
            elif cmd == 'EXPECT':
                t0 = time.time()
                while True:
                    m = re.search(arg.encode(), clean)
                    if m:
                        clean = clean[m.end():]
                        break
                    if time.time() - t0 > timeout:
                        print(f'MISS: {arg}')
                        print('--- screen tail ---')
                        print(clean[-2000:].decode('latin1', 'replace'))
                        return 1
                    clean += drain(fd, 0.5)
                print(f'HIT: {arg}', flush=True)
            elif cmd == 'LOG':
                print(f'# {arg}', flush=True)
        print('--- done ok ---')
        print('--- screen tail ---')
        print(cleantail(clean, 1500))
        return 0
    finally:
        try:
            os.close(fd)
        except OSError:
            pass


def cleantail(clean, n):
    return clean[-n:].decode('latin1', 'replace')


def drain(fd, wait):
    """Read available pty output for up to `wait` seconds, ANSI-stripped."""
    out = b''
    end = time.time() + wait
    while time.time() < end:
        r, _, _ = select.select([fd], [], [], max(0, end - time.time()))
        if not r:
            continue
        try:
            chunk = os.read(fd, 65536)
        except OSError:
            break
        if not chunk:
            break
        out += chunk
    return ANSI.sub(b'', out)


sys.exit(main())