#!/usr/bin/env python3
"""Give a data block a name, splitting it if the name belongs in the middle.

`nesdis.py` emits unreached bytes as `.byte` rows with a `; ---- data
$XXXX-$YYYY` header.  Pass 3 wants every table and data block to carry a
name, and a table often starts in the middle of such a row.  This tool
inserts the label at the right place: at the start of a row when the
address matches, or by splitting the row in two (each half keeps its own
audit comment and no byte changes) when it does not.

The bytes are untouched, so the build cannot change - `make` still has to
print BYTE-IDENTICAL.

Usage:
    python3 _tools/name_data_region.py <project dir> <addr>=<name> [...]
"""

import argparse
import re
import sys
from pathlib import Path

CPU_RE = re.compile(r"; CPU:\$([0-9A-Fa-f]{4}) PRG:[0-9A-Fa-f]{4}\s+"
                    r"((?:[0-9A-Fa-f]{2} ?)+)\s*$")
BYTE_RE = re.compile(r"^(\s*)\.byte\s+(.*?)\s{2,}; .*?CPU:\$([0-9A-Fa-f]{4}) PRG:")


def parse_bytes(text):
    vals = []
    for tok in text.split(","):
        tok = tok.strip()
        if not tok:
            continue
        vals.append(int(tok.lstrip("$"), 16))
    return vals


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("dir", help="project directory")
    ap.add_argument("names", nargs="+", metavar="ADDR=NAME")
    ap.add_argument("--write", action="store_true", help="apply (default: dry run)")
    ap.add_argument("--base", default="0xC000",
                    help="CPU address the PRG image is mapped at, for the PRG "
                         "offset in the audit comment (default 0xC000; use "
                         "0x8000 for a 32 KB NROM image mapped at $8000)")
    ap.add_argument("--segment", default=None,
                    help="only name data inside this ca65 segment.  A banked "
                         "image maps several banks into the same CPU window, so "
                         "without this the same address is named in every bank "
                         "that covers it")
    ap.add_argument("--file", default="prg.s",
                    help="the source file to name blocks in (default prg.s).  A "
                         "banked project keeps one file per bank, and the tool "
                         "works on one file at a time")
    args = ap.parse_args()

    prg = Path(args.dir) / args.file
    if not prg.exists():
        raise SystemExit(f"error: {prg} does not exist")
    lines = prg.read_text().split("\n")
    base = int(str(args.base), 0)

    wanted = {}
    for spec in args.names:
        a, _, name = spec.partition("=")
        wanted[int(a.lstrip("$"), 16)] = name
    # Sort so the labels land in address order.
    wanted = dict(sorted(wanted.items()))

    out = []
    done = set()
    segment = None
    for line in lines:
        if line.startswith(".segment"):
            segment = line.split('"')[1] if '"' in line else None
        m = BYTE_RE.match(line)
        if not m:
            out.append(line)
            continue
        if args.segment is not None and segment != args.segment:
            out.append(line)
            continue
        indent, body = m.group(1), m.group(2)
        cpu = int(m.group(3), 16)
        # Some sources carry a redundant provenance note in front of the
        # canonical audit comment ("; PRG $0647 | CPU $C647  ; CPU:$C647 ...").
        # The byte list ends at the first ';' either way.
        body = body.split(";")[0]
        vals = parse_bytes(body)

        hits = [a for a in wanted if cpu <= a <= cpu + len(vals) - 1]
        if not hits:
            out.append(line)
            continue

        def emit(vs, start):
            if not vs:
                return
            b = ", ".join(f"${v:02X}" for v in vs)
            r = " ".join(f"{v:02X}" for v in vs)
            out.append(f"{indent}.byte {b}  ; CPU:${start:04X} PRG:{start - base:04X}  {r}")

        pos = 0
        for a in hits:
            k = a - cpu
            emit(vals[pos:k], cpu + pos)
            out.append(f"{wanted[a]}:")
            done.add(a)
            pos = k
        emit(vals[pos:], cpu + pos)

    missing = set(wanted) - done
    if missing:
        raise SystemExit("error: no .byte row covers " +
                         ", ".join(f"${a:04X}" for a in sorted(missing)))

    print(f"{len(done)} data block(s) named")
    if args.write:
        prg.write_text("\n".join(out))
        print("written")
    return 0


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