"""
prepare_win_gfx.py

Convert Doom intermission label PNGs (win_gfx/) to VB 2bpp tile + map arrays.
Each image is 12px tall (2 tile rows, padded), variable width.

Output: src/vbdoom/assets/images/win_gfx.c and win_gfx.h
"""

import os
from PIL import Image

SCRIPT_DIR = os.path.dirname(os.path.abspath(__file__))
INPUT_DIR = os.path.join(SCRIPT_DIR, "win_gfx")
OUTPUT_C = os.path.join(SCRIPT_DIR, "src", "vbdoom", "assets", "images", "win_gfx.c")
OUTPUT_H = os.path.join(SCRIPT_DIR, "src", "vbdoom", "assets", "images", "win_gfx.h")

TILE_W = 8
TILE_H = 8

# Images to convert: (filename, C prefix)
LABELS = [
    ("WITIME.png",   "witime"),
    ("WISCRT2.png",  "wiscrt"),
    ("WIOSTK.png",   "wiostk"),
    ("WIOSTI.png",   "wiosti"),
    ("WIENTER.png",  "wienter"),
    ("WIF.png",      "wif"),
    ("WIFRGS.png",   "wifrgs"),
    ("WIMSTT.png",   "wimstt"),
    ("WIPCNT.png",   "wipcnt"),
]


def build_palette_remap(img):
    """Map PNG palette indices to VB brightness levels (0-3)."""
    pal = img.getpalette()
    if pal is None:
        # Not indexed -- convert from RGB
        return None

    used_indices = sorted(set(img.getdata()))
    idx_brightness = []
    for idx in used_indices:
        r = pal[idx * 3]
        idx_brightness.append((idx, r))

    idx_brightness.sort(key=lambda x: x[1])
    remap = {}
    for vb_level, (png_idx, _) in enumerate(idx_brightness):
        if vb_level > 3:
            vb_level = 3
        remap[png_idx] = vb_level

    return remap


def pixel_to_vb(img, px, py, remap):
    """Get VB brightness for pixel at (px, py)."""
    if px >= img.width or py >= img.height:
        return 0  # padding
    val = img.getpixel((px, py))
    if remap is not None:
        return remap.get(val, 0)
    # RGB mode
    if isinstance(val, tuple):
        r = val[0]
        if r < 32:
            return 0
        elif r < 96:
            return 1
        elif r < 192:
            return 2
        else:
            return 3
    return val & 3


def encode_tile_words(pixels):
    """Convert 64-pixel list to 4 VB 2bpp u32 words (sequential format)."""
    words = []
    for pair in range(4):
        row_a = pair * 2
        row_b = pair * 2 + 1
        u16_a = 0
        u16_b = 0
        for x in range(TILE_W):
            u16_a |= (pixels[row_a * TILE_W + x] << (x * 2))
            u16_b |= (pixels[row_b * TILE_W + x] << (x * 2))
        word = u16_a | (u16_b << 16)
        words.append(word)
    return words


def convert_image(filepath, prefix):
    """Convert one label PNG to tile data + map entries."""
    img = Image.open(filepath)
    w, h = img.size

    # Pad height to multiple of 8 (tiles are 8x8)
    tile_cols = (w + TILE_W - 1) // TILE_W
    tile_rows = (h + TILE_H - 1) // TILE_H

    remap = None
    if img.mode == 'P':
        remap = build_palette_remap(img)
    elif img.mode == 'RGBA':
        img = img.convert('RGB')

    # Build tile bank with dedup
    tile_bank = []      # list of word-tuples
    tile_lookup = {}    # pixel_tuple -> index
    map_entries = []    # one u16 per cell

    # Tile 0 = blank
    blank = tuple([0] * 64)
    blank_words = tuple(encode_tile_words(list(blank)))
    tile_bank.append(blank_words)
    tile_lookup[blank] = 0

    for ty in range(tile_rows):
        for tx in range(tile_cols):
            pixels = []
            for py in range(TILE_H):
                for px in range(TILE_W):
                    pixels.append(pixel_to_vb(img, tx * TILE_W + px, ty * TILE_H + py, remap))
            pixels = tuple(pixels)

            if pixels in tile_lookup:
                map_entries.append(tile_lookup[pixels])
            else:
                idx = len(tile_bank)
                tile_bank.append(tuple(encode_tile_words(list(pixels))))
                tile_lookup[pixels] = idx
                map_entries.append(idx)

    print(f"  {os.path.basename(filepath)}: {w}x{h} -> {tile_cols}x{tile_rows} tiles, "
          f"{len(tile_bank)} unique, {len(tile_bank)*16} bytes")

    return tile_bank, map_entries, tile_cols, tile_rows


def main():
    print("Converting win_gfx label images to VB 2bpp tile data...\n")

    all_results = []
    for fname, prefix in LABELS:
        fpath = os.path.join(INPUT_DIR, fname)
        if not os.path.isfile(fpath):
            print(f"  WARNING: {fpath} not found, skipping")
            continue
        result = convert_image(fpath, prefix)
        all_results.append((prefix, result))

    # Write C file
    with open(OUTPUT_C, 'w') as f:
        f.write("/* win_gfx -- Doom intermission label tiles\n")
        f.write(" * Auto-generated by prepare_win_gfx.py\n */\n\n")

        for prefix, (tile_bank, map_entries, tcols, trows) in all_results:
            num_words = len(tile_bank) * 4
            all_words = [w for tile in tile_bank for w in tile]

            f.write(f"/* {prefix}: {tcols}x{trows} tiles, {len(tile_bank)} unique */\n")
            f.write(f"const unsigned int {prefix}Tiles[{num_words}]"
                    f" __attribute__((aligned(4))) =\n{{\n")
            for i, word in enumerate(all_words):
                if i % 8 == 0:
                    f.write("\t")
                f.write(f"0x{word:08X}")
                if i < num_words - 1:
                    f.write(",")
                if i % 8 == 7:
                    f.write("\n")
            if num_words % 8 != 0:
                f.write("\n")
            f.write("};\n\n")

            f.write(f"const unsigned short {prefix}Map[{len(map_entries)}]"
                    f" __attribute__((aligned(4))) =\n{{\n")
            for i, entry in enumerate(map_entries):
                if i % 8 == 0:
                    f.write("\t")
                f.write(f"0x{entry:04X}")
                if i < len(map_entries) - 1:
                    f.write(",")
                if i % 8 == 7:
                    f.write("\n")
            if len(map_entries) % 8 != 0:
                f.write("\n")
            f.write("};\n\n")

    # Write H file
    with open(OUTPUT_H, 'w') as f:
        f.write("#ifndef __WIN_GFX_H__\n#define __WIN_GFX_H__\n\n")
        for prefix, (tile_bank, map_entries, tcols, trows) in all_results:
            num_words = len(tile_bank) * 4
            f.write(f"#define {prefix.upper()}_TILES  {len(tile_bank)}\n")
            f.write(f"#define {prefix.upper()}_TILE_BYTES  {len(tile_bank)*16}\n")
            f.write(f"#define {prefix.upper()}_COLS  {tcols}\n")
            f.write(f"#define {prefix.upper()}_ROWS  {trows}\n")
            f.write(f"extern const unsigned int {prefix}Tiles[{num_words}];\n")
            f.write(f"extern const unsigned short {prefix}Map[{len(map_entries)}];\n\n")
        f.write("#endif\n")

    print(f"\nWritten: {OUTPUT_C}")
    print(f"Written: {OUTPUT_H}")


if __name__ == "__main__":
    main()