#!/usr/bin/env python3
# extractFlame.py -- wire the engine-flame sprite cels from the C64 ROM
# image (stuff/spacetaxi/raw.bin) into the sprite sheet's flame row.
#
# The flame is C64 sprite 2, MULTICOLOR (verified: $D01C bit2 set). Its
# per-direction cel is chosen from the pointer table at $6DB0 (indexed by
# the direction mask $716A). Reading raw.bin at $6DB0 (2-byte load-addr
# header -> +2) gives, and the port's kFlameCelByDirMask mirrors exactly:
#   mask 1  UP        -> block $D8      mask 8  RIGHT     -> block $D7
#   mask 2  DOWN      -> block $D4      mask 9  UP+RIGHT  -> block $D3
#   mask 4  LEFT      -> block $D5      mask 10 DOWN+RIGHT-> block $D6
#   mask 5  UP+LEFT   -> block $D2
#   mask 6  DOWN+LEFT -> block $D1
# flameCels[0..7] load from .spr cels 18..25 (sheet row 2, cols 0..7), so
# the cel order below is the kFlameCelByDirMask index order.
#
# Sprite data lives at block*64 in VIC bank 0 ($3000-$37FF). Each C64
# multicolor sprite is 12 MC-pixels x 21 rows (3 bytes/row, 2 bits/pixel);
# each MC pixel is 2 hi-res pixels, so it maps 1 MC pixel -> 2 of the
# port's 24-wide 4bpp cel pixels. Colors are the live-level flame palette
# ($7D05->$D025=$02 red = MC0, $7D06->$D026=$07 yellow = MC1; sprite-2
# colour = orange), mapped to sprites.png's indexed palette.
#
# Patches ONLY the flame row (row 2) of sprites.png -- taxi (row 0),
# passenger (row 1), and font.png are left untouched. Re-run the .spr
# bake afterwards (make/<port>.mk STAXI_SPR rule).

import os
import sys

try:
    from PIL import Image
except ImportError:
    sys.exit("extractFlame.py: needs Pillow (pip install pillow)")

HERE = os.path.dirname(os.path.abspath(__file__))
RAW = os.path.join(HERE, "..", "..", "..", "stuff", "spacetaxi", "raw.bin")
PNG = os.path.join(HERE, "sprites", "sprites.png")

CEL = 24                 # cel is 24x24 px
LOAD_HDR = 2             # raw.bin 2-byte load-address header

# flameCels index -> C64 sprite block (from $6DB0 / kFlameCelByDirMask).
FLAME_BLOCKS = [0xD8, 0xD4, 0xD5, 0xD2, 0xD1, 0xD7, 0xD3, 0xD6]

# Multicolor 2-bit pixel -> palette index. The engine renders sprite
# cels through kDefaultPalette (stRender.c), whose index IS the C64 VIC-II
# colour code -- the sprite's own baked palette is ignored (loadSpriteSheet
# passes NULL). So the C64 flame colours map DIRECTLY by code:
#   00 transparent(0), 01 MC0=$D025=$02 red(2), 11 MC1=$D026=$07 yellow(7),
#   10 sprite-2 colour -> orange(8) (barely used by these cels).
# (NOT the sprites.png preview-palette indices, which differ.)
MC_TO_INDEX = {0: 0, 1: 2, 2: 8, 3: 7}

FLAME_ROW_Y = 2 * CEL    # sheet row 2 (flame row) top edge


def loadRaw():
    with open(RAW, "rb") as f:
        return f.read()


def celFromBlock(raw, block):
    # Return a 24x24 list-of-rows of palette indices for one flame block.
    base = LOAD_HDR + block * 64
    cel = [[0] * CEL for _ in range(CEL)]
    for row in range(21):                    # C64 sprite is 21 rows tall
        for byteIdx in range(3):             # 3 bytes/row
            b = raw[base + row * 3 + byteIdx]
            for pair in range(4):            # 4 MC pixels per byte (MSB first)
                mc = (b >> ((3 - pair) * 2)) & 3
                idx = MC_TO_INDEX[mc]
                px = (byteIdx * 4 + pair) * 2  # MC pixel -> 2 cel pixels
                cel[row][px] = idx
                cel[row][px + 1] = idx
    return cel


def main():
    raw = loadRaw()
    im = Image.open(PNG)
    if im.mode != "P":
        sys.exit(f"extractFlame.py: {PNG} is {im.mode}, expected indexed 'P'")
    px = im.load()
    nonEmpty = 0
    for celIdx, block in enumerate(FLAME_BLOCKS):
        cel = celFromBlock(raw, block)
        x0 = celIdx * CEL
        opaque = 0
        for y in range(CEL):
            for x in range(CEL):
                v = cel[y][x]
                px[x0 + x, FLAME_ROW_Y + y] = v
                if v != 0:
                    opaque += 1
        nonEmpty += 1 if opaque else 0
        print(f"  flameCels[{celIdx}] block ${block:02X}: {opaque} opaque px -> cel {18 + celIdx}")
    im.save(PNG)
    print(f"extractFlame: patched {nonEmpty}/8 flame cels into {os.path.relpath(PNG, HERE)}")


if __name__ == "__main__":
    main()