#!/usr/bin/env python3
"""Re-align canonical audit comments after a Pass 3 symbol substitution.

`nesdis.py` emits every instruction as

    f"    {text:<31}; CPU:${cpu:04X} PRG:{prg:04X}  {raw}"

so the audit comment starts at column 35 (0-indexed) for every instruction
whose code text is 31 characters or shorter.  `apply_symbols.py` replaces an
operand in place, so a substitution that lengthens the operand (e.g.
`$2001` -> `PPUMASK`) pushes the audit comment to the right and breaks the
column.  The repo convention (AGENTS.md) is that "comment/audit alignment is
re-done after substitutions".

This tool restores the canonical column.  It is whitespace-only: the CPU
address, PRG offset and byte list in the comment are never touched, so the
build cannot change.

Data directives (`.byte`, `.word`, `.dword`, `.incbin`) are left alone - the
emitter writes those with two spaces before the comment, not a padded
column.  Lines whose code text is longer than 31 characters are also left
alone, because the emitter cannot pad those into the column either.

Usage:
    python3 _tools/align_audit_comments.py "NROM/Pac-Man (U)"          # dry run
    python3 _tools/align_audit_comments.py "NROM/Pac-Man (U)" --write
"""

import argparse
import re
from pathlib import Path

# Canonical audit comment: ; CPU:$XXXX PRG:XXXX  AA BB
AUDIT_RE = re.compile(
    r"^(    \S.*?)\s*(; CPU:\$[0-9A-Fa-f]{4} PRG:[0-9A-Fa-f]{4}.*)$")

# Data directives keep the emitter's two-space form.
DATA_RE = re.compile(r"^    \.(byte|word|dword|incbin|res)\b")

CODE_WIDTH = 31          # f"    {text:<31}" -> ';' lands at column 35
COMMENT_COL = 4 + CODE_WIDTH
# The emitter's field is 31 characters wide *for the code text alone*, so an
# indented code of up to 4 + 31 characters still fits the column.  Testing the
# indented length against 31 (as this tool used to) left code texts of 28-31
# characters un-padded, which is what produced the ragged column-34/36/37
# comments in a hand-edited source.
MAX_CODE = 4 + CODE_WIDTH  # longest indented code that still fits the column


def align(text):
    """Return (new_text, changed_line_count)."""
    out = []
    changed = 0
    for line in text.split("\n"):
        m = AUDIT_RE.match(line)
        if m and not DATA_RE.match(line):
            code, comment = m.group(1), m.group(2)
            if len(code) <= MAX_CODE and line.index(";") != COMMENT_COL:
                new = code.ljust(COMMENT_COL) + comment
                if new != line:
                    changed += 1
                out.append(new)
                continue
        out.append(line)
    return "\n".join(out), changed


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("dir", help="project directory, e.g. 'NROM/Pac-Man (U)'")
    ap.add_argument("--write", action="store_true",
                    help="apply the change (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")

    text = prg.read_text()
    new, changed = align(text)
    if changed:
        print(f"re-aligned {changed} audit comment(s) to column {COMMENT_COL}")
        first = new.split("\n")[0]
        for a, b in zip(text.split("\n"), new.split("\n")):
            if a != b:
                print(f"  - {a}")
                print(f"  + {b}")
                break
    else:
        print("already aligned: no audit comment column drift found")

    if args.write and changed:
        prg.write_text(new)
        print("written")
    elif args.write:
        print("nothing to write")
    return 0


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