#!/usr/bin/env python3
# Convert ROM-extracted Space Taxi scenes directly into the runtime
# STL1 binary format (see stLevel.c). Uses the
# 6502 emulator + raw RAM dump to decompress each scene, then writes
# screen + color RAM bytes straight into the .dat layout the JoeyLib
# port loads at runtime.
#
# Replaces the .txt -> mkstlevel pipeline for canonical levels. The C
# mkstlevel tool stays in tree for hand-authored .txt levels.
#
# Usage:
#     romToLevel.py <raw-dump> <out-dir>            # all 24 levels
#     romToLevel.py <raw-dump> <out-dir> <scene_id> # one specific
#
# Writes:
#     <out-dir>/level01.dat ... level24.dat   (scene 0 = "A" = level01)
#     <out-dir>/tbank0.png                    (charset extracted from $2800)

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

from cpu6502 import Cpu6502
from dumpScenes import load_dump, run_scene


# STL2 binary layout (see stLevel.c):
#   "STL2"
#   nameLen + name
#   tileBank, music, bgColor, borderColor, taxiSpawnX, taxiSpawnY
#   xAccel, yAccel, xGrav (int8), yGrav (int8)
#   bgColor1, bgColor2, bgColor3, spriteMc0, spriteMc1, sprite0Color, sprite1Color
#   padCount + padCount * (letter, tileX, tileY, tileW)
#   fareCount + fareCount * (spawnPad, destPad)
#   tilemap[40 * 25]
#   colormap[40 * 25]

ST_TILEMAP_W      = 40
ST_PLAYFIELD_ROWS = 25
ST_TILE_PX        = 8

# VIC-II sprite-register origins: sprite-X $18 (24) is display column 0,
# sprite-Y $32 (50) is the first visible raster line. All in-game pad/taxi
# positions are stored in this sprite-register space, so tile conversion
# subtracts these origins. PAD_SURFACE_ORIGIN folds in the ~14px landing-
# gear offset (50-14=36) so a pad's row byte maps to its contact-surface
# tile row rather than the sprite-top row.
VIC_X_ORIGIN      = 24
VIC_Y_ORIGIN      = 50
PAD_SURFACE_ORIGIN = 36
LEVEL_LETTERS = [chr(ord('A') + i) for i in range(24)]


def write_level(path, level_index, after, screen, colormap, sprite_blocks):
    """Write an STL4 level .dat straight from post-decompression RAM.

    Everything stays in the C64's own units so the runtime can mirror the
    original routines byte for byte:
      name[22]        $7D78-$7D8D  level name (screen codes, ASCII-compatible)
      header[9]       $7D00-$7D08  VIC colours: border, bg, bg1-3, mc0, mc1,
                                   sprite-0, sprite-1
      padCount        $7D5A        pads in the table
      specialPad      $7D09        fare pads (== padCount, or padCount-1 when
                                   the last pad is the fuel station)
      fuelRate        $7D60        ticks of flight per half fuel cell
      spawn[5]        $7D5B-$7D5F  X frac, X lo, X hi, Y frac, Y hi (sprite coords)
      templates[8]    $7D8F-$7D96  Y accel, X accel, Y gravity, X gravity
                                   (16-bit little-endian each)
      pads[10][8]     $7D0A-$7D59  raw slots: x1 hi/lo, x2 hi/lo, row,
                                   passenger X frac, passenger X col, unused
      tilemap[1000]   screen RAM (all 25 rows, the HUD rows included)
      colormap[1000]  colour RAM low nibbles
      levelIndex      0..23 = A..X, 24 = title
      spriteCount     N, then N x (ptr, 63 bytes): the per-level sprite
                      bitmaps the level's extra scene loads at $2000+
      hookData[616]   $7D98-$7FFF: the level's hook code and tables
    """
    cells = ST_TILEMAP_W * ST_PLAYFIELD_ROWS
    assert len(screen)   == cells, f'screen is {len(screen)}, need {cells}'
    assert len(colormap) == cells, f'colormap is {len(colormap)}, need {cells}'
    with open(path, 'wb') as fp:
        fp.write(b'STL4')
        fp.write(bytes(after[0x7D78:0x7D78 + 22]))
        fp.write(bytes(after[0x7D00:0x7D09]))
        fp.write(bytes([after[0x7D5A], after[0x7D09], after[0x7D60]]))
        fp.write(bytes(after[0x7D5B:0x7D60]))
        fp.write(bytes(after[0x7D8F:0x7D97]))
        pad_count = min(after[0x7D5A], 10)
        for pad_idx in range(10):
            if pad_idx < pad_count:
                off = 0x7D0A + pad_idx * 8
                fp.write(bytes(after[off:off + 8]))
            else:
                fp.write(bytes(8))
        fp.write(bytes(screen))
        fp.write(bytes(colormap))
        fp.write(bytes([level_index]))
        fp.write(bytes([len(sprite_blocks)]))
        for ptr, data in sprite_blocks:
            fp.write(bytes([ptr]))
            fp.write(bytes(data))
        fp.write(bytes(after[0x7D98:0x8000]))


# Levels whose data blob flags an extra scene ($7D97 != 0) load a second
# scene from this table ($6F55, indexed by level) right after the level
# itself. Those extra scenes only write sprite bitmaps into $2000-$27FF
# (block pointers $80-$9F) for the level's hook code to animate.
_EXTRA_SCENE = [0, 0, 0, 0, 0, 0x1A, 0x1B, 0x1C, 0x1D, 0x1E, 0, 0x1F, 0, 0, 0,
                0x20, 0x21, 0, 0, 0, 0x22, 0, 0, 0, 0x23]


def extract_one_level(ram, scene_id, out_dir):
    """Decompress `scene_id` and write out_dir/level<NN>.dat (NN = scene+1)."""
    letter = LEVEL_LETTERS[scene_id]
    screen_buf, color_buf, after = run_scene(ram, scene_id)
    sprite_blocks = []
    if after[0x7D97] != 0 and _EXTRA_SCENE[scene_id] != 0:
        # Start the extra scene from a zeroed sprite window so the blocks
        # it writes are exactly the ones the level ships with.
        work = bytearray(after)
        for i in range(0x2000, 0x2800):
            work[i] = 0
        _s, _c, after2 = run_scene(bytes(work), _EXTRA_SCENE[scene_id])
        for ptr in range(0x80, 0xA0):
            block = bytes(after2[ptr * 64:ptr * 64 + 63])
            if any(block):
                sprite_blocks.append((ptr, block))
        # The level data itself (pads, hooks) is untouched by the extra
        # scene; keep reading it from the level's own image.
    cells = ST_TILEMAP_W * ST_PLAYFIELD_ROWS
    screen   = list(screen_buf[:cells])
    colormap = [(b & 0x0F) for b in color_buf[:cells]]
    level_num = scene_id + 1
    out_path = os.path.join(out_dir, f'level{level_num:02d}.dat')
    write_level(out_path, scene_id, after, screen, colormap, sprite_blocks)
    return out_path, after[0x7D5A], len(sprite_blocks)


def extract_screen_y(ram, out_dir):
    """Screen Y, the 25th screen ("MYSTERY SCREEN !!!"), is scene 24.

    $5010 dispatches on the screens-completed counter: at 24 it jumps to
    $5167, which loads #$59 ('Y') as the pending screen; the shift only
    ends when the counter reaches 25. So a full 24-hour shift plays A..X
    and then Y -- the 25 screens the intro advertises.
    """
    screen_buf, color_buf, after = run_scene(ram, 24)
    cells = ST_TILEMAP_W * ST_PLAYFIELD_ROWS
    screen   = list(screen_buf[:cells])
    colormap = [(b & 0x0F) for b in color_buf[:cells]]
    out_path = os.path.join(out_dir, 'level25.dat')
    write_level(out_path, 24, after, screen, colormap, [])
    return out_path


def extract_title(ram, out_dir):
    """The title screen is scene 25 (what $477A decompresses).
    Pads, spawn and templates are dummies."""
    screen_buf, color_buf, after = run_scene(ram, 25)
    cells = ST_TILEMAP_W * ST_PLAYFIELD_ROWS
    screen   = list(screen_buf[:cells])
    colormap = [(b & 0x0F) for b in color_buf[:cells]]
    out_path = os.path.join(out_dir, 'title.dat')
    # The title scene writes no level name, so $7D78 still holds whatever
    # the previously decompressed scene left there (PUZZLER, level H).
    # The title never draws the name; ship it blank rather than stale.
    work = bytearray(after)
    work[0x7D78:0x7D78 + 22] = b' ' * 22
    write_level(out_path, 25, work, screen, colormap, [])
    return out_path


def main():
    if len(sys.argv) < 3:
        print('usage: romToLevel.py <raw-dump> <out-dir> [scene_id]', file=sys.stderr)
        sys.exit(2)
    ram = load_dump(sys.argv[1])
    out_dir = sys.argv[2]
    os.makedirs(out_dir, exist_ok=True)

    if len(sys.argv) >= 4:
        scene_ids = [int(sys.argv[3], 0)]
    else:
        # Levels A..X = scenes 0..23. The title comes from screen RAM
        # at $0400 (the dump was taken on the title screen), not from
        # a decompressed scene.
        scene_ids = list(range(24))

    for sid in scene_ids:
        path, pad_count, nspr = extract_one_level(ram, sid, out_dir)
        size = os.path.getsize(path)
        print(f'  scene {sid:2d} (level {LEVEL_LETTERS[sid]}): pads={pad_count} sprites={nspr} size={size} -> {path}')

    # Screen Y and the title are not part of the A..X sweep.
    if len(sys.argv) < 4:
        y_path = extract_screen_y(ram, out_dir)
        print(f'  scene 24 (screen Y)      : -> {y_path} ({os.path.getsize(y_path)} bytes)')
        title_path = extract_title(ram, out_dir)
        print(f'  title (screen RAM $0400) : -> {title_path} ({os.path.getsize(title_path)} bytes)')


if __name__ == '__main__':
    main()