#!/usr/bin/env python3
"""Apply a Pass 2 name map (address -> name) to a project's prg.s.

The map is a TSV of `<CPU address><TAB><name>` lines, one per routine or
branch target.  This is the "fresh map derived from the committed split"
workflow the repo uses: the map is written by reading the disassembly, and
this tool only performs the mechanical rewrite.

Two things are rewritten:

  1. `Label_<addr>` label definitions and their references become the
     mapped name.
  2. A raw `$XXXX` operand on a `jsr` or `jmp` becomes the mapped name of
     the address it points at.  Only `jsr`/`jmp` are touched: a raw
     operand on a data reference (`lda $C253`) points at a table, not a
     routine, and naming it would be a lie.

Everything else is preserved byte for byte: the mnemonic, the whitespace and
the trailing `; CPU:$.... PRG:....  XX XX` audit comment.  Because every
replacement names a label that is defined in the same file, the assembled
bytes cannot change - `make` still has to print BYTE-IDENTICAL.

The tool refuses to write if any label in the map has no definition in the
file, or if a rename would collide with an existing label.

Banked cartridges
-----------------
A banked PRG (MMC3 and friends) has several labels at the same CPU address:
one per bank.  Those banks are told apart by a label prefix, and the map may
name them with a `<prefix>:<addr>` key:

    800F    UpdateSprites
    Bank4:800D  Bank4_DrawTile

`--bank-prefix` says which prefix belongs to which part of the file:

    --bank-prefix "0:,8000:Bank4_,A000:Bank5_"

i.e. the labels from file offset $8000 on are the `Bank4_` ones.  A
reference inside a file region is resolved against that region's bank, so
`Label_800F` means bank 0's label in the lower half and bank 4's label in
the upper half.

Usage:
    python3 _tools/apply_pass2_map.py <project dir> <map.tsv> [--write]
"""

import argparse
import re
import sys
from pathlib import Path

LABEL_RE = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*):\s*$")
ADDR_LABEL_RE = re.compile(r"^Label_([0-9A-F]{4})$")
CALL_RE = re.compile(r"^(\s*(?:jsr|jmp)\s+)\$([0-9A-F]{4})(\s*)$")
ADDR_RE = re.compile(r"; CPU:\$([0-9A-F]{4}) PRG:([0-9A-F]{4})(?: BANK:([0-9A-F]+))?")
# a label that belongs to a bank: optional prefix plus the address
# The address-derived label forms the repository uses: `Label_XXXX`,
# `Sub_XXXX` and `SubCXXXX`, each optionally with a bank prefix
# (`Bank4_XXXX`).
BANKED_LABEL_RE = re.compile(
    r"^(?:(?P<pre>[A-Za-z][A-Za-z0-9]*)_)?"
    r"(?:SubC|Sub|Label)_?(?P<addr>[0-9A-F]{4})$")


def read_map(path):
    """Return {(prefix, address): name}.  A bare address means the default
    (unprefixed) bank; `<prefix>:<addr>` names a label in a prefixed bank."""
    out = {}
    for raw in Path(path).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)
        key = parts[0].strip()
        name = parts[1].strip()
        prefix = ""
        if ":" in key:
            prefix, _, key = key.partition(":")
            prefix = prefix.strip()
        addr = int(key.strip().lstrip("$"), 16)
        if not re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", name):
            raise SystemExit(f"error: bad name {name!r} for ${addr:04X}")
        if (prefix, addr) in out and out[(prefix, addr)] != name:
            raise SystemExit(f"error: ${addr:04X} mapped twice "
                             f"({out[(prefix, addr)]} and {name})")
        out[(prefix, addr)] = name
    return out


def read_bank_prefixes(spec):
    """[(file offset, prefix)] sorted by offset."""
    out = []
    for tok in spec.split(","):
        tok = tok.strip()
        if not tok:
            continue
        off, _, pre = tok.partition(":")
        out.append((int(off, 16), pre.strip()))
    return sorted(out)


def prefix_at(bank_prefixes, file_offset):
    pre = ""
    for off, p in bank_prefixes:
        if file_offset >= off:
            pre = p
        else:
            break
    return pre


def file_offset_of(match):
    """File offset of an audit comment.  See derive_branch_names.file_offset_of."""
    if match.lastindex and match.lastindex >= 3 and match.group(3):
        return int(match.group(3), 16) + int(match.group(2), 16)
    return int(match.group(2), 16)


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("dir", help="project directory")
    ap.add_argument("map", help="address -> name TSV")
    ap.add_argument("--bank-prefix", default="",
                    help="comma list of FILEOFF:PREFIX, telling which labels "
                         "belong to which bank")
    ap.add_argument("--write", action="store_true", help="apply (default: dry run)")
    ap.add_argument("--equates", action="store_true",
                    help="name an address that carries no label with an equate in "
                         "locals.inc instead of a new label line - needed for a "
                         "project that uses ca65's '@local' labels, where a "
                         "mid-routine label would re-scope them")
    args = ap.parse_args()

    prg = Path(args.dir) / "prg.s"
    use_equates = args.equates or "@" in prg.read_text()
    if not prg.is_file():
        raise SystemExit(f"error: {prg} not found")

    names = read_map(args.map)
    bank_prefixes = read_bank_prefixes(args.bank_prefix)
    lines = prg.read_text().split("\n")

    # Which addresses actually carry a label definition?
    defined = set()
    for line in lines:
        m = LABEL_RE.match(line)
        if m:
            defined.add(m.group(1))

    # Which addresses carry an instruction at all?  A call target the
    # tracer never labelled still has an instruction there, and the map
    # gives it a name.
    insn_addrs = set()
    for line in lines:
        am = ADDR_RE.search(line)
        if am:
            insn_addrs.add(int(am.group(1), 16))

    # (prefix, address) -> label name, for every label in the file.
    addr_of = {}
    cur = None
    for line in lines:
        m = LABEL_RE.match(line)
        if m:
            cur = m.group(1)
            continue
        am = ADDR_RE.search(line)
        if am and cur:
            lm = BANKED_LABEL_RE.match(cur)
            if lm:
                key = (lm.group("pre") or "", int(lm.group("addr"), 16))
            else:
                key = ("", int(am.group(1), 16))
            addr_of.setdefault(key, cur)
            cur = None

    missing = []
    for (prefix, addr) in names:
        if (prefix, addr) in addr_of:
            continue
        if names[(prefix, addr)] in defined:
            continue
        if prefix == "" and addr in insn_addrs:
            continue
        missing.append((prefix, addr))
    if missing:
        raise SystemExit("error: map names an address with no label "
                         "definition: " + ", ".join(
                             f"{p or ''}${a:04X}" for p, a in missing))

    # The new names must not collide with a label that is not being renamed.
    being_renamed = set(addr_of.get(k) for k in names)
    existing = {n for n in defined if n not in being_renamed}
    clash = [names[k] for k in names if names[k] in existing]
    if clash:
        raise SystemExit("error: renamed label would collide: "
                         + ", ".join(sorted(set(clash))))

    # the addresses that already carry a label definition
    addr_with_label = set(addr_of)

    # old label spelling -> new name, for every label the map renames.
    # One combined pattern, so a 15,000-line file does not need 15,000
    # separate substitutions per label.
    renamed_labels = {}
    for key, old_name in addr_of.items():
        if key in names and names[key] != old_name:
            renamed_labels[old_name] = names[key]
    renamed_re = None
    if renamed_labels:
        pattern = "|".join(re.escape(n) for n in sorted(renamed_labels,
                                                       key=len, reverse=True))
        renamed_re = re.compile(r"\b(" + pattern + r")\b")

    out = []
    renamed_defs = 0
    renamed_refs = 0
    renamed_calls = 0
    renamed_headers = 0
    created_defs = 0
    equates = []
    cur = None
    cur_prefix = ""
    for line in lines:
        m = LABEL_RE.match(line)
        if m:
            old = m.group(1)
            lm = BANKED_LABEL_RE.match(old)
            key = None
            if lm:
                cand = (lm.group("pre") or "", int(lm.group("addr"), 16))
                if cand in addr_of and addr_of[cand] == old:
                    key = cand
            if key is not None and key in names and names[key] != old:
                out.append(names[key] + ":")
                renamed_defs += 1
                defined.discard(old)
                defined.add(names[key])
                cur = names[key]
                cur_prefix = key[0]
                continue
            out.append(line)
            cur = old
            cur_prefix = lm.group("pre") or "" if lm else cur_prefix
            continue

        am = ADDR_RE.search(line)
        line_prefix = cur_prefix
        if am:
            line_prefix = prefix_at(bank_prefixes, file_offset_of(am))

        # A call target that carries no label at all - the reclass can
        # leave a freshly decoded entry unlabelled - gets one from the map.
        # A project that uses ca65's '@local' labels cannot have a 'Name:'
        # line inserted mid-routine: that re-scopes every '@local' after the
        # split and breaks the forward '@refs'.  There the name becomes an
        # equate in locals.inc instead, which names the address without
        # touching the label scope.
        if am:
            a = int(am.group(1), 16)
            if ("", a) in names and ("", a) not in addr_with_label:
                if use_equates:
                    equates.append((names[("", a)], a))
                else:
                    out.append(f"{names[('', a)]}:")
                    addr_with_label.add(("", a))
                    created_defs += 1

        # The operand token lives in the code part; the trailing audit
        # comment is preserved byte for byte.
        code, sep, comment = line.partition(";")

        # The routine header block names the routine it introduces; it has
        # to follow the rename or the header would point at a label that
        # no longer exists.
        hm = re.match(r"^(\s*Routine:\s*)(\S+)(\s*)$", comment)
        if hm:
            old = hm.group(2)
            lm = BANKED_LABEL_RE.match(old)
            if lm:
                cand = (lm.group("pre") or "", int(lm.group("addr"), 16))
                if cand in names:
                    comment = hm.group(1) + names[cand] + hm.group(3)
                    renamed_headers += 1

        def repl(mo):
            nonlocal renamed_refs
            a = int(mo.group(1), 16)
            cand = (line_prefix, a)
            if cand in names:
                renamed_refs += 1
                return names[cand]
            return mo.group(0)
        code = re.sub(r"\bLabel_([0-9A-F]{4})\b", repl, code)

        # A label the map renames can carry any spelling - `SubC<addr>` is
        # the honest residue form - so its references have to follow it.
        if renamed_re is not None:
            code = renamed_re.sub(lambda mo: renamed_labels[mo.group(1)], code)

        # raw $XXXX on jsr/jmp
        cm = CALL_RE.match(code)
        if cm:
            a = int(cm.group(2), 16)
            # A call into $E000-$FFFF always lands in the fixed bank,
            # whichever bank the caller runs in.
            cand = ("", a) if 0xE000 <= a <= 0xFFFF else (line_prefix, a)
            if cand not in names and ("", a) in names and line_prefix:
                cand = ("", a)
            if cand in names:
                code = f"{cm.group(1)}{names[cand]}{cm.group(3)}"
                renamed_calls += 1
        out.append(code + sep + comment)

    print(f"{renamed_defs} label definitions renamed, {created_defs} created, "
          f"{renamed_refs} label references, {renamed_calls} raw call operands, "
          f"{renamed_headers} routine headers renamed "
          f"({len(names)} entries in the map)")
    if args.write:
        prg.write_text("\n".join(out))
        if equates:
            # The equates go in locals.inc, which prg.s includes.
            loc = prg.with_name("locals.inc")
            text = loc.read_text() if loc.is_file() else ""
            if not text.endswith("\n"):
                text += "\n"
            have = {m.group(1) for m in
                    re.finditer(r"^([A-Za-z_][A-Za-z0-9_]*)\s*=", text, re.M)}
            new = [(n, a) for n, a in equates if n not in have]
            if new:
                if "Sub-entry equates" not in text:
                    text += ("\n; Sub-entry equates: addresses that carry no label\n"
                             "; of their own.  They are equates rather than\n"
                             "; mid-routine labels so that no ca65 '@local' scope is\n"
                             "; split.\n")
                for name, addr in new:
                    text += f"{name:<24s} = ${addr:04X}\n"
                loc.write_text(text)
                print(f"{len(new)} equate(s) written to {loc.name}")
        print("written")
    return 0


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