#!/usr/bin/env python3
"""Re-classify a region of prg.s from data to code, or the other way round.

The Pass 1 tracer classifies a byte as data when nothing reaches it.  When
a dispatch table that the traced code *does* read points into such a
region, the region is code after all - the tracer simply cannot follow an
indirect jump.  This tool re-decodes such a region straight from the ROM
and rewrites the source, so the split matches the ROM.

The whole span of source lines covered by the region is replaced: the
data block header and its `.byte` lines, and any instruction lines that
the tracer mis-split.  The audit trail is preserved - every instruction
carries the canonical `; CPU:$XXXX PRG:XXXX  AA BB` comment with its
exact bytes, in the emitter's own comment column (35).  The bytes are the
ROM's, so the build cannot change - `make` still has to print
BYTE-IDENTICAL.

Operand text comes from the tracer's renderer, so zero-page operands come
out as `ZpXX`; hardware registers and absolute targets come out numeric.
Run `apply_symbols.py` afterwards to name the hardware registers, and a
pass-2 map to name the targets.

Usage:
    python3 _tools/recode_region.py <project dir> START-END [START-END ...]
                                  [--labels FILE] [--write]

`--labels` is a TSV of `<CPU address><TAB><name>`; a name marks a label to
emit at that address.
"""

import argparse
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import nesdis  # noqa: E402

CPU_RE = re.compile(r"; CPU:\$([0-9A-F]{4}) PRG:")


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("dir", help="project directory")
    ap.add_argument("regions", nargs="+", metavar="START-END",
                    help="CPU address range to re-decode as code")
    ap.add_argument("--labels", help="TSV of <address><TAB><name>")
    ap.add_argument("--write", action="store_true", help="apply (default: dry run)")
    args = ap.parse_args()

    prg_dir = Path(args.dir)
    blob = (prg_dir / "prg.bin").read_bytes()
    src = prg_dir / "prg.s"
    if not src.is_file():
        raise SystemExit(f"error: {prg_dir} does not look like a project")
    base = 0xC000

    labels = {}
    if args.labels:
        for raw in Path(args.labels).read_text().splitlines():
            line = raw.split("#", 1)[0].strip()
            if not line:
                continue
            parts = line.split("\t")
            if len(parts) < 2:
                parts = line.split(None, 1)
            labels[int(parts[0].strip().lstrip("$"), 16)] = parts[1].strip()

    regions = []
    for spec in args.regions:
        a, _, b = spec.partition("-")
        regions.append((int(a, 16), int(b, 16)))
    regions.sort()

    lines = src.read_text().split("\n")

    # The source line that carries each CPU address.
    line_of = {}
    for i, line in enumerate(lines):
        m = CPU_RE.search(line)
        if m:
            line_of.setdefault(int(m.group(1), 16), i)

    out = []
    i = 0
    changed = 0
    while i < len(lines):
        m = CPU_RE.search(lines[i])
        hit = None
        if m:
            a = int(m.group(1), 16)
            for ra, rb in regions:
                if ra <= a <= rb:
                    hit = (ra, rb)
                    break
        if not hit:
            out.append(lines[i])
            i += 1
            continue

        ra, rb = hit
        # Every source line whose address falls in the region, plus the
        # data-block header that introduces them and any `.byte` lines
        # that follow the last of them.
        inside = [n for n, line in enumerate(lines)
                  if (mm := CPU_RE.search(line))
                  and ra <= int(mm.group(1), 16) <= rb]
        start = min(inside)
        end = max(inside)
        while start > 0:
            prev = lines[start - 1].strip()
            if prev.startswith(".byte") or prev.startswith("; ---- data") \
               or prev.startswith(";   referenced by"):
                start -= 1
                continue
            break
        while end + 1 < len(lines) and lines[end + 1].lstrip().startswith(".byte"):
            end += 1

        pc = ra
        while pc <= rb:
            if pc in labels:
                out.append(f"{labels[pc]}:")
            ins = nesdis.render(blob, base, blob[pc - base], pc - base)
            if ins is None:
                raise SystemExit(f"error: illegal opcode at ${pc:04X}")
            raw = " ".join(f"{b:02X}" for b in blob[pc - base:pc - base + ins.length])
            text = f"{ins.mnemonic} {ins.text}".rstrip()
            out.append(f"    {text:<31}; CPU:${pc:04X} PRG:{pc - base:04X}  {raw}")
            pc += ins.length
        if pc - 1 != rb:
            raise SystemExit(f"error: region ${ra:04X}-${rb:04X} does not end "
                             f"on an instruction boundary (ended ${pc - 1:04X})")
        changed += 1
        i = end + 1

    print(f"re-classified {changed} region(s) as code")
    if args.write:
        src.write_text("\n".join(out))
        print("written")
    return 0


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