#!/usr/bin/env python3
"""Create a byte-identical build skeleton for an NES project folder.

This is Pass 0: the ROM image is split into its header / PRG / CHR / footer
binaries and included verbatim, so the project rebuilds byte-for-byte
identical to the original ROM from its very first commit.  Pass 1 later
replaces the `prg.s` include with the traced disassembly; every ROM byte
stays accounted for either way, which is what the repo's contract is.

The point of doing this before Pass 1 is that the *folder shape* and the
*build* are fixed and proven first, so the disassembly work afterwards is
only ever about content, never about plumbing.

Layout of the generated project folder (flat, matching the promoted tier):

    header.bin prg.bin chr.bin trailer.bin   binary splits of the ROM
    rom.nes                                  space-free copy for `cmp`
    header.s chr.s trailer.s prg.s           ca65 sources
    linker.cfg symbols.inc locals.inc Makefile
    README.md                                ROM facts + open work

Two layouts drive the linker config, selected with `--layout`:

  * `flat` (the default) puts the MEMORY areas at *file offsets* in the
    .nes file - HEADER at $0000, PRG at $0010, CHR after it.  This works
    for any ROM size, and it is the honest model for a banked mapper, where
    the CPU never sees the whole PRG image at once.  `NROM/Bomberman (U)` is
    the repo's precedent for it.

  * `cpu` puts the PRG area at the CPU address the image is mapped at,
    which is only possible when the whole image fits in the 6502's 64 KB
    address space: a 16 KB PRG at $C000 (NROM-128) or a 32 KB PRG at $8000
    (NROM-256, contiguous $8000-$FFFF).  The 39 `_wip/` projects that
    predate the flat layout use this one, and their PRG really is mapped
    at that address.  Re-running this tool on one of them with the default
    would switch it to the flat layout - still byte-identical, but it would
    lose the CPU-address mapping, so use `--layout cpu` for those.

  * Either way, ld65 writes the output file in MEMORY *declaration* order,
    not address order, so the areas are declared in the order the bytes
    appear in the file: HEADER, PRG, CHR, TRAILER.

Usage:
    python3 _tools/make_skeleton.py <project dir> [<project dir> ...] [--force]
"""

import argparse
import hashlib
import shutil
import sys
from pathlib import Path

REPO = Path(__file__).resolve().parent.parent

# The hardware-constants file copied into every generated project.  The
# promoted tier uses this exact file in 25 of its 35 projects.
SYMBOLS_SRC = REPO / "NROM/Bomberman (U)/symbols.inc"

# Mirroring is header byte 6 bit 0 only.  Bit 1 is the battery flag and bit 3
# is four-screen VRAM, so neither may be folded into this value: the README
# carries the arrangement here and the four-screen fact in its own row.
MIRRORING = {0: "Horizontal", 1: "Vertical"}

# Mapper number -> board name, for the mappers the repository has a folder
# for.  Anything else is named after its number (`MAPPER<NN>`), which is
# what `_wip/MAPPER13/` does for Videomation (U).
BOARDS = {
    0: "NROM",
    1: "MMC1",
    2: "UNROM",
    3: "CNROM",
    4: "MMC3",
    5: "MMC5",
}


def board_name(mapper):
    """A human-readable board name for a mapper number."""
    return BOARDS.get(mapper, f"MAPPER{mapper}")


class SkeletonError(Exception):
    pass


def parse_ines(data):
    """Split an iNES image into header / PRG / CHR / footer and decode it."""
    if len(data) < 16 or data[:4] != b"NES\x1a":
        raise SkeletonError("not an iNES image (bad magic)")

    prg_banks, chr_banks = data[4], data[5]
    f6, f7 = data[6], data[7]

    # iNES byte 6: bit 0 mirroring, bit 1 battery, bit 2 trainer,
    # bit 3 four-screen, bits 4-7 mapper bits 0-3.
    # iNES byte 7: bits 4-7 mapper bits 4-7.
    mapper = (f6 >> 4) | (f7 & 0xF0)

    off = 16
    trainer = b""
    if f6 & 0x04:
        trainer = data[off:off + 512]
        off += 512

    prg = data[off:off + prg_banks * 16384]
    off += prg_banks * 16384
    chr_ = data[off:off + chr_banks * 8192]
    off += chr_banks * 8192
    footer = data[off:]

    if len(prg) != prg_banks * 16384:
        raise SkeletonError("ROM is truncated: PRG shorter than the header declares")

    return {
        "header": data[:16],
        "trainer": trainer,
        "prg": prg,
        "chr": chr_,
        "footer": footer,
        "prg_banks": prg_banks,
        "chr_banks": chr_banks,
        "mapper": mapper,
        "mirroring": f6 & 0x01,
        "four_screen": bool(f6 & 0x08),
        "battery": bool(f6 & 0x02),
        "trainer_flag": bool(f6 & 0x04),
        "f6": f6,
        "f7": f7,
    }


def prg_area(prg_len, layout, trainer_len=0):
    """Return (start, size) for the PRG memory area.

    Two layouts are possible.

    `cpu`  - the PRG area is placed at the CPU address the image is mapped
             at, which is only possible when the whole image fits in the
             6502's 64 KB address space:
                 16 KB -> $C000 (NROM-128, the repo's convention)
                 32 KB -> $8000 (NROM-256, contiguous $8000-$FFFF)
             This is what the existing `_wip/` projects use, because their
             PRG really is mapped at that address.

    `flat` - the areas are placed at *file offsets*: HEADER at $0000, the
             trainer (if any) at $0010, PRG after that, CHR immediately
             after the PRG. The addresses are not CPU addresses; they are
             positions in the .nes file. This is the only option once the
             PRG exceeds 32 KB, because a 64 KB+ image does not fit the
             address space, and it is also the honest model for a banked
             mapper (MMC1, MMC3, ...), where the CPU never sees the whole
             image at once. `NROM/Bomberman (U)` is the repo's precedent for
             the flat layout.
    """
    if layout == "cpu":
        if prg_len == 16384:
            return 0xC000, 0x4000
        if prg_len == 32768:
            return 0x8000, 0x8000
        raise SkeletonError(
            f"PRG size {prg_len} does not fit a contiguous CPU window; "
            f"use the flat layout")
    # flat: file offsets.  The trainer sits between the header and the PRG.
    return 0x0010 + trainer_len, prg_len


def hexs(value):
    return f"${value:04X}"


def vectors(prg):
    """NMI / RESET / IRQ, read from the last six bytes of the PRG image.

    These live at $FFFA-$FFFF of the *fixed* bank, which for a banked image
    is the last bank in the file, so the end of the PRG image is the right
    place to read them from.
    """
    if len(prg) < 6:
        raise SkeletonError("PRG image too small for vectors")
    v = prg[-6:]
    return (int.from_bytes(v[0:2], "little"),
            int.from_bytes(v[2:4], "little"),
            int.from_bytes(v[4:6], "little"))


def linker_cfg(prg_start, prg_size, chr_len, has_chr, footer_len,
              trainer_len=0, chr_start=0x0000):
    """MEMORY areas, declared in the order the bytes appear in the .nes file.

    ld65 writes the output file in MEMORY declaration order, not address
    order, so this order is what reproduces the ROM exactly.
    """
    lines = [
        "MEMORY {",
        "  HEADER: start = $0000, size = $0010, fill = yes;",
    ]
    segs = [
        "  HEADER: load = HEADER, type = ro;",
    ]
    if trainer_len:
        lines.append(f"  TRAINER: start = $0010, size = {hexs(trainer_len)}, fill = yes;")
        segs.append("  TRAINER: load = TRAINER, type = ro;")
    lines.append(f"  PRG:    start = {hexs(prg_start)}, size = {hexs(prg_size)}, fill = yes;")
    segs.append("  PRG:    load = PRG,    type = ro;")
    if has_chr:
        lines.append(f"  CHR:    start = {hexs(chr_start)}, size = {hexs(chr_len)}, fill = yes;")
        segs.append("  CHR:    load = CHR,    type = ro;")
    if footer_len:
        lines.append(f"  TRAILER: start = $0000, size = {hexs(footer_len)}, fill = yes;")
        segs.append("  TRAILER: load = TRAILER, type = ro;")
    lines.append("}")
    lines.append("SEGMENTS {")
    lines.extend(segs)
    lines.append("}")
    return "\n".join(lines) + "\n"


def makefile(has_chr, has_footer, has_trainer=False):
    """The whole-ROM build: header + prg + chr (+ footer) linked, then cmp."""
    lines = [
        "CA65 ?= ca65",
        "LD65 ?= ld65",
        "",
        "all: reconstructed.nes verify",
        "",
        "prg.o: prg.s symbols.inc locals.inc prg.bin",
        "\t$(CA65) -g -o $@ $<",
        "",
        "header.o: header.s header.bin",
        "\t$(CA65) -g -o $@ $<",
        "",
    ]
    objs = ["header.o", "prg.o"]
    if has_trainer:
        lines += ["trainer.o: trainer.s trainer.bin", "\t$(CA65) -g -o $@ $<", ""]
        objs.insert(1, "trainer.o")
    if has_chr:
        lines += ["chr.o: chr.s chr.bin", "\t$(CA65) -g -o $@ $<", ""]
        objs.append("chr.o")
    if has_footer:
        lines += ["trailer.o: trailer.s trailer.bin", "\t$(CA65) -g -o $@ $<", ""]
        objs.append("trailer.o")
    lines += [
        "reconstructed.nes: " + " ".join(objs) + " linker.cfg",
        "\t$(LD65) -C linker.cfg -m reconstructed.map -Ln reconstructed.lbl "
        "-o $@ " + " ".join(objs),
        "",
        "verify: reconstructed.nes rom.nes",
        "\tcmp reconstructed.nes rom.nes",
        '\t@echo "Byte-identical ROM verified."',
        "",
        "clean:",
        "\trm -f *.o reconstructed.nes reconstructed_prg.bin reconstructed.map "
        "reconstructed.lbl",
        "",
        ".PHONY: all verify clean",
    ]
    return "\n".join(lines) + "\n"


def prg_source(project, prg_start, prg_size, board, vec, layout):
    nmi, rst, irq = vec
    if layout == "flat":
        placement = (f"{prg_size // 1024} KB at file offset {hexs(prg_start)} "
                     f"(the image is bank-switched; it does not fit the 6502 "
                     f"address space as one window)")
    else:
        placement = (f"{prg_size // 1024} KB mapped at "
                     f"{hexs(prg_start)}-{hexs(prg_start + prg_size - 1)}")
    return "\n".join([
        f"; {project} NES disassembly - Pass 0 (build skeleton)",
        ";",
        "; The PRG image is included verbatim, so the project rebuilds",
        "; byte-identical to the original ROM from the first commit.  Pass 1",
        "; replaces this block with the traced disassembly; no ROM byte is",
        "; ever deleted or reformatted, only re-classified as code or data.",
        ";",
        f"; Board:    {board}",
        f"; PRG:      {placement}",
        f"; Vectors:  NMI=${nmi:04X}  RESET=${rst:04X}  IRQ=${irq:04X}",
        "",
        '.setcpu "6502"',
        "",
        '.include "symbols.inc"',
        '.include "locals.inc"',
        "",
        '.segment "PRG"',
        '    .incbin "prg.bin"',
        "",
    ])


def readme(project, rom_name, info, prg_start, prg_size, board, vec, sha,
           layout):
    nmi, rst, irq = vec
    if layout == "flat":
        prg_desc = (f"{len(info['prg']) // 1024} KB "
                    f"({info['prg_banks']} x 16 KB), at file offset "
                    f"{hexs(prg_start)}")
    else:
        prg_desc = (f"{len(info['prg']) // 1024} KB "
                    f"({info['prg_banks']} x 16 KB), mapped at CPU "
                    f"{hexs(prg_start)}-{hexs(prg_start + prg_size - 1)}")
    facts = [
        ("Format", "iNES"),
        ("Mapper", f"{info['mapper']}"),
        ("Board", board),
        ("Mirroring", MIRRORING[info["mirroring"]]),
        ("Four-screen VRAM", "yes" if info["four_screen"] else "no"),
        ("Battery", "yes" if info["battery"] else "no"),
        ("PRG", prg_desc),
        ("CHR", f"{len(info['chr']) // 1024} KB ({info['chr_banks']} x 8 KB)"
                if info["chr"] else "none"),
        ("Trainer", "512 bytes" if info["trainer_flag"] else "none"),
        ("Footer", f"{len(info['footer'])} bytes past the CHR"
                   if info["footer"] else "none"),
        ("NMI vector", f"${nmi:04X}"),
        ("RESET vector", f"${rst:04X}"),
        ("IRQ vector", f"${irq:04X}"),
        ("ROM file", f"`{rom_name}`"),
        ("SHA-256", f"`{sha}`"),
    ]
    fact_rows = "\n".join(f"| {k} | {v} |" for k, v in facts)

    flags = []
    if info["f7"] & 0x0F:
        flags.append(f"header byte 7 is `${info['f7']:02X}`; its upper nybble "
                     f"selects mapper bits 4-7")
    if info["f6"] & 0x08:
        flags.append("header byte 6 bit 3 (four-screen VRAM) is set")
    if info["trainer_flag"]:
        flags.append("the header declares a 512-byte trainer, which follows "
                     "the header in the file")
    if info["footer"]:
        flags.append(f"{len(info['footer'])} bytes follow the CHR; they are "
                     "reproduced by the TRAILER memory area")
    if rst == 0 and nmi == 0 and irq == 0:
        flags.append("all three vectors are $0000 - the image carries no "
                     "vector table")
    flagtext = "\n".join(f"- {f}" for f in flags) if flags else "- none"

    sources = ["`header.s`", "`prg.s`"]
    if info["trainer_flag"]:
        sources.append("`trainer.s`")
    if info["chr"]:
        sources.append("`chr.s`")
    if info["footer"]:
        sources.append("`trailer.s`")

    files = [
        ("`prg.s`", "Pass 0 PRG include (replaced by the Pass 1 disassembly)"),
        ("`header.s`", "iNES header include"),
    ]
    if info["trainer_flag"]:
        files.append(("`trainer.s`", "trainer include"))
    if info["chr"]:
        files.append(("`chr.s`", "CHR include"))
    if info["footer"]:
        files.append(("`trailer.s`", "footer include"))
    files += [
        ("`prg.bin`", f"{len(info['prg'])} bytes "
                      f"({len(info['prg']) // 1024} KB) PRG image"),
    ]
    if info["chr"]:
        files.append(("`chr.bin`", f"{len(info['chr'])} bytes "
                                    f"({len(info['chr']) // 1024} KB) CHR image"))
    if info["footer"]:
        files.append(("`trailer.bin`", f"{len(info['footer'])} footer bytes"))
    files += [
        ("`rom.nes`", "byte-identical copy of the original ROM"),
        ("`linker.cfg`", "HEADER / PRG"
         + (" / CHR" if info["chr"] else "")
         + (" / TRAILER" if info["footer"] else "") + " memory areas"),
        ("`symbols.inc`", "NES hardware register names"),
        ("`locals.inc`", "zero-page aliases (empty at Pass 0)"),
        ("`Makefile`", "build and byte-identity check"),
    ]
    file_rows = "\n".join(f"| {k} | {v} |" for k, v in files)

    return f"""# {project}

## Status

**Pass 0 skeleton.** The project builds and verifies byte-identical to the
original ROM, but no disassembly exists yet: `prg.s` is a single
`.incbin "prg.bin"` of the PRG image.

`make` assembles {", ".join(sources)}, links them with `ld65` and compares the
result against `rom.nes`:

```
Byte-identical ROM verified.
```

## ROM

| Property | Value |
|----------|-------|
{fact_rows}

### Header notes

{flagtext}

## Open work

- **Pass 1** - run `nesdis.py` over `prg.bin` to produce the code/data
  split. The tracer reads the vectors from the last six bytes of the PRG
  image, so the split seeds correctly.
- **Pass 1** must emit the canonical audit comment
  `; CPU:$C123 PRG:0123  AD 21 03` at column 35 for every instruction.
- **Pass 2** - replace the address-derived `Label_XXXX` names and the raw
  `$XXXX` call operands.
- **Pass 3** - semantic zero-page aliases (applied only at the sites where
  their role is proven) and named data tables.
- **Promotion** to `NROM/<name>/` once Pass 3 is done, as `git mv`.

## Files

| File | Purpose |
|------|---------|
{file_rows}
"""


def find_rom(folder):
    """Pick the original ROM in the folder.

    `rom.nes` is the space-free copy the build compares against and
    `reconstructed.nes` is the build output, so neither is the source
    image.
    """
    build_outputs = {"rom.nes", "reconstructed.nes"}
    cands = sorted(p for p in folder.glob("*.nes") if p.name not in build_outputs)
    if not cands:
        raise SkeletonError(f"{folder}: no .nes image found")
    if len(cands) > 1:
        raise SkeletonError(
            f"{folder}: {len(cands)} candidate .nes files "
            f"({', '.join(p.name for p in cands)}) - remove the extras")
    return cands[0]


def build(folder, force, layout="flat", board=None):
    folder = Path(folder)
    if not folder.is_dir():
        raise SkeletonError(f"{folder}: not a directory")

    rom = find_rom(folder)
    data = rom.read_bytes()
    info = parse_ines(data)

    prg_start, prg_size = prg_area(len(info["prg"]), layout,
                                   len(info["trainer"]))
    chr_len = len(info["chr"])
    # In the flat layout the CHR sits immediately after the PRG in the file.
    chr_start = prg_start + prg_size
    vec = vectors(info["prg"])
    has_chr = bool(info["chr"])
    has_footer = bool(info["footer"])
    sha = hashlib.sha256(data).hexdigest()
    if board is None:
        board = board_name(info["mapper"])

    # Refuse to stomp a real disassembly.
    guarded = ["prg.s", "Makefile", "linker.cfg", "symbols.inc", "locals.inc"]
    existing = [g for g in guarded if (folder / g).exists()]
    if existing and not force:
        raise SkeletonError(
            f"{folder}: already has {', '.join(existing)} - pass --force to "
            f"overwrite (this replaces a real disassembly with a Pass 0 "
            f"include)")

    # Binary splits.
    (folder / "header.bin").write_bytes(info["header"])
    (folder / "prg.bin").write_bytes(info["prg"])
    if info["trainer_flag"]:
        (folder / "trainer.bin").write_bytes(info["trainer"])
    if has_chr:
        (folder / "chr.bin").write_bytes(info["chr"])
    if has_footer:
        (folder / "trailer.bin").write_bytes(info["footer"])
    shutil.copyfile(rom, folder / "rom.nes")

    # ca65 sources.
    (folder / "header.s").write_text(
        '.segment "HEADER"\n    .incbin "header.bin"\n')
    if info["trainer_flag"]:
        (folder / "trainer.s").write_text(
            '.segment "TRAINER"\n    .incbin "trainer.bin"\n')
    if has_chr:
        (folder / "chr.s").write_text(
            '.segment "CHR"\n    .incbin "chr.bin"\n')
    if has_footer:
        (folder / "trailer.s").write_text(
            '.segment "TRAILER"\n    .incbin "trailer.bin"\n')
    (folder / "prg.s").write_text(
        prg_source(folder.name, prg_start, prg_size, board, vec, layout))

    # Config.
    shutil.copyfile(SYMBOLS_SRC, folder / "symbols.inc")
    (folder / "locals.inc").write_text(
        "; Zero-page aliases for this project.\n"
        "; Empty at Pass 0.  Pass 1+ adds proven semantic names here;\n"
        "; addresses whose role is not yet determined keep the neutral ZpXX\n"
        "; form.\n")
    (folder / "linker.cfg").write_text(
        linker_cfg(prg_start, prg_size, chr_len, has_chr,
                   len(info["footer"]), len(info["trainer"]), chr_start))
    (folder / "Makefile").write_text(
        makefile(has_chr, has_footer, bool(info["trainer"])))
    (folder / "README.md").write_text(
        readme(folder.name, rom.name, info, prg_start, prg_size, board, vec,
               sha, layout))

    return {
        "project": folder.name,
        "board": board,
        "mapper": info["mapper"],
        "prg": len(info["prg"]),
        "chr": len(info["chr"]),
        "footer": len(info["footer"]),
        "sha": sha,
    }


def main():
    ap = argparse.ArgumentParser(
        description="Create a byte-identical build skeleton for an NES project.")
    ap.add_argument("folders", nargs="+", help="project folders")
    ap.add_argument("--force", action="store_true",
                    help="overwrite existing prg.s / Makefile / linker.cfg")
    ap.add_argument("--layout", choices=("flat", "cpu"), default="flat",
                    help="flat: MEMORY areas at .nes file offsets (works for "
                         "any ROM size, and is the honest model for a banked "
                         "mapper); cpu: the PRG area at the CPU address it is "
                         "mapped at (only possible for a 16 KB or 32 KB PRG)")
    args = ap.parse_args()

    results = []
    failed = False
    for f in args.folders:
        try:
            results.append(build(f, args.force, args.layout))
        except SkeletonError as exc:
            print(f"ERROR: {exc}", file=sys.stderr)
            failed = True

    for r in results:
        print(f"{r['project']}: {r['board']} mapper={r['mapper']} "
              f"prg={r['prg'] // 1024}K chr={r['chr'] // 1024}K "
              f"footer={r['footer']} sha256={r['sha'][:16]}")

    return 1 if failed else 0


if __name__ == "__main__":
    sys.exit(main())