#!/usr/bin/env python3
"""Execute the assembled 8086 runtime on a real CPU emulator and check it.

The runtime (shell/Runtime.mod) is assembled by RtProbe, which prints its
bytes as hex; this script loads them at offset 0 of a flat 64K segment - the
layout a .COM gets - calls each entry with a known argument and compares the
bytes it sends to INT 21h with what it expects.

That is the whole point of the exercise: the library is hand-assembled 8086,
so "it built" says nothing.  This says it *runs*.
"""
import re
import subprocess
import sys

from unicorn import Uc, UC_ARCH_X86, UC_MODE_16
from unicorn.x86_const import (
    UC_X86_REG_AX, UC_X86_REG_DX, UC_X86_REG_SP, UC_X86_REG_IP,
    UC_X86_REG_CS, UC_X86_REG_DS, UC_X86_REG_ES, UC_X86_REG_SS,
)
import unicorn

HERE = __file__.rsplit("/", 1)[0]
PROBE = HERE + "/rtprobe"

HDR = 0x200          # where the fake program header sits
DATA_BASE = 0x300
DATA_END = 0x320
STACK = 0xF000
SENTINEL = 0xBEEF   # "return address" that tells us an entry came back


def load_runtime():
    """Run RtProbe, parse its hex dump, return (bytes, {entry: offset})."""
    out = subprocess.run([PROBE], capture_output=True, text=True, check=True).stdout
    size = int(re.search(r"^(\d+) bytes", out, re.M).group(1))
    entries = {n: int(v) for v, n in re.findall(r"entry \d+ = (\d+)\s+\((\w+)\)", out)}
    blob = bytearray()
    for line in out.splitlines():
        m = re.match(r"^[0-9A-F]{8}  ((?:[0-9A-F]{2} )+)$", line)
        if m:
            blob += bytes.fromhex(m.group(1).replace(" ", ""))
    assert len(blob) == size, f"parsed {len(blob)} bytes, header says {size}"
    return bytes(blob), entries


class Machine:
    def __init__(self, blob):
        self.blob = blob
        self.out = bytearray()
        self.input = bytearray()
        self.uc = Uc(UC_ARCH_X86, UC_MODE_16)
        self.uc.mem_map(0, 0x110000)
        self.uc.mem_write(0, blob)
        # a program header word block: flag, code end, data base, data end
        self.uc.mem_write(HDR, b"\x01\x00\x34\x02\x00\x03\x20\x03\x00\x00")
        self.uc.mem_write(DATA_BASE, b"\xAA" * (DATA_END - DATA_BASE))  # poison
        for r in (UC_X86_REG_CS, UC_X86_REG_DS, UC_X86_REG_ES, UC_X86_REG_SS):
            self.uc.reg_write(r, 0)
        self.uc.hook_add(unicorn.UC_HOOK_INTR, self._intr)


    def _intr(self, mu, intno, _):
        if intno != 0x21:
            return
        ah = (mu.reg_read(UC_X86_REG_AX) >> 8) & 0xFF
        if ah == 0x02:                                   # display character
            self.out.append(mu.reg_read(UC_X86_REG_DX) & 0xFF)
        elif ah == 0x09:                                 # display $-string
            a = mu.reg_read(UC_X86_REG_DX)
            while True:
                b = mu.mem_read(a, 1)[0]
                if b == ord("$"):
                    break
                self.out.append(b)
                a += 1
        elif ah == 0x4C:                                 # exit
            mu.emu_stop()
        elif ah == 0x08:                                 # read char, no echo
            if self.input:
                mu.reg_write(UC_X86_REG_AX, (mu.reg_read(UC_X86_REG_AX) & 0xFF00) | self.input.pop(0))
            else:
                mu.reg_write(UC_X86_REG_AX, (mu.reg_read(UC_X86_REG_AX) & 0xFF00) | 0x1A)

    def call(self, entry, args=(), ax=0):
        """Call an entry with `args` pushed (caller-cleaned, like the compiler
        does).  Returns whatever it wrote to stdout."""
        self.out.clear()
        uc = self.uc
        sp = STACK - 4 * len(args) - 2
        words = [SENTINEL] + list(args)
        uc.mem_write(sp, b"".join(w.to_bytes(2, "little") for w in words))
        uc.reg_write(UC_X86_REG_SP, sp)
        uc.reg_write(UC_X86_REG_AX, ax)
        uc.reg_write(UC_X86_REG_IP, entry)
        # end = SENTINEL: the entry stops by RETurning to it.  Using 0 as the
        # end address would stop instantly for the entry that lives at 0.
        uc.emu_start(entry, SENTINEL, timeout=2_000_000, count=200000)
        if uc.reg_read(UC_X86_REG_IP) != SENTINEL:
            raise AssertionError(
                f"entry {entry} did not return (IP={uc.reg_read(UC_X86_REG_IP):#06x})")
        return bytes(self.out)


def main():
    blob, ent = load_runtime()
    m = Machine(blob)
    fails = []

    def check(name, got, want):
        if got == want:
            print(f"  ok   {name}: {got!r}")
        else:
            print(f"  FAIL {name}: got {got!r} want {want!r}")
            fails.append(name)

    # TU_InitMem must clear [data base, data end) and leave DS alone
    m.call(ent["initmem"], ax=HDR)
    cleared = m.uc.mem_read(DATA_BASE, DATA_END - DATA_BASE)
    check("initmem zeroes globals", cleared, b"\x00" * (DATA_END - DATA_BASE))

    # TU_WrInt: signed 16-bit decimal
    for v, want in [(0, b"0"), (1, b"1"), (7, b"7"), (10, b"10"), (999, b"999"),
                    (12345, b"12345"), (32767, b"32767"),
                    (-1, b"-1"), (-32768, b"-32768"), (-999, b"-999")]:
        check(f"wrint({v})", m.call(ent["wrint"], (v & 0xFFFF,)), want)

    check("wrchar('A')", m.call(ent["wrchar"], (ord("A"),)), b"A")
    check("wrchar('!')", m.call(ent["wrchar"], (ord("!"),)), b"!")
    check("wrbool(0)", m.call(ent["wrbool"], (0,)), b"FALSE")
    check("wrbool(1)", m.call(ent["wrbool"], (1,)), b"TRUE")
    check("wrbool(2)", m.call(ent["wrbool"], (2,)), b"TRUE")
    check("wrln", m.call(ent["wrln"]), b"\r\n")
    check("stackchk returns", m.call(ent["stackchk"]), b"")

    # a sequence, the way a program actually calls these
    m.out.clear()
    m.call(ent["wrint"], (42,))
    m.call(ent["wrchar"], (ord(" "),))
    m.call(ent["wrbool"], (1,))
    m.call(ent["wrln"])
    check("writeln(42) writeln TRUE", bytes(m.out), b"42 TRUE\r\n")

    # TU_RdInt / RdChar / RdBool / RdLn against supplied input
    store = 0x400
    for text, want in [(b"  42abc", 42), (b"-17 x", -17), (b"+5", 5),
                       (b"0", 0), (b"  007", 7), (b"1234", 1234)]:
        m.input = bytearray(text)
        m.uc.mem_write(store, b"\xEE\xEE")
        m.call(ent["rdint"], (store,))
        got = int.from_bytes(m.uc.mem_read(store, 2), "little", signed=True)
        check(f"rdint({text!r})", got, want)
        # the delimiter must be left for the following rdln
        rest = bytes(m.input)
        m.out.clear()
        m.call(ent["rdln"])
        check(f"rdln eats {rest!r}", bytes(m.out), b"")

    for text, want in [(b"Q", ord("Q")), (b"7", ord("7"))]:
        m.input = bytearray(text)
        m.uc.mem_write(store, b"\xEE\xEE")
        m.call(ent["rdchar"], (store,))
        check(f"rdchar({text!r})", int.from_bytes(m.uc.mem_read(store, 2), "little"), want)

    for text, want in [(b"T", 1), (b"y", 1), (b"1", 1), (b"F", 0), (b"n", 0), (b"0", 0)]:
        m.input = bytearray(text)
        m.uc.mem_write(store, b"\xEE\xEE")
        m.call(ent["rdbool"], (store,))
        check(f"rdbool({text!r})", int.from_bytes(m.uc.mem_read(store, 2), "little"), want)

    # end of input must terminate the read loops rather than spin
    m.input = bytearray()
    m.uc.mem_write(store, b"\xEE\xEE")
    m.call(ent["rdint"], (store,))
    check("rdint at EOF", int.from_bytes(m.uc.mem_read(store, 2), "little"), 0)
    m.call(ent["rdln"])

    print()
    if fails:
        print(f"RUNTIME: {len(fails)} FAILURE(S): {', '.join(fails)}")
        return 1
    print("RUNTIME: all checks passed")
    return 0


if __name__ == "__main__":
    sys.exit(main())