#!/usr/bin/env python3
"""Substitute semantic zero-page aliases into a project's prg.s.

`apply_symbols.py` only rewrites *numeric* operands, so it cannot rename an
operand that is already written as a neutral `ZpXX` name.  The multi-alias
Pass 3 work therefore needs a direct textual replacement at the specific
use sites, which is what this script does.

Two modes:

  --global NAME=OLD   replace every occurrence of the neutral name OLD with
                       NAME (word boundary), for addresses where *every* use
                       has the same proven role.

  --site NAME=OLD@CPU replace OLD with NAME only on the instruction line
                       whose audit comment carries CPU address CPU.

Both modes touch the operand token only.  The mnemonic, the whitespace and
the trailing `; CPU:$.... PRG:....  XX XX` audit comment are preserved byte
for byte, so the byte-identity audit trail survives.  Because every
replacement names a symbol with the same value as the one it replaces, the
assembled bytes cannot change - `make` still has to print BYTE-IDENTICAL.

Usage:
    python3 _tools/apply_zp_aliases.py <project dir> --alias PpuCtrlShadow=ZpC2 \
        --at PtrLo=Zp00@E742 ... [--write]
"""

import argparse
import re
from pathlib import Path

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


def parse_global(spec):
    name, _, old = spec.partition("=")
    return name.strip(), old.strip()


def parse_site(spec):
    name, _, rest = spec.partition("=")
    old, _, cpu = rest.partition("@")
    return name.strip(), old.strip(), cpu.strip().upper().lstrip("$")


def replace_token(text, old, new):
    """Replace one operand token.

    A bare identifier is matched on word boundaries so `Zp01` cannot match
    inside `Zp010`.  Anything with punctuation (an indirect operand such as
    `($0017)`) is matched literally - the caller passes the whole token
    including its parentheses, because dropping them would change the
    addressing mode.

    A leading `a:` ca65 prefix is kept: it forces absolute addressing, and
    dropping it would turn `sta a:Zp17,y` into a zero-page store and change
    the bytes.
    """
    old_prefix = old.startswith("a:")
    new_prefix = new.startswith("a:")
    if old_prefix and not new_prefix:
        new = "a:" + new
    if not old_prefix and re.fullmatch(r"[A-Za-z_][A-Za-z0-9_]*", old):
        return re.sub(r"\b" + re.escape(old) + r"\b", new, text)
    return text.replace(old, new)


def replace_in_code(line, old, new):
    """Replace a token in the code part of the line only.

    The trailing `; CPU:$.... PRG:....  XX XX` audit comment is preserved
    byte for byte, as the tool contract requires.
    """
    code, sep, comment = line.partition(";")
    return replace_token(code, old, new) + sep + comment


def replace_indirect(line, name):
    """Replace an indirect operand `($XX)` with `(name)`.

    The indirect jump reads a 16-bit pointer from the zero page, so the
    operand keeps its parentheses and only the address is named.
    """
    code, sep, comment = line.partition(";")
    code = re.sub(r"\(\$[0-9A-Fa-f]{2,4}\)", f"({name})", code)
    return code + sep + comment


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("dir", help="project directory, e.g. '_wip/NROM/Raid (U) [a2]'")
    ap.add_argument("--alias", action="append", default=[], metavar="NAME=OLD",
                    help="replace every occurrence of OLD with NAME")
    ap.add_argument("--at", action="append", default=[], metavar="NAME=OLD@CPU",
                    help="replace OLD with NAME on the instruction at CPU")
    ap.add_argument("--at-ind", action="append", default=[], metavar="NAME@CPU",
                    help="replace the indirect operand on the instruction at "
                         "CPU with (NAME) - for `jmp (zp)` style operands")
    ap.add_argument("--at-file",
                    help="TSV of <CPU address><TAB>OLD<TAB>NEW, one per site. "
                         "Address-based, so the file survives edits that move "
                         "lines around.")
    ap.add_argument("--sites", help="TSV of <line><TAB>OLD<TAB>NEW, one per site")
    ap.add_argument("--write", action="store_true", help="apply (default: dry run)")
    args = ap.parse_args()

    prg = Path(args.dir) / "prg.s"
    if not prg.is_file():
        raise SystemExit(f"error: {prg} not found")

    lines = prg.read_text().split("\n")
    aliases = [parse_global(s) for s in args.alias]
    sites = [parse_site(s) for s in args.at]
    ind_sites = []
    for spec in args.at_ind:
        name, _, cpu = spec.partition("@")
        ind_sites.append((name.strip(), cpu.strip().upper().lstrip("$")))

    # Line-based per-site choices: <line><TAB>OLD<TAB>NEW
    by_line = {}
    if args.sites:
        for raw in Path(args.sites).read_text().splitlines():
            line = raw.split("#", 1)[0].strip()
            if not line:
                continue
            parts = line.split("\t")
            if len(parts) < 3:
                parts = line.split()
            by_line[int(parts[0])] = (parts[1].strip(), parts[2].strip())

    # Address-based per-site choices: <CPU><TAB>OLD<TAB>NEW
    by_addr = {}
    if args.at_file:
        for raw in Path(args.at_file).read_text().splitlines():
            line = raw.split("#", 1)[0].strip()
            if not line:
                continue
            parts = line.split("\t")
            if len(parts) < 3:
                parts = line.split()
            cpu = parts[0].strip().upper().lstrip("$")
            by_addr[cpu] = (parts[1].strip(), parts[2].strip())

    out = []
    changed = 0
    for lineno, line in enumerate(lines, 1):
        new = line
        for name, old in aliases:
            new = replace_in_code(new, old, name)
        if sites or ind_sites:
            m = CPU_RE.search(new)
            if m:
                cpu = m.group(1).upper()
                for name, old, want in sites:
                    if want == cpu:
                        new = replace_in_code(new, old, name)
                for name, want in ind_sites:
                    if want == cpu:
                        new = replace_indirect(new, name)
        if by_addr:
            m = CPU_RE.search(new)
            if m:
                cpu = m.group(1).upper()
                if cpu in by_addr:
                    old, name = by_addr[cpu]
                    new = replace_in_code(new, old, name)
        if lineno in by_line:
            old, name = by_line[lineno]
            new = replace_in_code(new, old, name)
        if new != line:
            changed += 1
        out.append(new)

    print(f"{changed} line(s) changed")
    if args.write:
        prg.write_text("\n".join(out))
        print("written")
    else:
        for a, b in zip(lines, out):
            if a != b:
                print(f"  - {a}")
                print(f"  + {b}")
    return 0


if __name__ == "__main__":
    raise SystemExit(main())