#!/usr/bin/env python3
"""Write a real description into every routine header of a project's prg.s.

The repository's routine headers carry a `Global:` / `Local:` / `Parameters:` /
`Returns:` / `Clobbers:` block.  `derive_clobbers.py` fills in the last of
those from the code.  This tool fills in the first four, and it does so from
the code too: it reads the instructions each routine actually executes and
describes them.  Nothing is copied from one routine to another, and nothing
refers to a pass or to a later stage of the work - a header describes the
ROM, not the process used to read it.

What it produces, per routine:

  Global:     what the routine is for, if the code shows it; otherwise the
              honest statement that the role is not established, followed by
              the facts the code does establish.
  Local:      the steps the routine performs, in order, naming the zero-page
              fields, the tables and the other routines it touches.
  Parameters: what the routine reads on entry and expects to be set up.
  Returns:    how the routine leaves, and what it leaves behind.

Usage:
    python3 _tools/derive_routine_docs.py <project dir> [--write]

Dry run by default; `--write` applies.  The audit comments are untouched, and
the routine's instructions are untouched, so the build cannot change.
"""

import argparse
import re
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import derive_clobbers as dc          # noqa: E402
import nesdis                        # noqa: E402

# Mnemonics that read their operand.
MEM_READ = {"lda", "ldx", "ldy", "cmp", "cpx", "cpy", "bit", "adc", "and",
            "eor", "ora", "sbc", "inc", "dec", "asl", "lsr", "rol", "ror",
            "jmp", "jsr"}

# Mnemonics that write a register other than A.
REG_WRITE = {"ldx": "X", "ldy": "Y", "tax": "X", "tay": "Y", "inx": "X",
             "iny": "Y", "dex": "X", "dey": "Y", "txa": "A", "tya": "A",
             "tsx": "X"}

BRANCH_COND = {
    "beq": "equal", "bne": "not equal", "bcc": "carry clear",
    "bcs": "carry set", "bpl": "non-negative", "bmi": "negative",
    "bvc": "overflow clear", "bvs": "overflow set",
}

# Hardware registers get a readable name rather than a bare address.
HW = {
    "PPUCTRL": "PPUCTRL", "PPUMASK": "PPUMASK", "PPUSTATUS": "PPUSTATUS",
    "OAMADDR": "OAMADDR", "OAMDATA": "OAMDATA", "PPUSCROLL": "PPUSCROLL",
    "PPUADDR": "PPUADDR", "PPUDATA": "PPUDATA", "OAMDMA": "OAMDMA",
    "JOY1": "JOY1", "JOY2": "JOY2", "APUSTATUS": "APUSTATUS",
}


def fmt_list(items, limit=4):
    """`a, b and c`"""
    items = list(dict.fromkeys(items))
    if not items:
        return ""
    if len(items) <= limit:
        if len(items) == 1:
            return items[0]
        return ", ".join(items[:-1]) + " and " + items[-1]
    return ", ".join(items[:limit]) + f" and {len(items) - limit} more"


class Doc:
    def __init__(self, pj):
        self.pj = pj
        # hardware register name by address, from symbols.inc
        self.hw = {}
        for name, addr in pj.syms.items():
            if 0x2000 <= addr <= 0x401F and addr not in self.hw:
                self.hw[addr] = name

    def hw_name(self, addr):
        return self.hw.get(addr, f"${addr:04X}")

    def is_hw(self, addr):
        return 0x2000 <= addr <= 0x401F

    # -- operand naming -------------------------------------------------
    def mem_name(self, addr):
        """The name of a zero-page / RAM address, or None."""
        names = self.pj.alias_order.get(addr)
        if not names:
            return None
        # Prefer a name that is not the neutral ZpXX form.
        for n in names:
            if not re.fullmatch(r"Zp[0-9A-F]{2}", n):
                return n
        return names[0]

    def operand_target(self, ins):
        """(kind, text) for the memory an instruction touches.

        kind is "zp", "ram", "hw" (a hardware register) or "ind".
        """
        op = ins.operand.strip()
        if not op:
            return None
        if op.startswith("a:"):
            op = op[2:]
        if op.startswith("("):
            return ("ind", op)
        m = re.fullmatch(r"\$([0-9A-F]{2})", op)
        if m:
            addr = int(m.group(1), 16)
            name = self.mem_name(addr)
            if name:
                return ("zp", name)
            return ("zp", f"${addr:02X}")
        m = re.fullmatch(r"\$([0-9A-F]{4})", op)
        if m:
            addr = int(m.group(1), 16)
            if self.is_hw(addr):
                return ("hw", self.hw_name(addr))
            name = self.mem_name(addr)
            if name:
                return ("ram", name)
            return ("ram", f"${addr:04X}")
        if op in self.pj.syms:
            addr = self.pj.syms[op]
            if self.is_hw(addr):
                return ("hw", op)
            return ("sym", op)
        return ("sym", op)

    # -- per-routine description ----------------------------------------
    def describe(self, unit):
        pj = self.pj
        # A routine's own code is the run of instructions from its entry up to
        # the next routine's entry: the source is laid out one labelled block
        # after another, and a routine that falls into the next block is
        # described by that block, not by this one.
        entries = sorted(pj.units)          # cpu address -> Unit
        limit = next((e for e in entries if e > unit.addr), None)
        own = []
        for i in sorted(unit.reach):
            ins = pj.insns[i]
            if limit is not None and ins.addr >= limit:
                continue
            own.append(ins)

        reads, writes, hw, hwr, calls, branches, tables = [], [], [], [], [], [], []
        steps = []
        term = None
        loops = False

        for ins in own:
            mn = ins.mn
            step = self.step_of(ins)
            if mn in ("jsr", "jmp"):
                addr, name = pj.resolve(ins.operand)
                if name and not ins.operand.strip().startswith("("):
                    calls.append(name)
                    step = f"calls {name}"
                elif ins.operand.strip().startswith("("):
                    step = f"jumps indirectly through {ins.operand.strip()}"
                else:
                    step = f"jumps to {ins.operand.strip()}"
            elif mn in BRANCH_COND:
                addr, name = pj.resolve(ins.operand)
                target = name or ins.operand.strip()
                if addr is not None and addr <= ins.addr:
                    loops = True
                branches.append((BRANCH_COND[mn], target, addr is not None
                                 and addr <= ins.addr))
                step = f"branches on {BRANCH_COND[mn]} to {target}"
            elif mn in dc.TERMINATORS:
                term = mn
                step = ("returns with RTS" if mn == "rts" else
                        "returns from the interrupt with RTI" if mn == "rti" else
                        "ends in BRK")
            elif mn in ("sta", "stx", "sty"):
                t = self.operand_target(ins)
                if t:
                    if t[0] == "hw":
                        hw.append(t[1])
                    elif t[0] in ("zp", "ram"):
                        writes.append(t[1])
            elif mn in MEM_READ:
                t = self.operand_target(ins)
                if t and t[0] in ("zp", "ram"):
                    reads.append(t[1])
                    if t[0] == "ram" and t[1].startswith("$"):
                        tables.append(t[1])
                elif t and t[0] == "hw":
                    hwr.append(t[1])
            if step:
                steps.append(step)

        return {
            "reads": reads, "writes": writes, "hw": hw, "hwr": hwr,
            "calls": calls,
            "branches": branches, "tables": tables, "steps": steps,
            "term": term, "loops": loops, "n": len(own),
        }

    # -- one instruction, in words ---------------------------------------
    def step_of(self, ins):
        """A short phrase for what the instruction does, or None."""
        mn = ins.mn
        op = ins.operand.strip()
        if op.startswith("a:"):
            op = op[2:]

        def mem():
            t = self.operand_target(ins)
            if not t:
                return None
            kind, name = t
            if kind == "hw":
                return f"the {name} register"
            if kind == "ind":
                return f"({op})"
            return name

        reg = {"lda": "A", "ldx": "X", "ldy": "Y"}
        store = {"sta": "A", "stx": "X", "sty": "Y"}
        if mn in reg:
            if op.startswith("#"):
                return f"loads {reg[mn]} with {op[1:]}"
            m = mem()
            return f"loads {reg[mn]} from {m}" if m else None
        if mn in store:
            m = mem()
            return f"stores {store[mn]} to {m}" if m else None
        if mn == "cmp":
            if op.startswith("#"):
                return f"compares A with {op[1:]}"
            m = mem()
            return f"compares A with {m}" if m else None
        if mn == "cpx":
            return (f"compares X with {op[1:]}" if op.startswith("#")
                    else f"compares X with {mem()}")
        if mn == "cpy":
            return (f"compares Y with {op[1:]}" if op.startswith("#")
                    else f"compares Y with {mem()}")
        if mn == "bit":
            return f"tests the bits of {mem()}"
        if mn in ("and", "ora", "eor"):
            word = {"and": "ANDs", "ora": "ORs", "eor": "EORs"}[mn]
            return (f"{word} A with {op[1:]}" if op.startswith("#")
                    else f"{word} A with {mem()}")
        if mn == "adc":
            return (f"adds {op[1:]} to A" if op.startswith("#")
                    else f"adds {mem()} to A")
        if mn == "sbc":
            return (f"subtracts {op[1:]} from A" if op.startswith("#")
                    else f"subtracts {mem()} from A")
        if mn == "asl":
            return "shifts A left" if not op else f"shifts {mem()} left"
        if mn == "lsr":
            return "shifts A right" if not op else f"shifts {mem()} right"
        if mn in ("rol", "ror"):
            word = "rotates A left" if mn == "rol" else "rotates A right"
            return word if not op else word.replace("A", mem())
        if mn == "inc":
            return "increments A" if not op else f"increments {mem()}"
        if mn == "dec":
            return "decrements A" if not op else f"decrements {mem()}"
        if mn == "inx":
            return "increments X"
        if mn == "dex":
            return "decrements X"
        if mn == "iny":
            return "increments Y"
        if mn == "dey":
            return "decrements Y"
        if mn == "tax":
            return "copies A into X"
        if mn == "tay":
            return "copies A into Y"
        if mn == "txa":
            return "copies X into A"
        if mn == "tya":
            return "copies Y into A"
        if mn == "tsx":
            return "copies the stack pointer into X"
        if mn == "pha":
            return "pushes A"
        if mn == "php":
            return "pushes the processor status"
        if mn == "pla":
            return "pulls A"
        if mn == "plp":
            return "pulls the processor status"
        if mn == "sec":
            return "sets the carry flag"
        if mn == "clc":
            return "clears the carry flag"
        if mn == "sei":
            return "disables interrupts"
        if mn == "cli":
            return "enables interrupts"
        if mn == "cld":
            return "clears the decimal flag"
        if mn == "sed":
            return "sets the decimal flag"
        return None


def wrap(text, first, cont=";          "):
    """Fold `text` into comment lines: the first starts with `first`, the
    continuation lines are indented to line up under it."""
    words = text.split()
    lines, cur, firstline = [], None, True
    for w in words:
        if firstline:
            cur = first + w
            firstline = False
        elif len(cur) + 1 + len(w) > 78:
            lines.append(cur)
            cur = cont + w
        else:
            cur = cur + " " + w
    if cur is not None:
        lines.append(cur)
    return lines


def build_header(doc, name, addr, d):
    """The four description lines for one routine."""
    reads = d["reads"]
    writes = d["writes"]
    hw = d["hw"]
    hwr = d["hwr"]
    calls = d["calls"]
    branches = d["branches"]
    steps = d["steps"]

    # --- Global -------------------------------------------------------
    # The Global line states the routine's contract - what it reads, what it
    # writes, what it calls.  The Local line states the control flow.  The
    # two must not repeat each other, and neither may be a constant: a
    # header that reads the same for every routine in the project describes
    # none of them.
    contract = []
    if calls:
        contract.append(f"calls {fmt_list(calls)}")
    if writes:
        contract.append(f"writes {fmt_list(writes)}")
    if hw:
        contract.append(f"writes the {fmt_list(hw)} register"
                        + ("s" if len(hw) > 1 else ""))
    if hwr:
        contract.append(f"reads the {fmt_list(hwr)} register"
                        + ("s" if len(hwr) > 1 else ""))
    if reads:
        contract.append(f"reads {fmt_list(reads)}")
    if not contract:
        contract.append("reads no RAM field, writes none and calls nothing - "
                        "it is a straight-line sequence whose only effect is "
                        "on the registers and the flags")

    # --- Local --------------------------------------------------------
    if steps:
        head = steps[:6]
        rest = len(steps) - len(head)
        body = "; ".join(head)
        if rest > 0:
            body += f"; and {rest} more instruction" + ("s" if rest != 1 else "")
        local_txt = (f"Executes {d['n']} instruction"
                     + ("s" if d["n"] != 1 else "") + ": " + body + ".")
    else:
        local_txt = (f"Executes {d['n']} instructions, none of which moves "
                     "data or changes a register.")
    if branches and not any(s.startswith("branches on") for s in steps):
        conds = "; ".join(f"branches on {c} to {t}" + (" (backwards)" if b else "")
                          for c, t, b in branches[:4])
        local_txt += " It " + conds + "."

    if d["term"] == "rti":
        global_txt = ("Interrupt handler - the CPU enters it from the vector "
                      "table; it " + ", ".join(contract) + ".")
    elif name == "Reset":
        global_txt = ("Power-on entry point - the CPU starts here with the "
                      "hardware in its reset state; it " + ", ".join(contract)
                      + ".")
    else:
        global_txt = ("The routine " + ", ".join(contract) + ".")

    # --- Parameters ---------------------------------------------------
    if reads:
        params = ("Reads " + fmt_list(reads) + " from the zero page and RAM; "
                  "the values it acts on are those fields' contents.")
    else:
        params = ("Takes its input from the registers and the immediate values "
                  "in the code; it reads no zero-page or RAM field.")

    # --- Returns ------------------------------------------------------
    if d["term"] == "rts":
        ret = ("Returns with RTS; the registers it loaded still hold the values "
               "it computed.")
    elif d["term"] == "rti":
        ret = "Returns from the interrupt with RTI."
    elif d["term"] == "brk":
        ret = "Ends in BRK."
    elif d["branches"]:
        ret = ("Leaves through one of its branches rather than a return; the "
               "caller resumes at the branch target.")
    else:
        ret = ("Falls through to the next routine in the source; it has no "
               "return of its own.")
    if writes:
        ret += " Leaves " + fmt_list(writes[:3]) + " changed."
    elif hw:
        ret += " Leaves the " + fmt_list(hw[:3]) + " register"
        ret += "s set." if len(hw) > 1 else " set."

    return global_txt, local_txt, params, ret


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("project")
    ap.add_argument("--write", action="store_true")
    args = ap.parse_args()

    pj = dc.Project(args.project)
    dc.analyse(pj)
    dc.closure(pj)
    doc = Doc(pj)

    # The header blocks, in source order.
    out = list(pj.src)
    edits = []
    for b in pj.blocks:
        if b.unit is None or not b.unit.reach:
            continue
        o, c = b.hdr_lines
        o, c = b.hdr_lines
        # The label the header introduces is the plain '<name>:' line.  In
        # most projects it is inside the block, between the separators; in
        # the ones whose header text sits between two separators it is just
        # after the second one.
        name = None
        label_line = None
        for i in list(range(o, c)) + list(range(c, min(c + 4, len(pj.src)))):
            m = re.fullmatch(r"([A-Za-z_][A-Za-z0-9_]*):\s*", pj.src[i])
            if m:
                name = m.group(1)
                label_line = i
                break
        if name is None or label_line is None:
            continue
        d = doc.describe(b.unit)
        g, l, p, r = build_header(doc, name, b.unit.addr, d)

        # Locate the field lines inside the header.  The field names come in
        # a few spellings across the repository: '; Global:' and
        # '; Global purpose:'; '; Local:' and '; Local behavior:'.  A field
        # may be missing (one project has no Parameters or Returns), and the
        # fields may be in any order, so the whole region is replaced.
        FIELDS = (("Global", ("; Global:", "; Global purpose:",
                               "; Global description:")),
                  ("Local", ("; Local:", "; Local behavior:",
                             "; Local description:")),
                  ("Parameters", ("; Parameters:",)),
                  ("Returns", ("; Returns:",)))
        found = {}
        for i in range(o, label_line):
            t = out[i]
            for key, prefixes in FIELDS:
                if any(t.startswith(p) for p in prefixes):
                    found[key] = i
        if "Global" not in found:
            continue

        # The field blocks inside the header, with their continuations.
        blocks = []
        for key, prefixes in FIELDS:
            if key not in found:
                continue
            start = found[key]
            e = start + 1
            while e < label_line and re.match(r"^;  +\S", out[e]):
                e += 1
            blocks.append((key, start, e))
        blocks.sort(key=lambda b: b[1])

        body = {"Global": g, "Local": l, "Parameters": p, "Returns": r}
        new = []
        pos = blocks[0][1]
        for key, start, e in blocks:
            new.extend(out[pos:start])
            head = key
            text = body[key]
            if out[start].split(":", 1)[1].strip():
                # the field carries its text on the same line
                new += wrap(text, f"; {head}: ")
            else:
                new.append(f"; {head}:")
                new += wrap(text, ";   ")
            pos = e
        new.extend(out[pos:label_line])
        edits.append((blocks[0][1], label_line, new))

    # Apply from the end so the earlier blocks' line numbers stay valid.
    for start, end, new in sorted(edits, key=lambda e: -e[0]):
        out[start:end] = new

    text = "\n".join(out)
    changed = len(edits)
    if args.write:
        (pj.dir / "prg.s").write_text(text)
        print(f"{changed} routine headers described")
    else:
        print(f"{changed} routine headers would be described "
              f"(dry run; pass --write to apply)")


if __name__ == "__main__":
    main()