#!/usr/bin/env python3
"""Derive the `; Clobbers:` line of every routine header in a project's prg.s.

The repo's routine-header template (AGENTS.md) is "global + local description,
parameters, returns, clobbers".  A routine's *clobbers* are what the caller
loses across the call: the registers the routine writes, the processor flags
it changes, and the memory it writes - plus the same for the routines it
calls, because a callee's damage survives the return.

Everything here is derived from the source itself, not guessed:

  * the opcode comes from the audit comment's byte list, so the addressing
    mode is exact (`sta $0400,x` is ABSX, `sta Zp00` is ZP) and a data
    directive is never mistaken for an instruction;
  * the operand text is resolved through `locals.inc`, `symbols.inc` and the
    labels defined in `prg.s`, so a clobber can be named;
  * the control flow comes from the instructions themselves.  A *unit* is a
    named entry point in the source (a label that is followed by code); the
    code a unit reaches - the shared tails it branches into, the block it
    falls through into - is its own, while a `jsr`/`jmp` to another unit is a
    call, so the callee's clobbers count but its code is not merged in.

Register *preservation* is modelled too, so a routine that saves a register on
the stack and puts it back before returning (`tya`/`pha` ... `pla`/`tay`) is
not reported as clobbering it.  A small abstract interpreter carries the value
of A, X, Y and of each processor flag through the unit: a register survives
when its value at the end is the value it had at the entry, and a flag
survives when a `php`/`plp` pair puts it back.

A call into code the project has not decoded (a forward reference in
`locals.inc` that lands in a data region) can damage anything, so it is
reported as a clobber of every register and every flag.

Usage:
    python3 _tools/derive_clobbers.py "NROM/Balloon Fight (JU)"            # report
    python3 _tools/derive_clobbers.py "NROM/Balloon Fight (JU)" --write    # apply
"""

import argparse
import os
import re
import sys
from pathlib import Path

sys.path.insert(0, os.path.dirname(os.path.abspath(__file__)))
import nesdis  # noqa: E402

# Registers written by each mnemonic.
REG_WRITE = {
    "lda": "A", "pla": "A", "txa": "A", "tya": "A", "lax": "A",
    "ldx": "X", "inx": "X", "dex": "X", "tax": "X", "txs": "X",
    "ldy": "Y", "iny": "Y", "dey": "Y", "tay": "Y",
}

# Which processor flags each mnemonic sets or clears.
FLAG_BITS = {
    "N": ("lda", "ldx", "ldy", "tax", "tay", "txa", "tya", "tsx", "cmp",
          "cpx", "cpy", "adc", "and", "asl", "bit", "dec", "dex", "dey",
          "eor", "inc", "inx", "iny", "lsr", "ora", "rol", "ror", "sbc"),
    "Z": ("lda", "ldx", "ldy", "tax", "tay", "txa", "tya", "tsx", "cmp",
          "cpx", "cpy", "adc", "and", "asl", "bit", "dec", "dex", "dey",
          "eor", "inc", "inx", "iny", "lsr", "ora", "rol", "ror", "sbc"),
    "C": ("adc", "cmp", "cpx", "cpy", "sbc", "sec", "asl", "lsr", "rol", "ror"),
    "D": ("cld", "sed"),
    "I": ("cli", "sei", "plp", "brk"),
    "V": ("adc", "sbc", "bit", "plp"),
}
FLAGS = list(FLAG_BITS)

# Mnemonics that write their operand.
MEM_WRITE = {"sta", "stx", "sty", "inc", "dec", "asl", "lsr", "rol", "ror"}

# Mnemonics that only move a register around.
COPY = {"txa": ("X", "A"), "tya": ("Y", "A"), "tax": ("A", "X"),
        "tay": ("A", "Y")}

# Instructions that end the routine.
TERMINATORS = {"rts", "rti", "brk"}

AUDIT_RE = re.compile(
    r"^    ([A-Za-z@][A-Za-z0-9_]*)\s*(.*?)\s*; CPU:\$([0-9A-F]{4,6}) "
    r"PRG:([0-9A-F]{4,6})  ((?:[0-9A-F]{2} ?)+)$")
LABEL_RE = re.compile(r"^([A-Za-z_@][A-Za-z0-9_]*):\s*(.*)$")

ENTRY = {r: ("entry", r) for r in ("A", "X", "Y")}

# Memory regions, named rather than listed address by address.
REGIONS = (
    (0x0200, 0x02FF, "the OAM shadow"),
    (0x4000, 0x4017, "the APU registers"),
)


class Insn:
    __slots__ = ("addr", "length", "mn", "operand", "mode", "opcode", "bytes",
                 "succ", "calls", "unknown_call")

    def __init__(self, addr, length, mn, operand, mode, opcode, raw):
        self.addr = addr
        self.length = length
        self.mn = mn
        self.operand = operand
        self.mode = mode
        self.opcode = opcode
        self.bytes = raw
        self.succ = []          # (kind, instruction index); kind "fall"/"branch"
        self.calls = []         # instruction index of a called unit's entry
        self.unknown_call = False


class Unit:
    """A named entry point in the source and the code it reaches."""

    def __init__(self, name, entry, addr):
        self.name = name
        self.entry = entry
        self.addr = addr
        self.reach = set()
        self.calls = set()         # units called
        self.unknown_call = False
        self.touched = set()
        self.restored = set()
        self.flag_written = set()
        self.flag_restored = set()
        self.mem = []
        # transitive results
        self.eff_regs = set()
        self.eff_flags = set()
        self.eff_mem = set()
        self.eff_calls = set()


class Block:
    """A routine header in prg.s."""

    def __init__(self, name, hdr_lines, entry_addr):
        self.name = name
        self.hdr_lines = hdr_lines
        self.entry_addr = entry_addr
        self.unit = None


def load_symbols(proj):
    """Return (name -> address, address -> [names in file order])."""
    syms = {}
    order = {}
    for fname in ("locals.inc", "symbols.inc"):
        p = proj / fname
        if not p.is_file():
            continue
        for raw in p.read_text().split("\n"):
            line = raw.split(";", 1)[0].strip()
            if "=" not in line:
                continue
            name, _, val = line.partition("=")
            name = name.strip()
            val = val.strip()
            if not val.startswith("$"):
                continue
            try:
                addr = int(val[1:], 16)
            except ValueError:
                continue
            syms.setdefault(name, addr)
            order.setdefault(addr, []).append(name)
    return syms, order


class Project:
    def __init__(self, proj):
        self.dir = Path(proj)
        self.src = (self.dir / "prg.s").read_text().split("\n")
        self.syms, self.alias_order = load_symbols(self.dir)
        self.insns = []
        self.by_addr = {}
        self.labels = {}
        self.units = {}          # cpu address -> Unit
        self.blocks = []
        self._parse()

    def _parse(self):
        for line in self.src:
            m = AUDIT_RE.match(line)
            if not m:
                continue
            mn = m.group(1).lower()
            operand = m.group(2).strip()
            by = m.group(5).split()
            opcode = int(by[0], 16)
            entry = nesdis.OPCODES.get(opcode) or nesdis.UNDOC_OPCODES.get(opcode)
            mode = entry[1] if entry else nesdis.IMP
            addr = int(m.group(3), 16)
            ins = Insn(addr, len(by), mn, operand, mode, opcode, by)
            self.by_addr[addr] = len(self.insns)
            self.insns.append(ins)
        # labels: the address of the next instruction after the label
        for i, line in enumerate(self.src):
            m = LABEL_RE.match(line)
            if not m:
                continue
            for j in range(i, min(i + 4, len(self.src))):
                a = AUDIT_RE.match(self.src[j])
                if a:
                    self.labels.setdefault(m.group(1), int(a.group(3), 16))
                    break
        # a unit for every label that is followed by code
        for name, addr in self.labels.items():
            if addr in self.by_addr and addr not in self.units:
                self.units[addr] = Unit(name, self.by_addr[addr], addr)
        # header blocks
        # The header separator comes in several spellings across this
        # repository: ';=====', '; ------', ';---', '; ==== ' and so on.
        # A separator is a line whose characters after the ';' are all
        # dashes, equals signs or spaces, and which is long enough to be
        # one.
        sep = [i for i, l in enumerate(self.src)
               if l.startswith(";") and len(l.strip()) > 3
               and set(l.strip()[1:]) <= set("-= ")]
        GLOBAL = ("; Global:", "; Global purpose:", "; Global description:")
        if not sep:
            # One project (Bomberman (U)) writes its headers with no
            # separator line at all: the block runs from the '; Routine:'
            # line to the next one.
            starts = [i for i, l in enumerate(self.src)
                      if l.startswith("; Routine: ")]
            sep = []
            for k, s in enumerate(starts):
                sep.append(s)
                if k + 1 < len(starts):
                    sep.append(starts[k + 1])
            # the pairs must be (header start, header end); with no
            # separator the header end is the next '; Routine:' line
            sep = [s for s in starts]
            for k in range(len(sep) - 1):
                o, c = sep[k], sep[k + 1]
                body = self.src[o + 1:c]
                if not any(l.startswith(GLOBAL) for l in body):
                    continue
                name = None
                for h in body:
                    if h.startswith("; ") and not h.startswith(";  "):
                        t = h[2:].strip()
                        if t and not t.endswith(":") and ":" not in t:
                            name = t
                            break
                entry = self._entry_after(c)
                blk = Block(name, (o, c), entry)
                if entry is not None and entry not in self.units:
                    self.units[entry] = Unit(name, self.by_addr[entry], entry)
                blk.unit = self.units.get(entry)
                self.blocks.append(blk)
            sep = []
        # The field names come in a few spellings across the repository:
        # '; Global:' and '; Global purpose:'; '; Local:' and
        # '; Local behavior:'.  Both are recognised.
        GLOBAL = ("; Global:", "; Global purpose:", "; Global description:")
        # The field names come in a few spellings across the repository:
        # '; Global:' and '; Global purpose:'; '; Local:' and
        # '; Local behavior:'.  Both are recognised.
        GLOBAL = ("; Global:", "; Global purpose:", "; Global description:")
        for k in range(len(sep) - 1):
            o, c = sep[k], sep[k + 1]
            body = self.src[o + 1:c]
            if not any(l.startswith(GLOBAL) for l in body):
                continue
            name = None
            for h in body:
                if h.startswith("; ") and not h.startswith(";  "):
                    t = h[2:].strip()
                    if t and not t.endswith(":") and ":" not in t:
                        name = t
                        break
            entry = self._entry_after(c)
            blk = Block(name, (o, c), entry)
            if entry is not None and entry not in self.units:
                # `entry` is a CPU address; the unit's entry is the index of
                # the instruction at that address.
                self.units[entry] = Unit(name, self.by_addr[entry], entry)
            blk.unit = self.units.get(entry)
            self.blocks.append(blk)

    def _entry_after(self, line):
        for j in range(line, len(self.src)):
            m = AUDIT_RE.match(self.src[j])
            if m:
                addr = int(m.group(3), 16)
                return addr if addr in self.by_addr else None
        return None

    def resolve(self, operand):
        """(cpu address, name) of a branch/call operand, or (None, None)."""
        op = operand.strip()
        if op.startswith("("):
            return None, None
        if op.startswith("$"):
            try:
                return int(op[1:], 16), op
            except ValueError:
                return None, None
        if op in self.labels:
            return self.labels[op], op
        if op in self.syms:
            return self.syms[op], op
        return None, None


def build_cfg(pj):
    """Wire up the control flow of every instruction."""
    for ins in pj.insns:
        mn = ins.mn
        nxt = pj.by_addr.get(ins.addr + ins.length)
        if mn in TERMINATORS:
            continue
        if mn == "jmp":
            if ins.operand.strip().startswith("("):
                continue                        # indirect dispatch
            addr, _name = pj.resolve(ins.operand)
            tgt = pj.by_addr.get(addr) if addr is not None else None
            if tgt is None:
                ins.unknown_call = True
            elif addr in pj.units:
                ins.calls.append(tgt)            # a tail call to another unit
            else:
                ins.succ.append(("branch", tgt))  # a jump into a shared tail
            continue
        if mn == "jsr":
            addr, _name = pj.resolve(ins.operand)
            tgt = pj.by_addr.get(addr) if addr is not None else None
            if tgt is None:
                ins.unknown_call = True
            elif addr in pj.units:
                ins.calls.append(tgt)            # a call to another unit
            else:
                ins.succ.append(("branch", tgt))  # a call into a shared tail
            if nxt is not None:
                ins.succ.append(("fall", nxt))
            continue
        if ins.mode == nesdis.REL:
            if len(ins.bytes) < 2:
                # A branch whose audit comment does not carry the operand
                # byte: the target cannot be computed, so the branch is
                # treated as falling through.
                if nxt is not None:
                    ins.succ.append(("fall", nxt))
                continue
            rel = int(ins.bytes[1], 16)
            if rel > 0x7F:
                rel -= 0x100
            tgt = pj.by_addr.get(ins.addr + 2 + rel)
            if tgt is None:
                ins.unknown_call = True
            else:
                ins.succ.append(("branch", tgt))
            if nxt is not None:
                ins.succ.append(("fall", nxt))
            continue
        if nxt is not None:
            ins.succ.append(("fall", nxt))


def analyse(pj):
    """Fill in the per-unit clobber data."""
    build_cfg(pj)
    for unit in pj.units.values():
        seen, stack = set(), [unit.entry]
        while stack:
            i = stack.pop()
            if i in seen:
                continue
            seen.add(i)
            ins = pj.insns[i]
            for c in ins.calls:
                called = pj.units.get(pj.insns[c].addr)
                if called is not None and called is not unit:
                    unit.calls.add(called)
            if ins.unknown_call:
                unit.unknown_call = True
            for _kind, s in ins.succ:
                if s not in seen:
                    stack.append(s)
        unit.reach = seen
        unit.touched = set()
        unit.flag_written = set()
        unit.mem = []
        for i in sorted(seen):
            ins = pj.insns[i]
            mn = ins.mn
            if mn in REG_WRITE:
                unit.touched.add(REG_WRITE[mn])
            for bit, mns in FLAG_BITS.items():
                if mn in mns:
                    unit.flag_written.add(bit)
            if mn in MEM_WRITE:
                mem = mem_target(pj, ins)
                if mem:
                    unit.mem.append(mem)
        # register / flag preservation, over the reached code in address order
        val = {"A": ENTRY["A"], "X": ENTRY["X"], "Y": ENTRY["Y"]}
        flags = {f: ENTRY for f in FLAGS}
        stack = []
        for i in sorted(seen):
            mn = pj.insns[i].mn
            if mn == "pha":
                stack.append(("A", val["A"]))
            elif mn == "php":
                stack.append(("P", dict(flags)))
            elif mn == "pla":
                v = stack.pop() if stack else None
                val["A"] = v[1] if v and v[0] == "A" else ("?",)
            elif mn == "plp":
                v = stack.pop() if stack else None
                flags = (dict(v[1]) if v and v[0] == "P"
                         else {f: ("?",) for f in FLAGS})
            elif mn in COPY:
                s, d = COPY[mn]
                val[d] = val[s]
            else:
                if mn in REG_WRITE:
                    val[REG_WRITE[mn]] = ("?",)
                for bit, mns in FLAG_BITS.items():
                    if mn in mns:
                        flags[bit] = ("?",)
                if mn in ("jsr", "jmp"):
                    val = {"A": ("?",), "X": ("?",), "Y": ("?",)}
                    flags = {f: ("?",) for f in FLAGS}
        unit.restored = {k for k, v in val.items() if v == ENTRY[k]}
        unit.flag_restored = {f for f, v in flags.items() if v == ENTRY}


def mem_target(pj, ins):
    """(name, address) of the location an instruction writes."""
    base = ins.operand.strip()
    if base.startswith("("):
        base = base.strip("()").split(",")[0].strip()
    base = re.sub(r",[xy]$", "", base).strip()
    if base in ("a", "A"):
        return None
    if base.startswith("$"):
        try:
            addr = int(base[1:], 16)
        except ValueError:
            return None
        return (base, addr)
    addr = pj.syms.get(base, pj.labels.get(base))
    if addr is None:
        return None
    return (base, addr)


def closure(pj):
    """Fold each unit's callees' clobbers into the caller's."""
    state = {}

    def reach(u, seen):
        if id(u) in state:
            return state[id(u)]
        if id(u) in seen:
            return (set(), set(), set(), set())
        regs = set(u.touched) - set(u.restored)
        flags = set(u.flag_written) - set(u.flag_restored)
        mem = set(u.mem)
        calls = set()
        for other in u.calls:
            cr, cf, cm, cc = reach(other, seen | {id(u)})
            regs |= cr
            flags |= cf
            mem |= cm
            calls |= cc
            calls.add(other.name)
        if u.unknown_call:
            # code this project has not decoded can damage anything
            regs |= {"A", "X", "Y"}
            flags |= set(FLAGS)
        regs -= set(u.restored)
        flags -= set(u.flag_restored)
        state[id(u)] = (regs, flags, mem, calls)
        return state[id(u)]

    for u in pj.units.values():
        (u.eff_regs, u.eff_flags, u.eff_mem,
         u.eff_calls) = reach(u, frozenset())


def mem_name(name, addr):
    if name and not name.startswith("$"):
        return name
    for lo, hi, label in REGIONS:
        if lo <= addr <= hi:
            return label
    return name or f"${addr:04X}"


def fmt_regs(regs):
    order = [x for x in ("A", "X", "Y") if x in regs]
    if not order:
        return "no register"
    if len(order) == 1:
        return order[0]
    return ", ".join(order[:-1]) + " and " + order[-1]


def fmt_flags(flags):
    if not flags:
        return None
    order = [f for f in FLAGS if f in flags]
    if len(order) == len(FLAGS):
        return "all the processor flags"
    if len(order) == 1:
        return f"the {order[0]} flag"
    return "the " + ", ".join(order[:-1]) + " and " + order[-1] + " flags"


def memory_summary(pj, u, limit):
    """A list of the memory the unit writes, or a summary when it is long.

    Several names can share one address (`Temp12` and `RecordIndex` are both
    $12), so the entries are folded per address and the first name in
    `locals.inc` is used - the file lists the primary name first.
    """
    by_addr = {}
    for name, addr in u.eff_mem:
        if addr not in by_addr:
            by_addr[addr] = name
    entries = []
    for addr in sorted(by_addr):
        name = by_addr[addr]
        aliases = pj.alias_order.get(addr, [])
        if name in aliases:
            name = aliases[0]
        n = mem_name(name, addr)
        if n not in entries:
            entries.append(n)
    if not entries:
        return None
    if len(entries) <= limit:
        return ", ".join(entries)
    regions, rest = [], []
    for n in entries:
        (regions if n.startswith("the ") else rest).append(n)
    parts = list(dict.fromkeys(regions))
    if rest:
        parts.append("the RAM fields it writes")
    return " and ".join(parts)


def clobber_text(pj, u, limit=6):
    parts = [fmt_regs(u.eff_regs)]
    fl = fmt_flags(u.eff_flags)
    if fl:
        parts.append(fl)
    mem = memory_summary(pj, u, limit)
    if mem:
        parts.append(mem)
    if len(parts) == 1 and parts[0] == "no register":
        return "nothing"
    return ", ".join(parts)


def wrap(text, width=76, indent=";   "):
    """Wrap a sentence into the ';   ' comment lines of a header block."""
    words = text.split()
    lines, cur = [], ""
    for w in words:
        cand = (cur + " " + w) if cur else w
        if cur and len(indent) + len(cand) + 1 > width:
            lines.append(indent + cur)
            cur = w
        else:
            cur = cand
    if cur:
        lines.append(indent + cur)
    return lines


CLOBBER_HEAD = "; Clobbers:"


def clobber_block(pj, u, limit=6):
    text = clobber_text(pj, u, limit)
    if not text.endswith("."):
        text += "."
    return [CLOBBER_HEAD] + wrap(text)


def is_blank(line):
    """A header line that carries no text: empty, or just the ';' separator."""
    s = line.strip()
    return s == "" or s == ";"


def insert_clobbers(pj, limit=6):
    """Return the source with a Clobbers block in every routine header."""
    plan = []
    for b in pj.blocks:
        o, c = b.hdr_lines
        at = c
        for i in range(o + 1, c):
            if pj.src[i].startswith("; PRG offset"):
                at = i
                break
        plan.append((b, at))
    out, pos = [], 0
    for b, at in plan:
        out.extend(pj.src[pos:at])
        # drop every Clobbers block already in the header (its head, its text
        # and the blank ';' line that follows it) so the new one replaces them.
        # The head has two spellings in the repository: the bare '; Clobbers:'
        # followed by an indented list, and the one-line '; Clobbers: A, X; ...'.
        while True:
            j = len(out)
            while j > 0 and is_blank(out[j - 1]):
                j -= 1
            if j > 0 and out[j - 1].startswith(";   "):
                k = j
                while k > 0 and out[k - 1].startswith(";   "):
                    k -= 1
                if k > 0 and out[k - 1] == CLOBBER_HEAD:
                    out = out[:k - 1]
                    continue
            elif j > 0 and out[j - 1].startswith(CLOBBER_HEAD):
                out = out[:j - 1]
                continue
            break
        if b.unit is not None:
            out.extend(clobber_block(pj, b.unit, limit))
            out.append(";")
        pos = at
    out.extend(pj.src[pos:])
    return out


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("project")
    ap.add_argument("--write", action="store_true")
    ap.add_argument("--limit", type=int, default=6)
    args = ap.parse_args()
    pj = Project(args.project)
    analyse(pj)
    closure(pj)

    if args.write:
        out = insert_clobbers(pj, args.limit)
        (pj.dir / "prg.s").write_text("\n".join(out))
        print(f"{len(pj.blocks)} routine headers given a Clobbers block")
        return

    print(f"{len(pj.blocks)} routine headers, {len(pj.units)} units")
    for b in pj.blocks:
        u = b.unit
        print(f"--- {b.name}  (hdr {b.hdr_lines[0]+1}-{b.hdr_lines[1]+1}, "
              f"{len(u.reach) if u else 0} insns reached)")
        print(f"    clobbers: {clobber_text(pj, u, args.limit) if u else '-'}")
        print(f"    calls   : {', '.join(sorted(x for x in u.eff_calls if x)) or '-'}"
              f"{'  [undecoded callee]' if u and u.unknown_call else ''}")


if __name__ == "__main__":
    main()