#!/usr/bin/env python3
"""
Repair a damaged audit-comment trail from a ca65 listing.

Some project sources carry broken audit comments (doubled `??` forms,
comments shifted against their mnemonics after label insertion). This tool
rebuilds the canonical `; CPU:$XXXX PRG:XXXX AA BB` trail using the ca65 -l
listing of prg.s as ground truth for each source line's PRG offset and byte
span; byte values come from prg.bin. Non-audit semantic notes (e.g.
`; state $02`) between the code and the audit tail are preserved.

Sorting out the echo: ca65 prints a source line's text on the FIRST byte
batch it emits (continuation batches follow with empty text), and depth-1
records belong to the main file (depth-2 echoes come from `.include`).

Alignment is a monotonic two-pointer walk: listing groups are strict in-order
echoes of a subsequence of source lines, so each group consumes the next
source line equal to its echoed text.

Safety:
  * refuses to run if any listing byte batch disagrees with prg.bin
  * refuses to write if any `??`-carrying line would remain unmapped, or any
    rewritten comment disagrees with prg.bin
  * leaves the code text (before comments) of every line untouched

Usage: repair_audit_comments.py <project-dir> [--base 0xC000] [--write]
"""

import argparse
import re
import subprocess
import sys
import tempfile
from pathlib import Path

# ca65 2.19 prints `<offset>r <n>  <text>` for a record that emits no bytes
# and `<offset>  <n>  <bytes>  <text>` for one that does; older builds put
# the `r` on every record.  Both are accepted.
LISTING_RE = re.compile(r"^([0-9A-Fa-f]{6})r?\s+(\d+)\s(.*)$")
TOKEN_RE = re.compile(r"^([0-9A-Fa-f]{2}|rr)(?:\s+(.*))?$")

CANONICAL = "; CPU:${cpu:04X} PRG:{prg:04X}  {bytes}"

MNEMONICS = (
    "adc and asl bcc bcs beq bit bmi bne bpl brk bvc bvs clc cld cli clv "
    "cmp cpx cpy dec dex dey eor inc inx iny jmp jsr lda ldx ldy lsr nop "
    "ora pha php pla plp rol ror rti rts sbc sec sed sei sta stx sty "
    "tax tay tsx txa txs tya"
).split()
EMIT_RE = re.compile(
    r"^\s*(?:@?\w*:\s*)?(?:(?:" + "|".join(MNEMONICS) + r")\b|\.byte\b|\.word\b|\.dword\b)",
    re.I,
)

# Audit segments that get replaced: canonical form (with or without bytes,
# including `??`), `$XXXX: AA BB`, `XXXX: AA BB`, `PRG:$XXXX AA BB`.
# The address may be wider than four hex digits: a source whose last
# instruction runs past $FFFF carries a five-digit CPU address.
AUDIT_SEG_RE = re.compile(
    r"^(?:"
    r"CPU:\$[0-9A-Fa-f]{4,6}(?:\s+PRG:[0-9A-Fa-f]{4,6})?(?:\s+[0-9A-Fa-f?]{2})*"
    r"|\$?[0-9A-Fa-f]{4,6}:\s*[0-9A-Fa-f?]{2}(?:\s+[0-9A-Fa-f?]{2})*"
    r"|PRG:\$[0-9A-Fa-f]{4,6}\s+[0-9A-Fa-f?]{2}(?:\s+[0-9A-Fa-f?]{2})*"
    r")$"
)


def parse_listing(path, base=0):
    """Return in-order records (off, tokens_or_None, text) for main-file echoes.

    `off` is the PRG offset.  ca65 prints the address the segment is
    assembled at, so for a project whose PRG is mapped at a CPU address
    (`$C000`, `$8000`) the printed address is that base plus the PRG
    offset; a project whose linker places the PRG at a file offset prints
    the offset itself.
    """
    recs = []
    for line in open(path, errors="replace"):
        m = LISTING_RE.match(line.rstrip("\n"))
        if not m:
            continue
        if int(m.group(2)) != 1:  # include-file echo (locals.inc)
            continue
        rest = m.group(3).strip()
        toks, remain = [], rest
        while remain:
            t = TOKEN_RE.match(remain)
            if not t:
                break
            toks.append(t.group(1))
            remain = (t.group(2) or "").lstrip()
        off = int(m.group(1), 16)
        if base and off >= base:
            off -= base
        recs.append((off, toks, remain))
    return recs


def build_groups(recs):
    """Group main-file records: a record with text starts a source line;
    following text-less records with bytes continue it. A group is
    (offset, tokens, echo_text); token-less record sequences (labels,
    equates, comments) produce byte-less groups and are dropped."""
    groups = []
    cur = None
    for off, toks, text in recs:
        if text.strip():
            if cur:
                groups.append(cur)
            cur = [off, list(toks), text.strip()]
        elif toks:
            if cur is None:
                sys.exit(f"listing: byte batch @{off:04X} precedes any echo text")
            if off != cur[0] + len(cur[1]):
                sys.exit(f"listing: non-contiguous batch @{off:04X} after group @{cur[0]:04X}")
            cur[1].extend(toks)
        # text-less, token-less record: padding, ignore
    if cur:
        groups.append(cur)
    return [g for g in groups if g[1]]  # only groups that emit bytes


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("project", help="project directory")
    ap.add_argument("--base", default="0xC000")
    ap.add_argument("--write", action="store_true", help="write changes (default: dry run)")
    args = ap.parse_args()
    base = int(args.base, 0)
    proj = Path(args.project)
    prg_s = proj / "prg.s"
    prg = (proj / "prg.bin").read_bytes()
    src = prg_s.read_text(errors="replace").split("\n")

    with tempfile.TemporaryDirectory(dir="/tmp/opencode") as td:
        lst = Path(td) / "out.lst"
        r = subprocess.run(
            ["ca65", "-g", "-l", str(lst), "-o", str(Path(td) / "out.o"), "prg.s"],
            cwd=proj, capture_output=True, text=True,
        )
        if r.returncode:
            sys.exit(f"ca65 failed:\n{r.stderr}")
        recs = parse_listing(lst, base)

    groups = build_groups(recs)

    # Ground-truth check: every group's bytes must equal prg.bin at its offset.
    bad = 0
    for off, toks, _ in groups:
        b = bytes(int(t, 16) for t in toks if t != "rr")
        if b and b != prg[off:off + len(b)]:
            bad += 1
            print(f"GROUP MISMATCH @{off:04X}: {b.hex(' ').upper()} vs "
                  f"{prg[off:off+len(b)].hex(' ').upper()}")
    if bad:
        sys.exit(f"{bad} groups disagree with prg.bin; aborting")

    # Monotonic two-pointer alignment of echoed text to source lines.
    mapping = {}
    unmatched = []
    si = 0
    for off, toks, text in groups:
        j = si
        while j < len(src) and src[j].strip() != text:
            j += 1
        if j == len(src):
            unmatched.append(f"group @{off:04X} ({text[:50]!r}) matches no source line")
            continue
        mapping[j] = (off, len(toks))
        si = j + 1

    out = list(src)
    n_changed = n_shifted = 0
    problems = []
    for i, line in enumerate(src):
        code_part = line.split(";", 1)[0]
        if not EMIT_RE.match(code_part):
            continue
        if i not in mapping:
            if "??" in line:
                problems.append(f"line {i+1}: ?? but no listing group")
            continue
        off, n = mapping[i]
        canonical = CANONICAL.format(
            cpu=base + off, prg=off,
            bytes=" ".join(f"{b:02X}" for b in prg[off:off + n]),
        )
        m = re.search(r";\s*CPU:\$([0-9A-Fa-f]{4})", line)
        was_shifted = m is not None and int(m.group(1), 16) != base + off

        sep = line.find(";")
        tail_segs = line[sep:].split(";")[1:] if sep >= 0 else []
        kept = [s.strip() for s in tail_segs if s.strip() and not AUDIT_SEG_RE.match(s.strip())]
        base_line = code_part.rstrip()
        new_line = (f"{base_line}  ; {'; '.join(kept)}  {canonical}" if kept
                    else f"{base_line}  {canonical}")
        if new_line != line:
            out[i] = new_line
            n_changed += 1
            if was_shifted:
                n_shifted += 1

    residue = [i + 1 for i, l in enumerate(out) if "??" in l]
    print(f"groups: {len(groups)}, mapped: {len(mapping)}, unmatched: {len(unmatched)}")
    for u in unmatched[:10]:
        print("  UNMATCHED", u)
    for p in problems[:10]:
        print("  PROBLEM", p)
    for r in residue[:10]:
        print(f"  RESIDUE line {r}: {out[r-1].rstrip()}")
    print(f"lines rewritten: {n_changed} (shifted addresses fixed: {n_shifted})")
    print(f"??-lines: before {sum(1 for l in src if '??' in l)}, after {len(residue)}")

    if unmatched or residue or problems:
        sys.exit("refusing: unmapped groups or residual ?? lines (orphan mapping)")

    # Validate every canonical comment in the output against prg.bin.
    bad = 0
    for i, line in enumerate(out):
        m = re.search(r"; CPU:\$[0-9A-Fa-f]{4} PRG:([0-9A-Fa-f]{4})\s+((?:[0-9A-Fa-f]{2}(?: |$))+)$", line)
        if not m:
            continue
        off = int(m.group(1), 16)
        vals = bytes.fromhex(m.group(2).replace(" ", "").strip())
        if prg[off:off + len(vals)] != vals:
            bad += 1
            if bad <= 5:
                print(f"BYTE MISMATCH line {i+1} PRG {off:04X}: "
                      f"{m.group(2)} vs {prg[off:off+len(vals)].hex(' ').upper()}")
    if bad:
        sys.exit(f"{bad} output comments disagree with prg.bin; refusing to write")

    if not args.write:
        print("dry run: no files written")
        return
    (prg_s.parent / "prg.s.auditbak").write_text("\n".join(src), errors="replace")
    prg_s.write_text("\n".join(out), errors="replace")
    print(f"wrote {prg_s} (backup at prg.s.auditbak)")


if __name__ == "__main__":
    main()