pokemon_pinball_wiki/scripts/export_radium_section3_bitmaps.py

#!/usr/bin/env python3
"""Export preview PNGs for simple Radium section 3 bitmap records.

Ghidra xrefs show section 3 is a qword pointer table with count id 1. The
records addressed by that table start with:

    uint32 record_id
    uint32 flags
    uint16 width
    uint16 height

Most small records then contain width*height bytes of indexed pixel data plus
padding. This exporter writes grayscale/alpha PNG previews only for those
straightforward records and leaves compact 128x32-style payloads as undecoded.
It is static and does not execute target binaries.
"""

from __future__ import annotations

import argparse
import csv
import json
import struct
import zlib
from pathlib import Path


PNG_SIGNATURE = b"\x89PNG\r\n\x1a\n"


FIELDS = [
    "logical_index",
    "record_id",
    "flags",
    "record_offset",
    "next_record_offset",
    "record_length",
    "width",
    "height",
    "pixel_payload_length",
    "decoder_status",
    "png_path",
]


def u64(data: bytes, offset: int) -> int:
    return struct.unpack_from("<Q", data, offset)[0]


def png_chunk(kind: bytes, payload: bytes) -> bytes:
    crc = zlib.crc32(kind + payload) & 0xFFFFFFFF
    return struct.pack(">I", len(payload)) + kind + payload + struct.pack(">I", crc)


def write_rgba_png(path: Path, width: int, height: int, rgba: bytes) -> None:
    rows = []
    stride = width * 4
    for y in range(height):
        rows.append(b"\x00" + rgba[y * stride : (y + 1) * stride])
    ihdr = struct.pack(">IIBBBBB", width, height, 8, 6, 0, 0, 0)
    path.write_bytes(
        PNG_SIGNATURE
        + png_chunk(b"IHDR", ihdr)
        + png_chunk(b"IDAT", zlib.compress(b"".join(rows), level=9))
        + png_chunk(b"IEND", b"")
    )


def pixel_to_rgba(value: int) -> tuple[int, int, int, int]:
    if value == 0xFF:
        return (0, 0, 0, 0)
    if value <= 0x0F:
        gray = value * 17
    else:
        gray = value
    return (gray, gray, gray, 255)


def read_header(path: Path) -> dict[str, object]:
    with path.open(encoding="utf-8") as f:
        return json.load(f)


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--image", type=Path, default=Path("games/pokemon_pro/image.bin"))
    parser.add_argument("--header", type=Path, default=Path("analysis/radium-image/header.json"))
    parser.add_argument("--out-dir", type=Path, default=Path("analysis/radium-image/section3-bitmaps"))
    parser.add_argument(
        "--manifest",
        type=Path,
        default=Path("analysis/radium-image/section3_bitmap_manifest.tsv"),
    )
    args = parser.parse_args()

    data = args.image.read_bytes()
    header = read_header(args.header)
    section_offsets = {key: int(value) for key, value in header["section_offsets"].items()}
    section3_offset = section_offsets["section_3"]
    section8_offset = section_offsets["section_8"]
    count = int(header["counts_from_header"]["count_1_header10_high32"])
    pointers = [u64(data, section3_offset + index * 8) for index in range(count)]
    sorted_pointers = sorted(set(pointers))
    next_by_pointer = {
        pointer: sorted_pointers[index + 1] if index + 1 < len(sorted_pointers) else section8_offset
        for index, pointer in enumerate(sorted_pointers)
    }

    args.out_dir.mkdir(parents=True, exist_ok=True)
    rows: list[dict[str, object]] = []
    exported = 0
    for logical_index, pointer in enumerate(pointers):
        next_pointer = next_by_pointer[pointer]
        record_length = next_pointer - pointer
        row = {
            "logical_index": logical_index,
            "record_id": "",
            "flags": "",
            "record_offset": f"0x{pointer:x}",
            "next_record_offset": f"0x{next_pointer:x}",
            "record_length": record_length,
            "width": "",
            "height": "",
            "pixel_payload_length": "",
            "decoder_status": "truncated_header",
            "png_path": "",
        }
        if pointer + 12 > len(data) or record_length < 12:
            rows.append(row)
            continue

        record_id, flags = struct.unpack_from("<II", data, pointer)
        width, height = struct.unpack_from("<HH", data, pointer + 8)
        area = width * height
        payload_length = max(0, record_length - 12)
        row.update(
            {
                "record_id": record_id,
                "flags": flags,
                "width": width,
                "height": height,
                "pixel_payload_length": payload_length,
            }
        )
        if width == 0 or height == 0 or width > 512 or height > 512:
            row["decoder_status"] = "invalid_dimensions"
        elif payload_length < area:
            row["decoder_status"] = "undecoded_compact_payload"
        else:
            pixels = data[pointer + 12 : pointer + 12 + area]
            rgba = bytearray()
            for value in pixels:
                rgba.extend(pixel_to_rgba(value))
            png_path = args.out_dir / f"section3_{logical_index:04d}_id{record_id:04d}_{width}x{height}_f{flags}.png"
            write_rgba_png(png_path, width, height, bytes(rgba))
            row["decoder_status"] = "raw_indexed_preview"
            row["png_path"] = png_path.as_posix()
            exported += 1
        rows.append(row)

    args.manifest.parent.mkdir(parents=True, exist_ok=True)
    with args.manifest.open("w", newline="", encoding="utf-8") as f:
        writer = csv.DictWriter(f, fieldnames=FIELDS, delimiter="\t", extrasaction="ignore")
        writer.writeheader()
        writer.writerows(rows)
    status_counts: dict[str, int] = {}
    for row in rows:
        status = str(row["decoder_status"])
        status_counts[status] = status_counts.get(status, 0) + 1
    print(f"wrote {args.manifest} rows={len(rows)} exported_pngs={exported}")
    for status, count_value in sorted(status_counts.items()):
        print(f"{status}: {count_value}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())