#!/usr/bin/env python3
"""apply_symbols.py -- substitute symbol names for numeric operands in a
disassembled prg.s.

Names come from two include files sitting next to prg.s:
  symbols.inc  NES hardware constants ($2000-$4017)
  locals.inc   zero-page/RAM names plus "forward reference" aliases for
               routine addresses (e.g. AbsoluteDifference = $F08E)

Both are simple `Name = $HEX` tables.  A few addresses carry more than one
name (the same byte is reused by different subsystems); those are called
multi-alias addresses and cannot be resolved from the address alone.

Usage:
  apply_symbols.py <project-dir> [--sites sites.tsv] [--labels lbl] [--write] [--report]

  --sites   TSV of `line<TAB>name` overrides for multi-alias addresses.
            Lines not listed are left numeric (recorded as residue).
  --labels  ld65 `-Ln` label file (or TSV `name<TAB>$hex`). Operands that point
            at a label *defined in prg.s* ($8000+, e.g. a data table) are named
            with it. Only labels prg.s actually defines are used, which keeps
            segment labels from the CHR/header files out of the PRG.
  --write   rewrite prg.s in place, else print stats only.
  --report  list the multi-alias operands left numeric (the residue).

Only the operand token is rewritten: 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.

Addressing width is preserved: an operand originally written with four hex
digits is absolute, so a zero-page-valued symbol gets the ca65 `a:` prefix
to stop ca65 from re-optimising it into zero-page addressing.

Note: `symbols.inc` also holds controller bit masks (PAD_RIGHT=$01, PAD_A=$80, ...).
Those are constants, not addresses, so only $2000-$4017 is treated as a register
window; without that guard a plain zero-page operand like `$10` would be renamed to
a mask.
"""
import argparse
import re
import sys
from collections import Counter, defaultdict
from pathlib import Path

NAME_RE = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*)\s*=\s*" + re.escape("$") + r"([0-9A-Fa-f]+)\s*$")
# the controller's button bits, in the order the NES reports them
BUTTON_BIT = {"PAD_A": 0x01, "PAD_B": 0x02, "PAD_SELECT": 0x04,
              "PAD_START": 0x08, "PAD_UP": 0x10, "PAD_DOWN": 0x20,
              "PAD_LEFT": 0x40, "PAD_RIGHT": 0x80}
INSTR_RE = re.compile(r"^(\s*[a-zA-Z]{3}\s+)(\S+)(\s*)$")
PLAIN_RE = re.compile(r"^(a:)?(\$[0-9A-Fa-f]{2,4})(,[xXyY])?$")
IND_RE = re.compile(r"^(a:)?\((\$[0-9A-Fa-f]{2,4})(,[xXyY])?\)(,[yY])?$")


def load_names(path):
    by_name, by_addr = {}, defaultdict(list)
    for raw in Path(path).read_text().splitlines():
        m = NAME_RE.match(raw.split(";", 1)[0].strip())
        if m:
            by_name[m.group(1)] = int(m.group(2), 16)
            by_addr[int(m.group(2), 16)].append(m.group(1))
    return by_name, by_addr


LBL_RE = re.compile(r"^al ([0-9A-Fa-f]{4,6}) \.([A-Za-z_][A-Za-z0-9_]*)$")


def load_labels(path, defined):
    """Address -> label name from an ld65 label file or a `name<TAB>$hex` TSV.

    Restricted to labels the source actually defines and that live in the PRG
    ($8000+), so segment labels from the CHR/header files cannot leak in.
    """
    out = {}
    for raw in Path(path).read_text().splitlines():
        raw = raw.strip()
        if not raw:
            continue
        m = LBL_RE.match(raw)
        if m:
            addr, name = int(m.group(1), 16), m.group(2)
        else:
            parts = raw.split("\t")
            hexval = parts[1].strip().lstrip("$") if len(parts) == 2 else ""
            if len(parts) != 2 or not re.fullmatch(r"[0-9A-Fa-f]{2,6}", hexval):
                continue
            addr, name = int(hexval, 16), parts[0].strip()
        if addr >= 0x8000 and name in defined:
            out[addr] = name
    return out


def operand_address(op):
    """(address, kind, inner, outer, hex_width) for a numeric operand, else Nones.

    kind          'plain' or 'ind'
    inner         index token inside the parens ('' / ',x' / ',y')
    outer         trailing index token ('' / ',y')
    """
    m = PLAIN_RE.match(op)
    if m:
        hexdigits = m.group(2)[1:]
        return int(hexdigits, 16), "plain", (m.group(3) or ""), "", len(hexdigits)
    m = IND_RE.match(op)
    if m:
        hexdigits = m.group(2)[1:]
        return (int(hexdigits, 16), "ind", (m.group(3) or ""),
                (m.group(4) or ""), len(hexdigits))
    return None, None, None, None, None


def render(name, kind, inner, outer, width, value):
    if kind == "plain":
        base = name + inner
        if width == 4 and value < 0x100:
            return "a:" + base
        return base
    return "(" + name + inner + ")" + outer


def resolve_multi(names, code):
    """Pick one of several names for the same address.

    A single-bit test on the value - `and #$10` - means the value is a
    controller button, and the button names say which.  Otherwise the
    neutral `ZpXX` name is the honest one: the address's role is not
    established, and the neutral name says so.

    Returns None when neither rule applies.
    """
    buttons = {n for n in names if n.startswith("PAD_")}
    if len(buttons) == 1:
        bit = BUTTON_BIT.get(next(iter(buttons)))
        m = re.search(r"\band\s+#\$([0-9A-F]{2})\b", code)
        if bit is not None and m and int(m.group(1), 16) == bit:
            return next(iter(buttons))
    neutral = [n for n in names if re.fullmatch(r"Zp[0-9A-F]{2}", n)]
    if len(neutral) == 1:
        return neutral[0]
    return None


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("dir")
    ap.add_argument("--sites")
    ap.add_argument("--labels")
    ap.add_argument("--write", action="store_true")
    ap.add_argument("--report", action="store_true")
    args = ap.parse_args()

    d = Path(args.dir)
    prg = d / "prg.s"
    _, hw_addr = load_names(d / "symbols.inc")
    _, loc_addr = load_names(d / "locals.inc")

    # Only the register window of symbols.inc is an address space; the
    # controller bit masks ($01, $80, ...) are constants, not addresses.
    hw = {a: names[0] for a, names in hw_addr.items()
          if names and 0x2000 <= a <= 0x4017}
    single = {a: names[0] for a, names in loc_addr.items() if len(names) == 1}
    multi = {a: names for a, names in loc_addr.items() if len(names) > 1}
    for a in hw:                       # hardware names win where both exist
        single.pop(a, None)

    sites = {}
    if args.sites:
        for raw in Path(args.sites).read_text().splitlines():
            raw = raw.strip()
            if not raw or raw.startswith("#"):
                continue
            line, _, name = raw.partition("\t")
            sites[int(line)] = name.strip()

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

    labels = {}
    if args.labels:
        defined = {m.group(1) for m in
                   (re.match(r"^([A-Za-z_][A-Za-z0-9_]*):", l) for l in lines) if m}
        labels = load_labels(args.labels, defined)

    counts = Counter()
    residue = []
    out = []
    for lineno, line in enumerate(lines, 1):
        code, sep, comment = line.partition(";")
        m = INSTR_RE.match(code)
        if not m:
            out.append(line)
            continue
        op = m.group(2)
        addr, kind, inner, outer, width = operand_address(op)
        if addr is None:
            out.append(line)
            continue
        if lineno in sites:
            name, tag = sites[lineno], "site"
        elif addr in labels:
            name, tag = labels[addr], "label"
        elif addr in single:
            name, tag = single[addr], "single"
        elif addr in hw:
            name, tag = hw[addr], "hardware"
        elif addr in multi:
            # The address carries more than one name - the zero page is
            # reused across modes.  Two rules settle most of them:
            # a single-bit test on the value is a controller button, and a
            # neutral `ZpXX` name is the honest one otherwise.
            name = resolve_multi(multi[addr], code)
            if name is None:
                counts["residue"] += 1
                residue.append((lineno, op, multi[addr]))
                out.append(line)
                continue
            tag = "multi"
        else:
            out.append(line)
            continue
        newop = render(name, kind, inner, outer, width, addr)
        if newop == op:
            out.append(line)
            continue
        counts[tag] += 1
        out.append(m.group(1) + newop + m.group(3) + sep + comment)

    total = (counts["site"] + counts["single"] + counts["hardware"]
             + counts["label"])
    print(f"substituted {total} operands: single={counts['single']} "
          f"hardware={counts['hardware']} site={counts['site']} "
          f"label={counts['label']}")
    print(f"multi-alias residue left numeric: {counts['residue']}")
    if args.report:
        for lineno, op, names in residue:
            print(f"  line {lineno}: {op}  candidates={names}")

    if args.write:
        prg.write_text("\n".join(out))
    return 0


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