#!/usr/bin/env python3
"""Headlessly play through Raiders of the Lost Aisle: boot, start,
navigate room 0 to the toy, cross room 1 timing the guard, reach the
till in room 2 and win. Also separately test getting caught and
timing out. Render key moments to PNG for a visual check."""
import struct, sys, zlib
sys.path.insert(0, "/var/www/html/thespeccy-app/tools")
from z80mini import Z80

SCRATCH = "/tmp/claude-0/-var-www-html/bf8c776a-3726-43e2-a9e3-4d930143eaaa/scratchpad"
code = open("/var/www/html/thespeccy-app/tools/raiders.bin", "rb").read()

ADDR = {}
for line in open("/var/www/html/thespeccy-app/tools/raiders.sym"):
    parts = line.split()
    if len(parts) >= 3 and parts[1] == "EQU":
        ADDR[parts[0]] = int(parts[2].rstrip("H"), 16)


def fresh_cpu():
    cpu = Z80()
    cpu.mem[32768:32768 + len(code)] = code
    cpu.pc = ADDR["start"]
    cpu.on_out = lambda p, v: None
    return cpu


def render(cpu, name):
    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))
    big = []
    for r in rows:
        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))
    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", 512, 384, 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")


KEY = {"Q": (0xFB, 0), "A": (0xFD, 0), "O": (0xDF, 1), "P": (0xDF, 0), "S": (0xFD, 1)}


def press(cpu, *keys, frames=4):
    active = {KEY[k] for k in keys}
    def on_in(port):
        row = port >> 8
        val = 0xFF
        for r, b in active:
            if r == row:
                val &= ~(1 << b) & 0xFF
        return val
    cpu.on_in = on_in
    cpu.run_frames(frames)
    cpu.on_in = lambda port: 0xFF


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


def step(cpu, key, watch, guard=40):
    """Hold `key` (a single move key) until `watch` ('herou' or 'herov')
    changes by exactly one tile, polling frame-by-frame like a real
    player tapping the key — robust to the exact move-gate cadence.
    Not for a step that crosses a room door (position gets reset by
    the transition instead) — use step_through_door for that."""
    before = rd(cpu, watch)
    n = 0
    while rd(cpu, watch) == before:
        press(cpu, key, frames=1)
        n += 1
        assert n < guard, f"{key}: {watch} never changed from {before} after {guard} frames"
    after = rd(cpu, watch)
    assert abs(after - before) == 1, f"{key}: {watch} jumped {before}->{after}, expected a single tile"


def step_through_door(cpu, key, guard=40):
    """Hold `key` until the room counter changes (walking onto a door
    tile resets hero position to the new room's entry point, so we
    can't watch herou/herov for a simple ±1 change here)."""
    before = rd(cpu, "room")
    n = 0
    while rd(cpu, "room") == before:
        press(cpu, key, frames=1)
        n += 1
        assert n < guard, f"{key}: room never advanced from {before} after {guard} frames"


# ============================================================
# 1. Title screen and start
# ============================================================
cpu = fresh_cpu()
cpu.on_in = lambda port: 0xFF
cpu.run_frames(20)
render(cpu, "rd_title")
assert rd(cpu, "state") == 0
press(cpu, "S", frames=3)
assert rd(cpu, "state") == 1, f"state={rd(cpu,'state')}"
assert rd(cpu, "room") == 0 and rd(cpu, "herou") == 4 and rd(cpu, "herov") == 0
print(f"PASS boot+start: room=0 hero=(4,0) time={rd(cpu,'timeleft')}")
render(cpu, "rd_room0")

# ============================================================
# 2. Navigate room 0 to the toy at (0,4), avoiding walls at
#    (1,3),(2,1),(2,2): go P (u+1..) no wait hero starts u=4 already
#    max; path is O(u-1) x4 to reach u=0, then P... let's recompute
#    using the actual key mapping: P=u+1, O=u-1, A=v+1, Q=v-1.
#    From (4,0) to (0,4): O x4 (u:4->0), then A x4 (v:0->4), but the
#    room's walls are (1,3),(2,1),(2,2) -- check the path (u=0..)
#    row-by-row-first isn't blocked since those walls are at u=1,2
#    not u=0. Actually go u down to 3 first via O, then v across via
#    A on row u=3 (clear), then O once more to u=... let's just do:
#    O,O,O (u:4->1), then A,A,A,A (v:0->4) along u=... wait row u=1
#    has a wall at v=3! Use u=3 for the crossing row instead:
#    O (u:4->3), A,A,A,A (v:0->4 along row3, clear), O,O,O (u:3->0).
# ============================================================
step(cpu, "O", "herou")
assert rd(cpu, "herou") == 3, f"herou={rd(cpu,'herou')}"
for _ in range(4):
    step(cpu, "A", "herov")
assert rd(cpu, "herov") == 4, f"herov={rd(cpu,'herov')}"
for _ in range(3):
    step(cpu, "O", "herou")
assert rd(cpu, "herou") == 0 and rd(cpu, "herov") == 4
assert rd(cpu, "toyflag") == 1, "should have picked up the toy"
print(f"PASS toy pickup: hero=({rd(cpu,'herou')},{rd(cpu,'herov')}) toyflag=1")
render(cpu, "rd_gottoy")

# back to the door at (4,4): u:0->3 normally, then the final u:3->4
# step lands on the door and transitions rooms
for _ in range(3):
    step(cpu, "P", "herou")
assert rd(cpu, "herou") == 3
step_through_door(cpu, "P")
assert rd(cpu, "room") == 1, f"room={rd(cpu,'room')} (door should have advanced it)"
assert rd(cpu, "herou") == 2 and rd(cpu, "herov") == 0, "should have entered room 1 at its entry point"
print(f"PASS room transition 0->1: hero=({rd(cpu,'herou')},{rd(cpu,'herov')})")
render(cpu, "rd_room1")

# ============================================================
# 3. Room 1: corridor u=2, v 0..4, guard oscillates v 1..3 at u=2.
#    Wait for the guard to be away from v=1, then dash across.
# ============================================================
guard_positions = []
for _ in range(40):
    cpu.run_frames(1)
    guard_positions.append(rd(cpu, "guardv"))
print(f"  guard v-trace (40 frames): {guard_positions}")
assert min(guard_positions) >= 1 and max(guard_positions) <= 3, "guard should stay within v=1..3"
assert len(set(guard_positions)) > 1, "guard should actually be moving"

# Cross the corridor exactly like a real (imperfect) player would:
# just keep trying to advance toward the door. If the guard catches
# us, checkcatch resets to the entry and we simply resume — the catch
# path itself is verified deterministically in section 5 below, so
# this only needs to prove the crossing eventually succeeds.
catches = 0
was_at_entry = True
guard_budget = 0
while rd(cpu, "room") == 1:
    press(cpu, "A", frames=1)
    guard_budget += 1
    assert guard_budget < 3000, "never made it through the corridor"
    at_entry = rd(cpu, "herov") == 0
    if at_entry and not was_at_entry:
        catches += 1
        assert catches < 20, "caught far too many times — corridor may be unwinnable"
    was_at_entry = at_entry
assert rd(cpu, "room") == 2
print(f"PASS room transition 1->2: crossed after {catches} guard catch(es), "
      f"time left={rd(cpu,'timeleft')}")
render(cpu, "rd_room2")

# ============================================================
# 4. Room 2: entry (2,0), till at (4,4), obstacles (2,2),(3,2).
#    Path: A(v0->1), P(u2->3), P(u3->4), A,A,A (v1->4)
# ============================================================
step(cpu, "A", "herov")
step(cpu, "P", "herou")
step(cpu, "P", "herou")
step(cpu, "A", "herov")
step(cpu, "A", "herov")
# final step lands on the till, which wins outright (no room change)
n = 0
while rd(cpu, "state") != 2:
    press(cpu, "A", frames=1)
    n += 1
    assert n < 40, "never reached the till"
assert rd(cpu, "state") == 2, f"state={rd(cpu,'state')} (should have won)"
print(f"PASS WIN: hero reached the till with the toy, state=2")
render(cpu, "rd_win")

# ============================================================
# 5. Separately: getting caught costs time and resets position
# ============================================================
cpu2 = fresh_cpu()
press(cpu2, "S", frames=3)
step(cpu2, "O", "herou")
for _ in range(4):
    step(cpu2, "A", "herov")
for _ in range(3):
    step(cpu2, "O", "herou")
for _ in range(3):
    step(cpu2, "P", "herou")
step_through_door(cpu2, "P")
assert rd(cpu2, "room") == 1
# the entry tile is deliberately safe (see checkcatch) -- step off it
# before forcing a collision, or this wouldn't test anything
step(cpu2, "A", "herov")
assert rd(cpu2, "herov") == 1
time_before = rd(cpu2, "timeleft")
cpu2.mem[ADDR["guardv"]] = rd(cpu2, "herov")
cpu2.run_frames(2)
time_after = rd(cpu2, "timeleft")
assert rd(cpu2, "herou") == 2 and rd(cpu2, "herov") == 0, "should have been sent back to room 1's entry"
assert time_after == max(0, time_before - 10), f"time {time_before}->{time_after}, expected -10"
print(f"PASS caught-by-guard: time {time_before}->{time_after}, hero reset to entry")

# ============================================================
# 6. Timeout -> lose
# ============================================================
cpu3 = fresh_cpu()
press(cpu3, "S", frames=3)
cpu3.mem[ADDR["timeleft"]] = 1
cpu3.run_frames(60)   # a bit over 1 second of ticks
assert rd(cpu3, "state") == 2 and rd(cpu3, "timeleft") == 0
print("PASS timeout: state=2 (game over), timeleft=0")
render(cpu3, "rd_lose")

print("\nALL CHECKS PASSED")