#!/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:[0-9A-Fa-f]{4}\s+"
                     r"((?:[0-9A-Fa-f]{2} ?)+)\s*)$")


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)")
    args = ap.parse_args()

    prg = Path(args.dir) / "prg.s"
    lines = prg.read_text().split("\n")

    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()
    for line in lines:
        m = BYTE_RE.match(line)
        if not m:
            out.append(line)
            continue
        indent, body, comment = m.group(1), m.group(2), m.group(3)
        cpu = int(re.search(r"CPU:\$([0-9A-Fa-f]{4})", comment).group(1), 16)
        vals = parse_bytes(body)
        raw = re.search(r"((?:[0-9A-Fa-f]{2} ?)+)\s*$", comment).group(1).strip()

        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 - 0xC000: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())