#!/usr/bin/env python3
"""Split an iNES or NES 2.0 ROM into header, trainer, PRG, and CHR files."""

import argparse
import sys
from pathlib import Path

HEADER_SIZE = 16
TRAINER_SIZE = 512
PRG_BANK_SIZE = 16 * 1024
CHR_BANK_SIZE = 8 * 1024


def decode_nes2_size(lsb: int, ms_nibble: int, unit: int, kind: str) -> int:
    """Return the NES 2.0 ROM size in bytes."""
    if ms_nibble != 0x0F:
        return ((ms_nibble << 8) | lsb) * unit

    # NES 2.0 exponent/multiplier encoding:
    # size = 2^E * (M * 2 + 1), where E = lsb >> 2 and M = lsb & 3.
    exponent = lsb >> 2
    multiplier = (lsb & 0x03) * 2 + 1
    size = (1 << exponent) * multiplier
    if size <= 0:
        raise ValueError(f"Invalid NES 2.0 {kind} size")
    return size


def parse_header(data: bytes) -> dict:
    if len(data) < HEADER_SIZE:
        raise ValueError("File is too small to contain a 16-byte NES header")

    if data[:4] != b"NES\x1a":
        raise ValueError("Not an iNES/NES 2.0 ROM (missing NES 1A signature)")

    flags6 = data[6]
    flags7 = data[7]
    nes2 = (flags7 & 0x0C) == 0x08
    has_trainer = bool(flags6 & 0x04)

    mapper = (flags6 >> 4) | (flags7 & 0xF0)
    submapper = None

    if nes2:
        mapper |= (data[8] & 0x0F) << 8
        submapper = data[8] >> 4
        prg_size = decode_nes2_size(data[4], data[9] & 0x0F, PRG_BANK_SIZE, "PRG-ROM")
        chr_size = decode_nes2_size(data[5], (data[9] >> 4) & 0x0F, CHR_BANK_SIZE, "CHR-ROM")
    else:
        prg_size = data[4] * PRG_BANK_SIZE
        chr_size = data[5] * CHR_BANK_SIZE

    return {
        "format": "NES 2.0" if nes2 else "iNES",
        "mapper": mapper,
        "submapper": submapper,
        "has_trainer": has_trainer,
        "prg_size": prg_size,
        "chr_size": chr_size,
    }


def write_file(path: Path, data: bytes) -> None:
    path.write_bytes(data)
    print(f"  {path.name:<16} {len(data):>8} bytes")


def split_fixed_banks(
    data: bytes,
    offset: int,
    total_size: int,
    bank_size: int,
    prefix: str,
    output_dir: Path,
) -> int:
    """Split a ROM region into fixed-size banks and return the new offset."""
    end = offset + total_size
    if end > len(data):
        raise ValueError(
            f"ROM is truncated while reading {prefix.upper()}: "
            f"header requires {total_size} bytes at file offset 0x{offset:X}, "
            f"but only {len(data) - offset} bytes remain"
        )

    region = data[offset:end]

    # Normal iNES sizes divide exactly into these units. NES 2.0's
    # exponent/multiplier encoding may not, so preserve any final partial bank.
    for bank_no, start in enumerate(range(0, total_size, bank_size)):
        chunk = region[start:start + bank_size]
        path = output_dir / f"{prefix}_{bank_no:02X}.bin"
        write_file(path, chunk)

    return end


def main() -> int:
    parser = argparse.ArgumentParser(
        description="Split an NES ROM into PRG-ROM and CHR-ROM banks."
    )
    parser.add_argument("rom", type=Path, help="input .nes ROM")
    parser.add_argument(
        "-o", "--output", type=Path,
        help="output directory (default: <romname>_split)",
    )
    args = parser.parse_args()

    if not args.rom.is_file():
        print(f"Error: ROM not found: {args.rom}", file=sys.stderr)
        return 1

    try:
        data = args.rom.read_bytes()
        info = parse_header(data)

        output_dir = args.output
        if output_dir is None:
            output_dir = args.rom.parent / f"{args.rom.stem}_split"
        output_dir.mkdir(parents=True, exist_ok=True)

        print(f"ROM:       {args.rom}")
        print(f"Size:      {len(data)} bytes")
        print(f"Format:    {info['format']}")
        print(f"Mapper:    {info['mapper']}")
        if info["submapper"] is not None:
            print(f"Submapper: {info['submapper']}")
        print(f"Trainer:   {'yes' if info['has_trainer'] else 'no'}")
        print(f"PRG-ROM:   {info['prg_size']} bytes ({info['prg_size'] / 1024:g} KiB)")
        print(f"CHR-ROM:   {info['chr_size']} bytes ({info['chr_size'] / 1024:g} KiB)")
        print(f"Output:    {output_dir}")
        print()

        print("Header:")
        write_file(output_dir / "header.bin", data[:HEADER_SIZE])

        offset = HEADER_SIZE

        if info["has_trainer"]:
            trainer_end = offset + TRAINER_SIZE
            if trainer_end > len(data):
                raise ValueError("ROM is truncated inside the 512-byte trainer")
            print("\nTrainer:")
            write_file(output_dir / "trainer.bin", data[offset:trainer_end])
            offset = trainer_end

        print("\nPRG-ROM:")
        if info["prg_size"]:
            offset = split_fixed_banks(
                data, offset, info["prg_size"], PRG_BANK_SIZE, "prg", output_dir
            )
        else:
            print("  none")

        print("\nCHR-ROM:")
        if info["chr_size"]:
            offset = split_fixed_banks(
                data, offset, info["chr_size"], CHR_BANK_SIZE, "chr", output_dir
            )
        else:
            print("  none (cartridge may use CHR-RAM)")

        if offset < len(data):
            print("\nTrailing data:")
            write_file(output_dir / "extra.bin", data[offset:])

        print("\nDone.")
        return 0

    except (OSError, ValueError) as exc:
        print(f"Error: {exc}", file=sys.stderr)
        return 1


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