#!/usr/bin/env python3
"""Add byte-identity audit comments to a project's prg.s.

Why this exists
---------------
Most projects in this repo carry a trailing audit comment on every
instruction, e.g.

    lda #$00          ; CPU:$C000 PRG:0000  A9 00

That comment is the byte-identity audit trail: it ties a source line back to
exact ROM bytes. Two projects lack it (Pac-Man has none at all, Excitebike has
it on only ~19% of code lines), so a source line in those files cannot be
mapped to ROM bytes without re-assembling.

How it works
------------
`ca65 -l` produces a listing with one entry per *emitted* source line:

    000186r 1  A9 00                lda #$00
    ^offset ^n   ^bytes              ^source text (echoed verbatim)

The listing has no source line numbers, but it does echo the source text, and
the entries appear in source order. So we recover the mapping by walking both
sequences and matching the echoed text, tolerating the lines ca65 skips.

Byte *values* are read from prg.bin at the listing offset, never from the
listing itself: ca65 prints `rr rr` for relocatable words (e.g. `.word Label`)
because the value is only known at link time. The count `n` and the token count
of the byte column give the length; prg.bin gives the truth.

Correctness
-----------
The script never edits instruction text, only appends a comment, so assembly is
unaffected. It then re-assembles and checks byte identity, and independently
verifies that every comment it wrote matches prg.bin at the offset it claims.
A disagreement is a hard error, not a warning.

Usage
-----
    python3 _tools/add_audit_comments.py "Excitebike (JU)"          # apply
    python3 _tools/add_audit_comments.py "Excitebike (JU)" --dry-run
    python3 _tools/add_audit_comments.py "Excitebike (JU)" --force   # re-comment
"""

import argparse
import difflib
import os
import re
import shutil
import subprocess
import sys
import tempfile

REPO = os.path.dirname(os.path.dirname(os.path.abspath(__file__)))

# <6-hex offset>r <count> <rest>
LISTING_RE = re.compile(r"^([0-9A-Fa-f]{6})r\s+(\d+)\s(.*)$")

# The byte column: space separated 2-hex tokens or `rr` for relocatable, then
# two-or-more spaces before the echoed source text.
BYTECOL_RE = re.compile(
    r"^\s*((?:[0-9A-Fa-f]{2}|rr)(?: (?:[0-9A-Fa-f]{2}|rr))*)?\s\s+(.*)$"
)

# A line that already carries some spelling of audit comment. Several
# spellings are in use across the repo; all of them count as "already done".
HAS_AUDIT = re.compile(r";\s*(?:CPU:\$|PRG:|OFF:)|;\s*\$?[0-9A-Fa-f]{4}:\s")

# A source line that emits bytes.
EMITS_RE = re.compile(
    r"^(?:[a-z]{3}\b|\.(?:byte|word|dword|addr|res|incbin)\b|"
    r"(?:[A-Za-z_][\w]*)\s*[:=]\s*\.(?:byte|word|res|incbin)\b|\w+\s*=|=)"
)


def run(cmd, cwd):
    p = subprocess.run(cmd, cwd=cwd, capture_output=True, text=True)
    if p.returncode != 0:
        sys.stderr.write(f"$ {' '.join(cmd)}\n{p.stdout}{p.stderr}\n")
        raise SystemExit(f"command failed: {' '.join(cmd)}")
    return p.stdout


def load_equates(*paths):
    """Map symbol name -> integer value from `NAME = $xx` lines."""
    eq = {}
    rx = re.compile(r"^\s*([A-Za-z_@][\w.]*)\s*=\s*\$?([0-9A-Fa-f]+)\s*(?:;.*)?$")
    for p in paths:
        if not os.path.exists(p):
            continue
        with open(p, errors="replace") as fh:
            for line in fh:
                m = rx.match(line)
                if m:
                    try:
                        eq[m.group(1)] = int(m.group(2), 16)
                    except ValueError:
                        pass
    return eq


def instruction_length(mnem, operand, equates=None):
    """Length in bytes of one 6502 instruction, from mnemonic + operand text.

    Used only by the .org fallback path, where the ca65 listing cannot be
    trusted for offsets. Addressing-mode table, official opcodes only: the
    projects this fallback targets do not use undocumented opcodes in the
    regions it annotates.
    """
    equates = equates or {}
    op = operand.lstrip()
    imm = op.startswith("#")
    indirect = op.startswith("(")
    core = op.strip()
    if core.startswith("("):
        core = core[1:]
    bits = core.split(",")
    indexed_x = bits[-1].strip().lower() == "x" if len(bits) > 1 else False
    indexed_y = bits[-1].strip().lower() == "y" if len(bits) > 1 else False
    core = bits[0].strip()
    # "(ZpFA),y" -> the closing paren sits on the pointer field, not the end.
    if core.endswith(")"):
        core = core[:-1].strip()

    # A symbolic operand is zero page if its equate is < $100. Without the
    # equate we cannot tell, so default to absolute (the 3-byte form), which
    # is what an unresolved `ZpXX` name is in practice.
    if re.fullmatch(r"[A-Za-z_@][\w.]*", core):
        val = equates.get(core)
        zp = val is not None and val < 0x100
        abs_ = not zp
    else:
        zp = re.fullmatch(r"\$[0-9A-Fa-f]{1,2}", core) is not None
        abs_ = re.fullmatch(r"\$[0-9A-Fa-f]{3,4}", core) is not None

    table = {
        # implied
        "asl": 1, "lsr": 1, "rol": 1, "ror": 1, "bcc": 2, "bcs": 2, "beq": 2,
        "bmi": 2, "bne": 2, "bpl": 2, "bvc": 2, "bvs": 2, "clc": 1, "cld": 1,
        "cli": 1, "clv": 1, "dex": 1, "dey": 1, "inx": 1, "iny": 1, "nop": 1,
        "pha": 1, "php": 1, "pla": 1, "plp": 1, "rti": 1, "rts": 1, "sec": 1,
        "sed": 1, "sei": 1, "tax": 1, "tay": 1, "tsx": 1, "txa": 1, "txs": 1,
        "tya": 1,
        # accumulator / implied-by-A
        "adc": None, "and": None, "cmp": None, "eor": None, "lda": None,
        "ora": None, "sbc": None, "sta": None,
        # branches
        "jmp": None, "jsr": 3,
        # stack
        "ldx": None, "ldy": None, "stx": None, "sty": None,
    }
    if mnem in ("jmp",):
        return 3 if not indirect else 5
    if mnem == "jsr":
        return 3
    base = table.get(mnem, None)
    if mnem in ("adc", "and", "cmp", "eor", "lda", "ora", "sbc", "sta"):
        if imm:
            return 2
        if indirect:
            # (zp,x) is 2 bytes; (zp),y is 2; (abs),y is 6; (abs) is 5.
            if indexed_x:
                return 2
            if indexed_y:
                return 6 if abs_ else 2
            return 6 if abs_ else 5
        if indexed_x or indexed_y:
            return 2 if zp else 3
        return 2 if zp else 3
    if mnem in ("ldx", "ldy"):
        if imm:
            return 2
        if indexed_y if mnem == "ldx" else indexed_x:
            return 2 if zp else 3
        if indexed_x or indexed_y:
            return 2 if zp else 3
        return 2 if zp else 3
    if mnem in ("stx", "sty"):
        if indexed_x or indexed_y:
            return 2 if zp else 3
        return 2 if zp else 3
    return base if base else 2


def offsets_from_labels(src_lines, lbl_path, base, cpu_base):
    """Fallback for sources using `.org`, where ca65's listing offsets restart.

    The linker's .lbl file gives every label's exact address. We align labels to
    their defining source line, then walk forward accumulating instruction
    lengths to recover each emitting line's PRG offset.
    """
    labels = {}
    with open(lbl_path, errors="replace") as fh:
        for line in fh:
            parts = line.split()
            if len(parts) >= 3 and parts[0] == "al":
                addr = int(parts[1], 16)
                name = parts[2].lstrip(".")
                labels[name] = addr

    label_re = re.compile(r"^([A-Za-z_@][\w.]*)\s*:")
    mapping = {}
    unresolved = []
    for i, raw in enumerate(src_lines):
        s = raw.strip()
        if not s or s.startswith(";"):
            continue
        m = label_re.match(s)
        if not m:
            continue
        name = m.group(1)
        if name.startswith("@"):
            continue  # cheap local labels are not exported to the .lbl file
        if name in labels:
            mapping[i] = labels[name] - cpu_base
        else:
            unresolved.append(name)
    return mapping, unresolved


def propagate_offsets(src_lines, anchors, equates=None):
    """Fill in PRG offsets for emitting lines between known label anchors.

    `anchors` maps source line index -> PRG offset for label definitions. We
    walk forward from each anchor, adding the length of every emitting line,
    until the next anchor. Data directives (.byte/.word/.res) are counted too.
    """
    data_re = re.compile(r"^\.(?:byte|word|dword|addr|res|incbin)\b")
    mnem_re = re.compile(r"^([a-z]{3})\b\s*(.*)$")
    lengths = {}
    idxs = sorted(anchors)
    for n, start in enumerate(idxs):
        off = anchors[start]
        i = start + 1
        stop = idxs[n + 1] if n + 1 < len(idxs) else len(src_lines)
        while i < stop:
            s = src_lines[i].strip()
            if not s or s.startswith(";"):
                i += 1
                continue
            if data_re.match(s):
                parts = s.split(None, 1)
                d = parts[0].lower()
                if d == ".byte":
                    nbytes = len([x for x in parts[1].split(",") if x.strip()]) if len(parts) > 1 else 0
                elif d in (".word", ".addr"):
                    nbytes = 2 * (len([x for x in parts[1].split(",") if x.strip()]) if len(parts) > 1 else 0)
                elif d == ".res":
                    mm = re.search(r"\$?([0-9A-Fa-f]+)", parts[1]) if len(parts) > 1 else None
                    nbytes = int(mm.group(1), 16) if mm else 0
                else:  # .incbin
                    nbytes = 0
                anchors[i] = off
                lengths[i] = nbytes
                off += nbytes
                i += 1
                continue
            mm = mnem_re.match(s)
            if mm and not s.endswith("=") and ":" not in s.split()[0]:
                nbytes = instruction_length(mm.group(1), mm.group(2), equates)
                anchors[i] = off
                lengths[i] = nbytes
                off += nbytes
            i += 1
    return lengths


def parse_listing(path):
    """Return [(prg_offset, n_tokens_or_None, echoed_text)] in listing order."""
    recs = []
    with open(path, errors="replace") as fh:
        for line in fh:
            m = LISTING_RE.match(line.rstrip("\n"))
            if not m:
                continue
            rest = m.group(3)
            bm = BYTECOL_RE.match(rest)
            if bm and bm.group(1):
                tokens = bm.group(1).split()
                text = bm.group(2).rstrip()
                recs.append((int(m.group(1), 16), len(tokens), text))
            else:
                # No byte column: a comment, equate, directive or label. The
                # echoed text is whatever follows the count.
                recs.append((int(m.group(1), 16), None, rest.rstrip()))
    return recs


def align(src_lines, recs):
    """Map source line index -> (prg_offset, nbytes) using echoed text.

    ca65 skips some source lines (blank lines, continuation lines, macro
    expansions) and emits entries for others only once, so this has to be a
    real sequence alignment rather than a lockstep zip. We use difflib's
    opcodes: the listing and the source should otherwise be identical, so the
    few genuine insertions/deletions are cheap to find and everything else
    lines up exactly.

    Returns (mapping, n_unmatched).
    """
    src_norm = [l.strip() for l in src_lines]
    rec_norm = [t.strip() for _, _, t in recs]

    matcher = difflib.SequenceMatcher(a=src_norm, b=rec_norm, autojunk=False)
    mapping = {}
    unmatched = 0
    for tag, i1, i2, j1, j2 in matcher.get_opcodes():
        if tag == "equal":
            for k in range(i2 - i1):
                mapping[i1 + k] = (recs[j1 + k][0], recs[j1 + k][1])
        else:
            # 'replace' can still line up when the two sequences are the same
            # length but differ in text; 'delete'/'insert' do not.
            if tag == "replace" and (i2 - i1) == (j2 - j1):
                for k in range(i2 - i1):
                    mapping[i1 + k] = (recs[j1 + k][0], recs[j1 + k][1])
            else:
                unmatched += (j2 - j1)
    return mapping, unmatched


def audit_comment(prg, off, nbytes, base, cpu_base):
    """Build `; CPU:$xxxx PRG:xxxx AA BB` from prg.bin, or None."""
    if nbytes is None or nbytes == 0:
        return None
    if off + nbytes > len(prg):
        return None
    vals = prg[off:off + nbytes]
    hexs = " ".join(f"{b:02X}" for b in vals)
    return f"; CPU:${cpu_base + off:04X} PRG:{off:04X}  {hexs}"


def resolve_project(repo, name):
    """Find a project folder under its mapper directory (or _wip/<MAPPER>/)."""
    for mapper in ("NROM", "AOROM", "CNROM", "MMC1", "MMC3", "MMC5", "UNROM"):
        for tier in ("", "_wip"):
            q = os.path.join(repo, tier, mapper, name)
            if os.path.isdir(q):
                return q
    raise SystemExit(f"no such project dir: {os.path.join(repo, '*/', name)}")


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("project")
    ap.add_argument("--dry-run", action="store_true")
    ap.add_argument("--force", action="store_true",
                    help="re-comment lines that already have an audit comment")
    ap.add_argument("--prg-name", default="prg.bin")
    ap.add_argument("--base", default="0xC000",
                    help="CPU address corresponding to PRG offset 0")
    args = ap.parse_args()

    cpu_base = int(args.base, 0)
    copilot = resolve_project(REPO, args.project)

    prg_path = os.path.join(copilot, args.prg_name)
    src_path = os.path.join(copilot, "prg.s")
    prg = open(prg_path, "rb").read()
    src_lines = open(src_path, errors="replace").read().splitlines()
    equates = load_equates(
        os.path.join(copilot, "symbols.inc"),
        os.path.join(copilot, "locals.inc"),
        src_path,
    )

    with tempfile.TemporaryDirectory() as td:
        # ca65 uses the -l name verbatim (no .lst suffix appended).
        lst = os.path.join(td, "listing")
        run(["ca65", "-g", "-l", lst, "-o", os.path.join(td, "prg.o"), "prg.s"], copilot)
        # Parse inside the context: the listing is deleted on exit.
        recs = parse_listing(lst)

    mapping, unmatched = align(src_lines, recs)

    # A source using `.org` makes ca65's listing offsets restart at zero, so
    # the listing carries no usable addresses. Detect that and fall back to
    # deriving offsets from the linker's label file.
    if recs and max((r[0] for r in recs), default=0) < 0x100 and len(mapping) > 0:
        with tempfile.TemporaryDirectory() as td:
            if not [f for f in os.listdir(copilot) if f.endswith(".o")]:
                # Flat layout: a clean project has no .o files yet; build once.
                print("note: no .o files present; building project first")
                run(["make"], copilot)
            run(["ld65", "-C", "linker.cfg", "-m", os.path.join(td, "m.map"),
                 "-Ln", os.path.join(td, "m.lbl"), "-o", os.path.join(td, "out.nes"),
                 *sorted(f for f in os.listdir(copilot)
                         if f.endswith(".o") and f != "reconstructed.o")], copilot)
            lbl = os.path.join(td, "m.lbl")
            if os.path.exists(lbl):
                print("note: source uses `.org`; deriving offsets from label file")
                anchors, unresolved = offsets_from_labels(src_lines, lbl, 0, cpu_base)
                unmatched = len(unresolved)
                lengths = propagate_offsets(src_lines, anchors, equates)
                mapping = {i: (anchors[i], lengths.get(i)) for i in anchors}

    out = list(src_lines)
    added = skipped = failed = 0
    for idx, (off, nbytes) in mapping.items():
        line = src_lines[idx]
        s = line.strip()
        if not s or s.startswith(";"):
            continue
        if not EMITS_RE.match(s):
            continue
        if not args.force and HAS_AUDIT.search(line):
            skipped += 1
            continue
        cmt = audit_comment(prg, off, nbytes, cpu_base, cpu_base)
        if cmt is None:
            continue
        out[idx] = line.rstrip() + ("  " if line.strip() else " ") + cmt
        added += 1

    # Verify every comment we are about to write against prg.bin.
    bad = 0
    for line in out:
        m = re.search(r";\s*CPU:\$([0-9A-F]{4})\s+PRG:([0-9A-F]{4})\s+((?:[0-9A-F]{2} ?)+)$", line)
        if not m:
            continue
        off = int(m.group(2), 16)
        vals = bytes(int(x, 16) for x in m.group(3).split())
        if prg[off:off + len(vals)] != vals:
            bad += 1
            if bad <= 5:
                sys.stderr.write(f"BYTE MISMATCH PRG {off:04X}: comment {m.group(3)} "
                                 f"vs prg.bin {prg[off:off+len(vals)].hex(' ').upper()}\n")
    if bad:
        raise SystemExit(f"{bad} comments disagree with {args.prg_name}; refusing to write")

    print(f"{args.project}: {added} comments added, {skipped} already present, "
          f"{unmatched} listing entries unmatched")

    if args.dry_run:
        return

    shutil.copyfile(src_path, src_path + ".preaudit")
    with open(src_path, "w") as fh:
        fh.write("\n".join(out) + "\n")
    print(f"wrote {src_path} (backup at {src_path}.preaudit)")


if __name__ == "__main__":
    main()