#!/usr/bin/env python3
# Drive the captured Space Taxi RAM through the 6502 emulator to
# decompress every scene the game knows how to load.
#
# Approach:
#   1. Load a VICE memory dump as flat 64KB RAM. The dump captured at
#      runtime contains the decompressor at $9656, the scene pointer
#      table at $9600, and (we hope) the compressed data the pointers
#      reach into.
#   2. Sanity check: run scene ID 1 (the captured live level) through
#      $6283 and compare the resulting screen buffer at $7530-$7929
#      with the live screen RAM at $0400-$07E7 already in the dump.
#      If they match, the emulator + dump together can reproduce any
#      scene. If not, we know something's missing (likely hidden RAM
#      under ROMs that VICE didn't capture).
#   3. If sanity passes: iterate scene IDs 0..23 (lettered levels)
#      plus known special IDs (title=$24, level-complete=$19) and
#      dump each scene's screen + color buffer.

import sys
import os
sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))

from cpu6502 import Cpu6502


def load_dump(path):
    with open(path, 'rb') as f:
        raw = f.read()
    if len(raw) == 0x10000:
        return bytearray(raw)
    if len(raw) == 0x10002:
        # 2-byte LE load-addr header
        return bytearray(raw[2:])
    raise SystemExit(f'{path}: unexpected dump size {len(raw)}')


# Memory map landmarks (from MECHANICS.md):
ADDR_DECOMP_ENTRY   = 0x6283   # bank-switch wrapper around $9656
SCREEN_BUF_BASE     = 0x7530   # decompressed screen RAM target
SCREEN_BUF_SIZE     = 0x400    # 4 pages = 1024 bytes
COLOR_BUF_BASE      = 0x7918   # decompressed color RAM target
COLOR_BUF_SIZE      = 0x400


def run_scene(ram, scene_id):
    """Run $6283(X=scene_id) on a fresh copy of `ram`. Returns a tuple
    (screen_buf, color_buf, final_ram) so the caller can compare or
    examine state."""
    work = bytearray(ram)        # copy so callers can keep the original
    cpu  = Cpu6502(work)
    cpu.a = scene_id
    cpu.x = scene_id
    cpu.y = 0
    cpu.call(ADDR_DECOMP_ENTRY)
    screen = bytes(work[SCREEN_BUF_BASE:SCREEN_BUF_BASE + SCREEN_BUF_SIZE])
    color  = bytes(work[COLOR_BUF_BASE:COLOR_BUF_BASE  + COLOR_BUF_SIZE])
    return screen, color, work


def compare_screen(decomp_screen, live_screen, label):
    """Compare 1024-byte decompressed screen vs live screen RAM.
    Only the first 1000 bytes are visible (40x25). Bytes 1000..1023
    are unused tail / overlap with color RAM, so we ignore them."""
    a = decomp_screen[:1000]
    b = live_screen[:1000]
    if a == b:
        print(f'  {label}: MATCH')
        return True
    # Show how many cells differ and a sample
    diffs = [(i, a[i], b[i]) for i in range(1000) if a[i] != b[i]]
    print(f'  {label}: {len(diffs)} / 1000 cells differ')
    for i, av, bv in diffs[:6]:
        row, col = divmod(i, 40)
        print(f'    cell ({col:2d},{row:2d}): decomp=${av:02X} live=${bv:02X}')
    return False


def main():
    if len(sys.argv) < 2:
        print('usage: dumpScenes.py <dump-path> [scene-id]', file=sys.stderr)
        sys.exit(2)
    ram = load_dump(sys.argv[1])

    # Figure out what the captured state is, then pick the scene to
    # test against. $7215 = current level number (0..23 = A..X),
    # but the title screen / specials don't update it -- the
    # most-recent JSR $6283 used scene $24 (title) or $19 (other
    # special). Default to the level-number scene, but allow override.
    state_level = ram[0x7215]
    if len(sys.argv) >= 3:
        scene = int(sys.argv[2], 0)
    else:
        # Heuristic: if state sentinel $07A6 says we're on the title,
        # the live screen matches scene 36 ($24). Otherwise we're
        # mid-gameplay and the live screen matches scene = level.
        # Title state in this game: top-left tile is the inner frame
        # corner $5D / $5C. For now just try scene 36 then scene N.
        scene = 36
    print(f'=== Sanity test: scene {scene} (${scene:02X}) vs live screen RAM ===')
    print(f'  captured state: $7215={state_level} $07A6=${ram[0x07A6]:02X}')
    try:
        scr, col, _ = run_scene(ram, scene)
    except Exception as e:
        print(f'  emulator raised: {type(e).__name__}: {e}')
        return
    live_scr = bytes(ram[0x0400:0x0400 + 1024])
    live_col = bytes(b & 0x0F for b in ram[0xD800:0xD800 + 1024])
    compare_screen(scr, live_scr, 'screen')
    masked_col = bytes(b & 0x0F for b in col)
    compare_screen(masked_col, live_col, 'color ')


if __name__ == '__main__':
    main()