#!/usr/bin/env python3
"""Headlessly verify lyingport: the border only responds to even port
numbers, and a keyboard IN read is entirely governed by A (the upper
address byte), not by the immediate low byte written in the source.
Asserts the exact register values the program computes, then renders
a screenshot -- text and border colour both -- as visual proof."""
import os, struct, sys, zlib
TOOLS_DIR = os.path.dirname(os.path.abspath(__file__))
sys.path.insert(0, TOOLS_DIR)
from z80mini import Z80

SCRATCH = os.environ.get("SPECCY_TEST_SCRATCH", os.path.join(TOOLS_DIR, "_renders"))
os.makedirs(SCRATCH, exist_ok=True)
code = open(os.path.join(TOOLS_DIR, "lyingport.bin"), "rb").read()

ADDR = {}
for line in open(os.path.join(TOOLS_DIR, "lyingport.sym")):
    parts = line.split()
    if len(parts) >= 3 and parts[1] == "EQU":
        ADDR[parts[0]] = int(parts[2].rstrip("H"), 16)

BASIC_RET = 0x1234

# z80mini has no real ROM loaded, so RST 0x10 (PRINT-A) and CALL 0x1601
# (CHAN-OPEN) are stubbed rather than actually executed. CHAN-OPEN is a
# no-op (there's only one "channel" here: the screen); PRINT-A is stood
# in for with a real glyph blit, using the same public-domain 8x8 font
# already checked into fontdata.inc for the polyglot-screen article, so
# what lands in screen memory -- and what the render below shows -- is
# genuine ROM-shaped bitmap data, not a placeholder.
FONT = {}
_code = 32
for _line in open(os.path.join(TOOLS_DIR, "fontdata.inc")):
    _line = _line.strip()
    if not _line.startswith("defb"):
        continue
    _bytes = [int(b, 16) for b in _line.split(";")[0].split("defb")[1].split(",") if b.strip()]
    FONT[_code] = _bytes
    _code += 1


def bitmap_off(y, cx):
    return 16384 + (((y & 0xC0) << 5) | ((y & 7) << 8) | ((y & 0x38) << 2) | cx)


def fresh_cpu():
    cpu = Z80()
    cpu.mem[32768:32768 + len(code)] = code
    cpu.pc = ADDR["start"]
    cpu.sp = 0xFFF0
    cpu.mem[cpu.sp] = BASIC_RET & 0xFF
    cpu.mem[cpu.sp + 1] = BASIC_RET >> 8

    cursor = {"row": 0, "col": 0}

    def newline():
        cursor["col"] = 0
        cursor["row"] = (cursor["row"] + 1) % 24  # wrap rather than scroll -- good enough here

    def print_char():
        ch = cpu.a
        if ch == 13:
            newline()
            return
        row, col = cursor["row"], cursor["col"]
        glyph = FONT.get(ch, FONT[32])
        for gy in range(8):
            cpu.mem[bitmap_off(row * 8 + gy, col)] = glyph[gy]
        cpu.mem[22528 + row * 32 + col] = 0x47  # bright white ink on black paper
        cursor["col"] += 1
        if cursor["col"] >= 32:
            newline()

    cpu.call_stubs[0x1601] = lambda: None   # CHAN-OPEN: only one channel here
    cpu.call_stubs[0x0010] = print_char     # RST 0x10 target: PRINT-A
    return cpu


def rd(cpu, name):
    return cpu.mem[ADDR[name]]


# Q lives on row 0xFB, bit 0 -- same convention as the other tools/
# tests (see test_gigrush.py's KEY dict). Held down for the whole run.
def on_in(port):
    row = port >> 8
    if row == 0xFB:
        return 0xFE  # bit 0 clear = Q pressed, everything else up
    return 0xFF


border = {"value": 7, "log": []}


def on_out(port, value):
    # real ULA behaviour: only bit 0 of the port byte is decoded for
    # the border/MIC/EAR register -- everything else is ignored.
    if port & 1 == 0:
        border["value"] = value & 7
    border["log"].append((port, value, border["value"]))


def render(cpu, name, border_colour):
    PAL = [(0,0,0),(0,0,215),(215,0,0),(215,0,215),(0,215,0),(0,215,215),(215,215,0),(215,215,215)]
    BRT = [(0,0,0),(0,0,255),(255,0,0),(255,0,255),(0,255,0),(0,255,255),(255,255,0),(255,255,255)]
    def off(y, cx): return 16384 + (((y & 0xC0) << 5) | ((y & 7) << 8) | ((y & 0x38) << 2) | cx)
    rows = []
    for y in range(192):
        r = bytearray()
        for cx in range(32):
            bits = cpu.mem[off(y, cx)]
            attr = cpu.mem[22528 + (y // 8) * 32 + cx]
            pal = BRT if attr & 0x40 else PAL
            ink, paper = pal[attr & 7], pal[(attr >> 3) & 7]
            for bit in range(8):
                r += bytes(ink if bits & (0x80 >> bit) else paper)
        rows.append(bytes(r))

    # pad with a genuine border margin using the tracked border colour
    # -- the whole point of this program is what happened to that
    # register, so the screenshot has to show it, not just the text.
    bc = PAL[border_colour]
    border_px = 24
    width = 256 + border_px * 2
    padded = []
    pad_row = bytes(bc) * width
    for _ in range(border_px):
        padded.append(pad_row)
    for r in rows:
        padded.append(bytes(bc) * border_px + bytes(r) + bytes(bc) * border_px)
    for _ in range(border_px):
        padded.append(pad_row)

    big = []
    for r in padded:
        r2 = b"".join(r[i:i + 3] * 2 for i in range(0, len(r), 3))
        big += [r2, r2]
    raw = b"".join(b"\x00" + r for r in big)
    def chunk(t, d):
        c = t + d
        return struct.pack(">I", len(d)) + c + struct.pack(">I", zlib.crc32(c))
    h, w = len(big), width * 2
    with open(f"{SCRATCH}/{name}.png", "wb") as f:
        f.write(b"\x89PNG\r\n\x1a\n")
        f.write(chunk(b"IHDR", struct.pack(">IIBBBBB", w, h, 8, 2, 0, 0, 0)))
        f.write(chunk(b"IDAT", zlib.compress(raw, 9)))
        f.write(chunk(b"IEND", b""))
    print(f"  render -> {name}.png ({w}x{h})")


def run(cpu):
    cpu.on_in = on_in
    cpu.on_out = on_out
    guard = 2_000_000
    while cpu.pc != BASIC_RET and guard:
        cpu.step()
        guard -= 1
    assert guard, "program never returned to BASIC"


def main():
    cpu = fresh_cpu()
    run(cpu)

    # --- Part 1: only even port numbers reach the border ---
    # every write in the program: (254,2) (252,4) (0,5) (128,6) (255,1)
    ports_written = [p for p, v, b in border["log"]]
    assert ports_written == [254, 252, 0, 128, 255], ports_written
    # the odd one (255) must NOT have changed the register: final
    # colour is 6 (yellow), the last EVEN write, not 1 (blue).
    assert border["value"] == 6, f"border ended up {border['value']}, expected 6 (255 should have been ignored)"
    print(f"  ok: border survived four even ports, ignored the odd one -> final colour {border['value']}")

    # --- Part 2: IN is governed by A, not by the immediate low byte ---
    row1, row2, row3 = rd(cpu, "row1"), rd(cpu, "row2"), rd(cpu, "row3")
    assert row1 == 0xFE, f"row1 = {row1:02X}, expected FE (Q's row, Q held)"
    assert row2 == 0xFF, f"row2 = {row2:02X}, expected FF (a row with no key held)"
    assert row3 == row1 == 0xFE, f"row3 = {row3:02X}, expected to match row1 (FE) -- N should not matter"
    assert row2 != row1, "row2 should differ from row1 -- A picked a different row"
    print(f"  ok: row1={row1:02X} row2={row2:02X} row3={row3:02X} -- A decided the row, N never did")

    render(cpu, "lyingport", border["value"])


if __name__ == "__main__":
    main()
    print("lyingport: all checks passed")
