#!/usr/bin/env python3

import argparse
import math
from pathlib import Path

from PIL import Image, ImageDraw, ImageFont


# =============================================================================
# NES CHR constants
# =============================================================================

CHR_BANK_SIZE = 0x2000
PATTERN_TABLE_SIZE = 0x1000

TILE_SIZE_BYTES = 16
TILE_WIDTH = 8
TILE_HEIGHT = 8

TILES_PER_PATTERN_TABLE = 256


# =============================================================================
# NES RGB palette
#
# Common RGB approximation of the NES 64-color system palette.
# CHR-ROM itself does NOT store these colors. CHR pixels contain values 0-3,
# which are mapped through NES palette RAM into one of these 64 colors.
# =============================================================================

NES_RGB = [
    # $00-$0F
    (84, 84, 84),
    (0, 30, 116),
    (8, 16, 144),
    (48, 0, 136),
    (68, 0, 100),
    (92, 0, 48),
    (84, 4, 0),
    (60, 24, 0),
    (32, 42, 0),
    (8, 58, 0),
    (0, 64, 0),
    (0, 60, 0),
    (0, 50, 60),
    (0, 0, 0),
    (0, 0, 0),
    (0, 0, 0),

    # $10-$1F
    (152, 150, 152),
    (8, 76, 196),
    (48, 50, 236),
    (92, 30, 228),
    (136, 20, 176),
    (160, 20, 100),
    (152, 34, 32),
    (120, 60, 0),
    (84, 90, 0),
    (40, 114, 0),
    (8, 124, 0),
    (0, 118, 40),
    (0, 102, 120),
    (0, 0, 0),
    (0, 0, 0),
    (0, 0, 0),

    # $20-$2F
    (236, 238, 236),
    (76, 154, 236),
    (120, 124, 236),
    (176, 98, 236),
    (228, 84, 236),
    (236, 88, 180),
    (236, 106, 100),
    (212, 136, 32),
    (160, 170, 0),
    (116, 196, 0),
    (76, 208, 32),
    (56, 204, 108),
    (56, 180, 204),
    (60, 60, 60),
    (0, 0, 0),
    (0, 0, 0),

    # $30-$3F
    (236, 238, 236),
    (168, 204, 236),
    (188, 188, 236),
    (212, 178, 236),
    (236, 174, 236),
    (236, 174, 212),
    (236, 180, 176),
    (228, 196, 144),
    (204, 210, 120),
    (180, 222, 120),
    (168, 226, 144),
    (152, 226, 180),
    (160, 214, 228),
    (160, 162, 160),
    (0, 0, 0),
    (0, 0, 0),
]


GRAYSCALE = [
    (0, 0, 0),
    (85, 85, 85),
    (170, 170, 170),
    (255, 255, 255),
]


# =============================================================================
# CHR decoder
# =============================================================================

def decode_tile(chr_data, tile_index):
    """
    Decode one NES 8x8 2bpp CHR tile.

    Each tile occupies 16 bytes:

        bytes 0-7   bitplane 0
        bytes 8-15  bitplane 1

    Each resulting pixel has a value from 0 through 3.
    """

    offset = tile_index * TILE_SIZE_BYTES
    tile = chr_data[offset:offset + TILE_SIZE_BYTES]

    if len(tile) != TILE_SIZE_BYTES:
        raise ValueError(
            f"Incomplete CHR tile ${tile_index:03X}"
        )

    pixels = []

    for y in range(TILE_HEIGHT):
        plane0 = tile[y]
        plane1 = tile[y + 8]

        row = []

        for x in range(TILE_WIDTH):
            bit = 7 - x

            low = (plane0 >> bit) & 1
            high = (plane1 >> bit) & 1

            pixel = low | (high << 1)

            row.append(pixel)

        pixels.append(row)

    return pixels


# =============================================================================
# Palette handling
# =============================================================================

def nes_palette_to_rgb(indices):
    """
    Convert four NES palette indices into RGB colors.
    """

    if len(indices) != 4:
        raise ValueError(
            "Palette must contain exactly four NES colors"
        )

    return [
        NES_RGB[index & 0x3F]
        for index in indices
    ]


def parse_hex_byte(value):
    """
    Parse a command-line hexadecimal byte.

    Accepted examples:

        0F
        FF
        0x20
        $26
    """

    value = value.strip()

    if value.startswith("$"):
        value = value[1:]

    if value.lower().startswith("0x"):
        value = value[2:]

    number = int(value, 16)

    if not 0 <= number <= 0xFF:
        raise argparse.ArgumentTypeError(
            f"Value out of byte range: {value}"
        )

    return number


# =============================================================================
# Tile rendering
# =============================================================================

def make_tile_image(tile, scale, colors=None):
    """
    Render one decoded 8x8 tile.

    colors[0..3] maps CHR pixel values 0..3 to RGB colors.
    """

    if colors is None:
        colors = GRAYSCALE

    image = Image.new(
        "RGB",
        (TILE_WIDTH, TILE_HEIGHT),
        colors[0]
    )

    pixels = image.load()

    for y in range(TILE_HEIGHT):
        for x in range(TILE_WIDTH):
            pixels[x, y] = colors[tile[y][x]]

    if scale != 1:
        image = image.resize(
            (
                TILE_WIDTH * scale,
                TILE_HEIGHT * scale
            ),
            Image.Resampling.NEAREST
        )

    return image


# =============================================================================
# Clean sprite sheet
# =============================================================================

def make_clean_sheet(
    chr_data,
    columns=16,
    scale=4,
    colors=None
):
    """
    Render CHR tiles as a compact sprite sheet.
    """

    tile_count = len(chr_data) // TILE_SIZE_BYTES
    rows = math.ceil(tile_count / columns)

    rendered_width = TILE_WIDTH * scale
    rendered_height = TILE_HEIGHT * scale

    sheet = Image.new(
        "RGB",
        (
            columns * rendered_width,
            rows * rendered_height
        ),
        (0, 0, 0)
    )

    for tile_index in range(tile_count):
        tile = decode_tile(
            chr_data,
            tile_index
        )

        tile_image = make_tile_image(
            tile,
            scale,
            colors
        )

        column = tile_index % columns
        row = tile_index // columns

        x = column * rendered_width
        y = row * rendered_height

        sheet.paste(
            tile_image,
            (x, y)
        )

    return sheet


# =============================================================================
# Labeled sprite sheet
# =============================================================================

def make_labeled_sheet(
    chr_data,
    columns=16,
    tile_scale=4,
    label_height=12,
    tile_base=0,
    colors=None
):
    """
    Render a debugging sheet with each tile labeled by hexadecimal tile ID.
    """

    tile_count = len(chr_data) // TILE_SIZE_BYTES
    rows = math.ceil(tile_count / columns)

    rendered_width = TILE_WIDTH * tile_scale
    rendered_height = TILE_HEIGHT * tile_scale

    cell_width = max(
        rendered_width + 4,
        40
    )

    cell_height = (
        rendered_height +
        label_height +
        6
    )

    sheet = Image.new(
        "RGB",
        (
            columns * cell_width,
            rows * cell_height
        ),
        (24, 24, 24)
    )

    draw = ImageDraw.Draw(sheet)
    font = ImageFont.load_default()

    for tile_index in range(tile_count):
        tile = decode_tile(
            chr_data,
            tile_index
        )

        tile_image = make_tile_image(
            tile,
            tile_scale,
            colors
        )

        column = tile_index % columns
        row = tile_index // columns

        cell_x = column * cell_width
        cell_y = row * cell_height

        tile_x = (
            cell_x +
            (cell_width - rendered_width) // 2
        )

        tile_y = cell_y + 2

        sheet.paste(
            tile_image,
            (tile_x, tile_y)
        )

        display_index = tile_base + tile_index

        label = f"{display_index:03X}"

        bbox = draw.textbbox(
            (0, 0),
            label,
            font=font
        )

        text_width = bbox[2] - bbox[0]

        text_x = (
            cell_x +
            (cell_width - text_width) // 2
        )

        text_y = (
            cell_y +
            rendered_height +
            4
        )

        draw.text(
            (text_x, text_y),
            label,
            fill=(255, 255, 255),
            font=font
        )

    return sheet


# =============================================================================
# iNES extraction
# =============================================================================

def extract_chr_banks_from_ines(rom):
    """
    Extract every 8 KiB CHR-ROM bank from an iNES ROM.
    """

    if len(rom) < 16:
        raise ValueError(
            "File is too small to contain an iNES header"
        )

    if rom[:4] != b"NES\x1A":
        raise ValueError(
            "Input does not contain a valid iNES header"
        )

    prg_banks = rom[4]
    chr_banks = rom[5]
    flags6 = rom[6]

    if chr_banks == 0:
        raise ValueError(
            "ROM contains no CHR-ROM. "
            "This cartridge may use CHR-RAM."
        )

    trainer_size = (
        512
        if flags6 & 0x04
        else 0
    )

    prg_size = (
        prg_banks * 0x4000
    )

    chr_offset = (
        16 +
        trainer_size +
        prg_size
    )

    banks = []

    for bank_index in range(chr_banks):

        start = (
            chr_offset +
            bank_index * CHR_BANK_SIZE
        )

        end = (
            start +
            CHR_BANK_SIZE
        )

        if end > len(rom):
            raise ValueError(
                f"ROM truncated while reading "
                f"CHR bank ${bank_index:02X}"
            )

        banks.append(
            rom[start:end]
        )

    return banks


# =============================================================================
# Raw CHR extraction
# =============================================================================

def extract_raw_chr_banks(data):
    """
    Split raw CHR data into 8 KiB banks.
    """

    if len(data) % CHR_BANK_SIZE != 0:
        raise ValueError(
            "Raw CHR data size must be a multiple "
            f"of ${CHR_BANK_SIZE:04X} bytes"
        )

    banks = []

    for offset in range(
        0,
        len(data),
        CHR_BANK_SIZE
    ):
        banks.append(
            data[
                offset:
                offset + CHR_BANK_SIZE
            ]
        )

    return banks


# =============================================================================
# GIF output
# =============================================================================

def save_gif(image, filename):
    """
    Save a static GIF.
    """

    image.save(
        filename,
        format="GIF",
        optimize=False
    )

    print(
        f"    {filename} "
        f"({image.width}x{image.height})"
    )


# =============================================================================
# Pattern table generation
# =============================================================================

def process_pattern_table(
    pattern_data,
    output_prefix,
    bank_index,
    table_index,
    scale,
    columns,
    colors,
    suffix=""
):
    """
    Produce clean and labeled sheets for a 4 KiB pattern table.
    """

    tile_base = (
        table_index * 0x100
    )

    base_name = (
        f"{output_prefix}"
        f"_chr{bank_index:02X}"
        f"_pt{table_index}"
        f"{suffix}"
    )

    clean = make_clean_sheet(
        pattern_data,
        columns=columns,
        scale=scale,
        colors=colors
    )

    save_gif(
        clean,
        f"{base_name}.gif"
    )

    labeled = make_labeled_sheet(
        pattern_data,
        columns=columns,
        tile_scale=scale,
        tile_base=tile_base,
        colors=colors
    )

    save_gif(
        labeled,
        f"{base_name}_labels.gif"
    )


# =============================================================================
# Full CHR bank rendering
# =============================================================================

def process_chr_bank(
    bank,
    bank_index,
    output_prefix,
    scale,
    columns,
    colors,
    suffix=""
):
    """
    Generate:
        full 8 KiB CHR bank
        pattern table 0
        pattern table 1
        labeled pattern table 0
        labeled pattern table 1
    """

    full_name = (
        f"{output_prefix}"
        f"_chr{bank_index:02X}"
        f"_full"
        f"{suffix}"
        f".gif"
    )

    full_sheet = make_clean_sheet(
        bank,
        columns=columns,
        scale=scale,
        colors=colors
    )

    save_gif(
        full_sheet,
        full_name
    )

    pattern0 = bank[
        0x0000:
        0x1000
    ]

    pattern1 = bank[
        0x1000:
        0x2000
    ]

    process_pattern_table(
        pattern0,
        output_prefix,
        bank_index,
        0,
        scale,
        columns,
        colors,
        suffix
    )

    process_pattern_table(
        pattern1,
        output_prefix,
        bank_index,
        1,
        scale,
        columns,
        colors,
        suffix
    )


# =============================================================================
# Full 32-byte palette RAM mode
# =============================================================================

def process_palette_ram(
    banks,
    palette_ram,
    output_prefix,
    scale,
    columns
):
    """
    Render each CHR bank through all four background palettes and all four
    sprite palettes contained in a 32-byte NES palette-RAM image.
    """

    if len(palette_ram) != 32:
        raise ValueError(
            "Palette RAM must contain exactly 32 bytes"
        )

    print()
    print("Generating background palette sheets")

    for palette_number in range(4):

        start = palette_number * 4

        palette_indices = (
            palette_ram[
                start:
                start + 4
            ]
        )

        colors = nes_palette_to_rgb(
            palette_indices
        )

        suffix = (
            f"_bg{palette_number}"
        )

        print()
        print(
            f"  BG{palette_number}: " +
            " ".join(
                f"${value:02X}"
                for value
                in palette_indices
            )
        )

        for bank_index, bank in enumerate(banks):
            process_chr_bank(
                bank,
                bank_index,
                output_prefix,
                scale,
                columns,
                colors,
                suffix
            )

    print()
    print("Generating sprite palette sheets")

    for palette_number in range(4):

        start = (
            0x10 +
            palette_number * 4
        )

        palette_indices = (
            palette_ram[
                start:
                start + 4
            ]
        )

        colors = nes_palette_to_rgb(
            palette_indices
        )

        suffix = (
            f"_sp{palette_number}"
        )

        print()
        print(
            f"  SP{palette_number}: " +
            " ".join(
                f"${value:02X}"
                for value
                in palette_indices
            )
        )

        for bank_index, bank in enumerate(banks):
            process_chr_bank(
                bank,
                bank_index,
                output_prefix,
                scale,
                columns,
                colors,
                suffix
            )


# =============================================================================
# Main
# =============================================================================

def main():

    parser = argparse.ArgumentParser(
        description=(
            "NES CHR-ROM GIF visualizer "
            "for reverse engineering"
        )
    )

    parser.add_argument(
        "input",
        help="iNES ROM or raw CHR binary"
    )

    parser.add_argument(
        "--output-prefix",
        default=None,
        help="Output filename prefix"
    )

    parser.add_argument(
        "--raw",
        action="store_true",
        help="Treat input as raw CHR data"
    )

    parser.add_argument(
        "--scale",
        type=int,
        default=4,
        help="Tile scale factor (default: 4)"
    )

    parser.add_argument(
        "--columns",
        type=int,
        default=16,
        help="Tiles per row (default: 16)"
    )

    parser.add_argument(
        "--palette",
        nargs=4,
        type=parse_hex_byte,
        metavar=("C0", "C1", "C2", "C3"),
        help=(
            "Render using one NES four-color palette. "
            "Example: --palette 0F 20 0F 06"
        )
    )

    parser.add_argument(
        "--palette-ram",
        nargs=32,
        type=parse_hex_byte,
        metavar="XX",
        help=(
            "Complete 32-byte NES palette RAM image. "
            "Generates output for all four BG and "
            "all four sprite palettes."
        )
    )

    args = parser.parse_args()

    if args.scale < 1:
        parser.error(
            "--scale must be at least 1"
        )

    if args.columns < 1:
        parser.error(
            "--columns must be at least 1"
        )

    if (
        args.palette is not None and
        args.palette_ram is not None
    ):
        parser.error(
            "--palette and --palette-ram "
            "cannot be used together"
        )

    input_path = Path(
        args.input
    )

    data = input_path.read_bytes()

    if args.output_prefix:
        output_prefix = (
            args.output_prefix
        )
    else:
        output_prefix = (
            input_path.stem
        )

    # -------------------------------------------------------------------------
    # Extract CHR
    # -------------------------------------------------------------------------

    if args.raw:

        banks = extract_raw_chr_banks(
            data
        )

    elif data[:4] == b"NES\x1A":

        banks = extract_chr_banks_from_ines(
            data
        )

    else:

        print(
            "No iNES header detected. "
            "Treating input as raw CHR."
        )

        banks = extract_raw_chr_banks(
            data
        )

    print(
        f"Input:      {input_path}"
    )

    print(
        f"CHR banks:  {len(banks)}"
    )

    # -------------------------------------------------------------------------
    # Full palette RAM mode
    # -------------------------------------------------------------------------

    if args.palette_ram is not None:

        process_palette_ram(
            banks,
            args.palette_ram,
            output_prefix,
            args.scale,
            args.columns
        )

        return

    # -------------------------------------------------------------------------
    # Single-palette / grayscale mode
    # -------------------------------------------------------------------------

    if args.palette is not None:

        colors = nes_palette_to_rgb(
            args.palette
        )

        suffix = "_color"

        print(
            "Palette:    " +
            " ".join(
                f"${value:02X}"
                for value
                in args.palette
            )
        )

    else:

        colors = GRAYSCALE
        suffix = ""

        print(
            "Palette:    grayscale CHR values"
        )

    for bank_index, bank in enumerate(banks):

        print()
        print(
            f"CHR bank ${bank_index:02X}"
        )

        print(
            f"  Size:     ${len(bank):04X}"
        )

        print(
            f"  Tiles:    "
            f"{len(bank) // TILE_SIZE_BYTES}"
        )

        process_chr_bank(
            bank,
            bank_index,
            output_prefix,
            args.scale,
            args.columns,
            colors,
            suffix
        )


if __name__ == "__main__":
    main()