#!/usr/bin/env python3

import sys
from pathlib import Path

def split_nes_rom(filename):
    data = Path(filename).read_bytes()

    if len(data) < 16:
        raise ValueError("File is too small to be a valid NES ROM.")

    header = data[:16]

    if header[0:4] != b"NES\x1A":
        raise ValueError("Invalid iNES header.")

    prg_banks = header[4]
    chr_banks = header[5]

    prg_size = prg_banks * 16384      # 16 KB per bank
    chr_size = chr_banks * 8192       # 8 KB per bank

    print(f"PRG Banks : {prg_banks}")
    print(f"CHR Banks : {chr_banks}")
    print(f"PRG Size  : {prg_size} bytes")
    print(f"CHR Size  : {chr_size} bytes")

    prg_start = 16
    prg_end = prg_start + prg_size

    chr_start = prg_end
    chr_end = chr_start + chr_size

    prg = data[prg_start:prg_end]
    chr = data[chr_start:chr_end]

    Path("header.bin").write_bytes(header)
    Path("prg.bin").write_bytes(prg)
    Path("chr.bin").write_bytes(chr)

    print("Created:")
    print("  header.bin")
    print("  prg.bin")
    print("  chr.bin")

if __name__ == "__main__":
    if len(sys.argv) != 2:
        print(f"Usage: {sys.argv[0]} rom.nes")
        sys.exit(1)

    split_nes_rom(sys.argv[1])