#!/usr/bin/env python3
# extractFromDump.py -- convert a VICE C64 memory dump into JoeyLib
# Space Taxi assets.
#
# Reads:   /tmp/spacetaxi/mem0000-<labelN>.bin (VICE `save` output;
#          2-byte LE start-addr header + 64 KB RAM).
# Writes:  examples/spacetaxi/assets/tiles/tilebank<N>.png
#          examples/spacetaxi/assets/levels/level<NN>.txt
#          examples/spacetaxi/assets/sprites/sprites.png  (level1 only)
#          examples/spacetaxi/assets/font.png             (level1 only)
#
# Per-level:
#   - Custom charset at $2800-$2FFF (2 KB, 256 chars * 8 bytes mono)
#     becomes a tile bank PNG: 40 cols x 25 rows of 8x8 cells, with
#     tile index N at cell (N%40, N/40). The PNG uses the EGA-ish
#     16-color palette we already install on stage.
#   - Screen RAM at $0400-$07E7 (1000 bytes = 40x25 char codes) and
#     color RAM at $D800-$DBE7 (1000 bytes = 40x25 palette nybbles)
#     become the tilemap+colormap in the level .txt source.
#
# Sprite data at $3000-$37FF + sprite pointer table at $07F8-$07FF
# is decoded once (it's the same across levels for the active set):
#   - sprite 0 / 7: the taxi cab body & shadow
#   - sprite 1..6: passenger frames / animation cels
# 64 bytes per sprite, 24x21 1-bit, MSB-leftmost. Color comes from
# $D027-$D02E (one color per sprite).
#
# Usage:
#   python3 extractFromDump.py stuff/spacetaxi/mem0000-level1.bin level01 1
#   python3 extractFromDump.py stuff/spacetaxi/mem0000-level2.bin level02 2
# args: <dump-path> <level-name-without-extension> <tilebank-id>
#
# Dumps live in stuff/spacetaxi/ (committed-ish, repo-local) rather
# than /tmp so they survive reboots and stay with the project.

import os
import struct
import sys
from collections import Counter

from PIL import Image


# C64 standard color palette (EGA-ish 16-color JoeyLib mapping).
# Indices 0..15 match the C64 VIC-II color register order so a value
# read from $D8xx or $D02x maps directly to a JoeyLib palette slot.
# Colors below approximate Pepto's standard PAL palette in $RGB form.
C64_PALETTE_RGB = [
    (0x00, 0x00, 0x00),  # 0  black
    (0xFF, 0xFF, 0xFF),  # 1  white
    (0x88, 0x39, 0x32),  # 2  red
    (0x67, 0xB6, 0xBD),  # 3  cyan
    (0x8B, 0x3F, 0x96),  # 4  purple
    (0x55, 0xA0, 0x49),  # 5  green
    (0x40, 0x31, 0x8D),  # 6  blue
    (0xBF, 0xCE, 0x72),  # 7  yellow
    (0x8B, 0x54, 0x29),  # 8  orange
    (0x57, 0x42, 0x00),  # 9  brown
    (0xB8, 0x69, 0x62),  # 10 light red
    (0x50, 0x50, 0x50),  # 11 dark gray
    (0x78, 0x78, 0x78),  # 12 mid gray
    (0x94, 0xE0, 0x89),  # 13 light green
    (0x78, 0x69, 0xC4),  # 14 light blue
    (0x9F, 0x9F, 0x9F),  # 15 light gray
]


ROOT = os.path.dirname(os.path.abspath(__file__))


def load_dump(path):
    with open(path, "rb") as f:
        raw = f.read()
    start = raw[0] | (raw[1] << 8)
    if start != 0:
        raise SystemExit(f"{path}: unexpected start ${start:04X}, want $0000")
    if len(raw) - 2 != 0x10000:
        raise SystemExit(f"{path}: not a 64K dump (got {len(raw) - 2} bytes)")
    return raw[2:]


def new_canvas(w, h):
    img = Image.new("P", (w, h), 0)
    flat = bytearray()
    for r, g, b in C64_PALETTE_RGB:
        flat.extend((r, g, b))
    flat.extend(b"\x00" * (768 - len(flat)))
    img.putpalette(bytes(flat))
    return img


def paint_charset_glyph(px, mem, src_code, bx, by):
    """Plot the 8x8 glyph for C64 screen code `src_code` at pixel
    (bx, by) of the indexed-palette PIL image. Foreground = palette
    index 1, background untouched (already index 0 from new_canvas).
    """
    src = 0x2800 + src_code * 8
    for r in range(8):
        byte = mem[src + r]
        for c in range(8):
            if byte & (0x80 >> c):
                px[bx + c, by + r] = 1


def extract_tilebank(mem, out_path, code_to_slot=None):
    """Render the custom charset at $2800-$2FFF as a 320x200 PNG.
    When `code_to_slot` is None, lays out all 256 chars at their raw
    C64 positions (40 cols x 25 rows). When `code_to_slot` is given,
    writes each charset glyph at the slot the level loader expects
    -- i.e. the bank tile at position N is the glyph for screen-code
    `code_for_slot[N]`. That keeps level data and tilebank aligned
    after the level extractor's frequency-sorted code remapping.
    """
    img = new_canvas(320, 200)
    px = img.load()
    if code_to_slot is None:
        for ch in range(256):
            paint_charset_glyph(px, mem, ch, (ch % 40) * 8, (ch // 40) * 8)
    else:
        # Invert: slot -> code. Multiple codes can collide into slot 0
        # (background); only the first one we see "owns" that slot for
        # painting purposes -- the rest just fall through to the same
        # blank tile. Highest priority is the explicit space mapping.
        slot_to_code = {}
        for code, slot in code_to_slot.items():
            slot_to_code.setdefault(slot, code)
        for slot, code in slot_to_code.items():
            if slot >= 256:
                continue
            paint_charset_glyph(px, mem, code, (slot % 40) * 8, (slot // 40) * 8)
    img.save(out_path, "PNG", optimize=True)


def build_code_to_slot(mem):
    """Build the screen-code -> alphabet-slot map shared between the
    tilebank emission and the level emission. The slot index is what
    the runtime sees in level.tilemap[]; the tilebank must mirror this
    ordering so a level-data slot-N cell renders the charset glyph
    that was actually used in the original screen RAM at that code.
    """
    screen = mem[0x0400:0x0400 + 1000]
    codes = sorted(Counter(screen).items(), key=lambda kv: (-kv[1], kv[0]))
    code_to_slot = {}
    next_slot = 0
    for code, _count in codes:
        if code == 0x20:                  # space -> always index 0
            code_to_slot[code] = 0
            continue
        if code in code_to_slot:
            continue
        if next_slot >= len(ALPHABET):
            code_to_slot[code] = 0
            continue
        # Reserve slot 0 for background (space). The very first
        # non-space code claims slot 1, etc.
        if next_slot == 0:
            next_slot = 1
        code_to_slot[code] = next_slot
        next_slot += 1
    return code_to_slot


def screencode_to_char(code):
    """Convert a C64 screen code into a printable tilemap character
    using the encoding the .txt format already understands:
        0      -> '.' (treated as 0)
        1..9   -> '1'..'9'  (digits encode the lower 10 indices)
                  but we shift to use up most of the index space below.
    For arbitrary 0..255 we need a 256-symbol alphabet -- not feasible
    in a single character per cell. So we encode as TWO chars per cell
    using base-16 hex would still only cover 0..255. Decision: keep
    .txt cells one-char and map screen codes via a custom 64-char
    alphabet plus an @char directive table emitted into the file so
    the loader knows exactly which screen code each char represents.
    """
    # This function is unused as a one-char encoder; we encode via the
    # generic alphabet below instead.
    return code


# Generic 1-char alphabet (62 codes); the level loader's index map
# was: '.' or ' '=0, '0'..'9'=0..9, 'A'..'Z'=10..35, 'a'..'z'=36..61.
# That's only 62 codes, not 256. The original C64 levels use ~30
# distinct screen codes per level so most levels fit. We map the N
# most-frequent screen codes to those alphabet slots.
ALPHABET = "0123456789ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz"


def extract_level(mem, level_name, tilebank_id, out_path, code_to_alpha):
    """Render screen RAM + color RAM as an STL1-format .txt source
    using the prebuilt screen-code -> alphabet-slot mapping. Caller
    must pass the same mapping that was used to reorder the tilebank
    PNG; otherwise the level tilemap and tilebank diverge.
    """
    screen = mem[0x0400:0x0400 + 1000]   # 40x25
    color  = mem[0xD800:0xD800 + 1000]   # 40x25, low nybble = palette slot
    codes  = sorted(Counter(screen).items(), key=lambda kv: (-kv[1], kv[0]))

    # Full 25-row C64 screen. The title screen uses all 25 rows
    # (PUBLISHED ... SOFTWARE credit wraps across rows 21, 22, 24
    # with frame chars interleaved). Gameplay HUD draws over the
    # bottom 3 rows at runtime.
    playfield_rows = 25

    lines = []
    lines.append(f"# Extracted from VICE memory dump. Edit at your own risk;")
    lines.append(f"# regenerating will overwrite. Keep the @tile mappings in")
    lines.append(f"# sync with the tilebank PNG (extracted from $2800).")
    lines.append(f"")
    lines.append(f"@name {level_name.upper()}")
    lines.append(f"@tilebank {tilebank_id}")
    lines.append(f"@music 0")
    bg = mem[0xD021] & 0x0F
    border = mem[0xD020] & 0x0F
    lines.append(f"@bgColor {bg}")
    lines.append(f"@borderColor {border}")

    # Best-effort taxi spawn: middle-top of playfield. The original
    # script-driven spawn position would need to be lifted out of
    # game state at dump time; for now hand-edit if you want it
    # to match the cracker's preferred spawn.
    lines.append(f"@taxiSpawn 18 2")
    lines.append(f"")

    # Pads: scan the playfield for "landable surfaces" -- runs of
    # >= 3 identical non-background tile cells with two or more
    # empty rows directly above (so the taxi can approach from the
    # top). Per-row scanning runs top-to-bottom; for the same
    # surface tile across rows we only take the topmost contiguous
    # platform edge.
    bg_code = codes[0][0] if codes else 0x20
    def is_bg(code):
        return code == bg_code or code == 0x20

    pads = []
    seen_used = [False] * (playfield_rows * 40)
    for row in range(2, playfield_rows):
        line = screen[row * 40:(row + 1) * 40]
        prev = screen[(row - 1) * 40:row * 40]
        prev2 = screen[(row - 2) * 40:(row - 1) * 40]
        col = 0
        while col < 40 and len(pads) < 8:
            if is_bg(line[col]):
                col += 1
                continue
            j = col
            while j < 40 and line[j] == line[col] and not seen_used[row * 40 + j]:
                j += 1
            width = j - col
            # Approachable from above? At least 2/3 of cells in the
            # 2 rows above this stretch must be background.
            above_clear = 0
            for k in range(col, j):
                if is_bg(prev[k]):
                    above_clear += 1
                if is_bg(prev2[k]):
                    above_clear += 1
            if width >= 3 and above_clear >= width:
                pads.append((row, col, width))
                for k in range(col, j):
                    seen_used[row * 40 + k] = True
            col = j

    lines.append("# Pads auto-detected from landable surfaces (>= 3 wide,")
    lines.append("# approachable from above). Verify and adjust as needed.")
    pad_letter = ord("A")
    for (row, col, width) in pads:
        if pad_letter > ord("H"):
            break
        lines.append(f"@pad {chr(pad_letter)} {col} {row} {width}")
        pad_letter += 1
    if pad_letter == ord("A"):
        lines.append("# (No landable surfaces auto-detected; emitting placeholders.)")
        lines.append("@pad A 4 10 4")
        lines.append("@pad B 32 10 4")
        pad_letter = ord("C")

    lines.append("")
    # Fare list: simple round-robin between detected pads.
    pad_count = pad_letter - ord("A")
    if pad_count >= 2:
        for k in range(pad_count):
            src = chr(ord("A") + k)
            dst = chr(ord("A") + ((k + 1) % pad_count))
            lines.append(f"@fare {src} {dst} 45")
    else:
        lines.append("@fare A B 45")
    lines.append("")

    # Emit the @tile mapping for round-trip clarity (and so the
    # author can read a level .txt and tell which screen code maps
    # to which char). Format: @tile <char> <screencode> -- the
    # mkstlevel parser ignores these today but they're useful
    # documentation, and a future loader could honor them.
    lines.append("# Screen-code -> alphabet-char mapping (one per unique code).")
    inv = {v: k for k, v in code_to_alpha.items()}
    for idx in sorted(inv):
        ch = '.' if idx == 0 else ALPHABET[idx]
        code = inv[idx]
        lines.append(f"# @tile '{ch}' = $C64_screencode_{code:02X}")
    lines.append("")

    # Tilemap (22 rows of 40 chars each).
    lines.append("@tilemap")
    for r in range(playfield_rows):
        row = screen[r * 40:(r + 1) * 40]
        s = ""
        for code in row:
            idx = code_to_alpha.get(code, 0)
            s += '.' if idx == 0 else ALPHABET[idx]
        lines.append(s)

    lines.append("")
    lines.append("@colormap")
    for r in range(playfield_rows):
        row = color[r * 40:(r + 1) * 40]
        s = "".join(
            ('0' if (c & 0x0F) == 0
             else ALPHABET[(c & 0x0F)]
             if (c & 0x0F) < len(ALPHABET) else '0')
            for c in row
        )
        lines.append(s)

    with open(out_path, "w") as fp:
        fp.write("\n".join(lines) + "\n")


def extract_sprites(mem, out_path):
    """Render C64 hardware sprites into the JoeyLib-expected sprite
    sheet layout (320x200, 16-color indexed PNG):
        y=  0..23  : 1 taxi cel (24x24, x = 0)
                       sprite 0 ptr (in level1 dump: $DC -> $3700)
                     The adjacent slot $DD is the landing-gear-down
                     pose, not a thrust frame -- ignored for now.
        y= 24..47  : 4 passenger cels (16x16 each, x = 0, 16, 32, 48)
    Sprite 3 ptr ($07FB) is the passenger data; sprites 3-6 share it.
    """
    img = new_canvas(320, 200)
    px = img.load()

    # Shared multicolor sprite colors from VIC registers.
    mcm_color0 = mem[0xD025] & 0x0F     # 2-bit value 01
    mcm_color1 = mem[0xD026] & 0x0F     # 2-bit value 11
    mcm_mask   = mem[0xD01C]            # bit per sprite: 1 = multicolor

    def paint_sprite_at(data_addr, sprite_color, dst_x, dst_y, dst_w, dst_h, multicolor=False):
        # Mono C64 sprite is 24x21, 1 bit per pixel, 3 bytes per row.
        # Multicolor C64 sprite is 12 "double-wide pixels" x 21 rows,
        # 2 bits per pixel, still 3 bytes per row. Color mapping:
        #   00 -> transparent
        #   01 -> mcm_color0 ($D025)
        #   10 -> sprite_color (sprite's own $D027+sp)
        #   11 -> mcm_color1 ($D026)
        if sprite_color == 0:
            sprite_color = 1
        sx_max = min(24, dst_w)
        sy_max = min(21, dst_h)
        for r in range(sy_max):
            if multicolor:
                for byte_idx in range(3):
                    byte = mem[data_addr + r * 3 + byte_idx]
                    # 4 multicolor pixels per byte; each spans 2
                    # hardware pixels wide.
                    for pp in range(4):
                        bits = (byte >> ((3 - pp) * 2)) & 0x03
                        if bits == 0:
                            continue
                        if bits == 1:
                            col = mcm_color0
                        elif bits == 2:
                            col = sprite_color
                        else:
                            col = mcm_color1
                        sx0 = byte_idx * 8 + pp * 2
                        if sx0 + 1 >= sx_max:
                            break
                        px[dst_x + sx0,     dst_y + r] = col
                        px[dst_x + sx0 + 1, dst_y + r] = col
            else:
                for byte_idx in range(3):
                    byte = mem[data_addr + r * 3 + byte_idx]
                    for bit in range(8):
                        sx = byte_idx * 8 + bit
                        if sx >= sx_max:
                            break
                        if byte & (0x80 >> bit):
                            px[dst_x + sx, dst_y + r] = sprite_color

    def paint_c64_sprite(sp_index, dst_x, dst_y, dst_w, dst_h):
        ptr = mem[0x07F8 + sp_index]
        if ptr == 0:
            return False
        sprite_color = mem[0xD027 + sp_index] & 0x0F
        is_multicolor = bool(mcm_mask & (1 << sp_index))
        paint_sprite_at(ptr * 64, sprite_color, dst_x, dst_y, dst_w, dst_h,
                        multicolor=is_multicolor)
        return True

    # Single taxi cel from sprite 0's actual pointer.
    paint_c64_sprite(0, 0, 0, 24, 24)

    # 4 passenger cels = the shared C64 sprite 3 (sprites 3..6 all
    # point to the same data in the level-1 dump).
    for cel in range(4):
        paint_c64_sprite(3, cel * 16, 24, 16, 16)

    img.save(out_path, "PNG", optimize=True)


def extract_font(mem, out_path):
    """Render the custom charset (which IS the in-game font in C64
    text-mode games) as the 320x200 font.png. Same layout as
    tilebank, but here we want each glyph at cell (ascii%40,
    ascii/40) for jlDrawText's ASCII map.

    For Space Taxi specifically, the charset doubles as the level
    tile bank AND any in-screen text Sierra-style. We map the C64
    screen codes used for ASCII display (codes 1..26 = A..Z,
    32..63 = punctuation/digits) directly to their ASCII positions
    in the JoeyLib font sheet.
    """
    img = new_canvas(320, 200)
    px = img.load()

    def render_at(ascii_code, charset_idx):
        if ascii_code < 32 or ascii_code >= 128:
            return
        src = 0x2800 + (charset_idx & 0xFF) * 8
        col = ascii_code % 40
        row = ascii_code // 40
        bx = col * 8
        by = row * 8
        for r in range(8):
            byte = mem[src + r]
            for c in range(8):
                if byte & (0x80 >> c):
                    px[bx + c, by + r] = 1

    # C64 upper-only screen codes: 1..26 = A..Z, 27=[ 28=£ 29=] 30=↑ 31=←
    # 32 = space, 33..63 = !"# ... ?, etc. We map ASCII directly.
    for c in range(32, 64):
        render_at(c, c)                       # symbols + digits 1:1
    for c, code in enumerate("ABCDEFGHIJKLMNOPQRSTUVWXYZ"):
        render_at(ord(code), 1 + c)           # screen code 1..26 -> A..Z

    img.save(out_path, "PNG", optimize=True)


def main():
    if len(sys.argv) != 4:
        print("usage: extractFromDump.py <dump-path> <level-name> <tilebank-id>", file=sys.stderr)
        sys.exit(2)
    dump_path     = sys.argv[1]
    level_name    = sys.argv[2]
    tilebank_id   = int(sys.argv[3])

    mem = load_dump(dump_path)
    print(f"loaded {dump_path}: 64 KB OK")

    tiles_dir   = os.path.join(ROOT, "assets", "tiles")
    sprites_dir = os.path.join(ROOT, "assets", "sprites")
    levels_dir  = os.path.join(ROOT, "assets", "levels")
    os.makedirs(tiles_dir,   exist_ok=True)
    os.makedirs(sprites_dir, exist_ok=True)
    os.makedirs(levels_dir,  exist_ok=True)

    # Build the screen-code -> bank-slot mapping once and feed it to
    # both extractors so the tilemap and tilebank stay aligned.
    code_to_slot = build_code_to_slot(mem)

    # 8.3 DOS naming -- "tilebank0" is 9 chars and fopen fails under
    # DOSBox strict 8.3 mode. Keep the short form across the pipeline.
    tilebank_path = os.path.join(tiles_dir, f"tbank{tilebank_id}.png")
    extract_tilebank(mem, tilebank_path, code_to_slot)
    print(f"  wrote {tilebank_path}")

    level_path = os.path.join(levels_dir, f"{level_name}.txt")
    extract_level(mem, level_name, tilebank_id, level_path, code_to_slot)
    print(f"  wrote {level_path}")

    # Sprites + font come from level 1 only (sprite pointer table is
    # script-driven so it differs by scene; we use the first dump's
    # set as the canonical sheet, plus its charset as the font).
    if tilebank_id == 1:
        sprite_path = os.path.join(sprites_dir, "sprites.png")
        extract_sprites(mem, sprite_path)
        print(f"  wrote {sprite_path}")

        font_path = os.path.join(ROOT, "assets", "font.png")
        extract_font(mem, font_path)
        print(f"  wrote {font_path}")


if __name__ == "__main__":
    main()