#!/usr/bin/env python3
"""Derive a project's power-on flow analysis from its disassembly.

`PROMPTS.md` asks every project README to carry a power-on flow analysis: the
RESET/NMI/main-loop walk.  This reads that walk out of the disassembly rather
than out of memory:

  * the RESET entry, from the vector table,
  * the routines the reset path calls, in the order it calls them, with the
    addresses,
  * the NMI handler and what it calls,
  * the main loop - the routine the reset path ends up spinning in, found by
    following the jumps at the end of the chain,
  * the IRQ handler, and whether the game uses it.

Everything it prints comes from the source: the call graph is walked from the
reset entry, and the names are the labels in the file.

Usage:
    python3 _tools/derive_poweron.py <project dir> [--depth 4]
"""

import argparse
import re
import sys
from collections import OrderedDict
from pathlib import Path

AUDIT = re.compile(
    r"^    ([A-Za-z@][A-Za-z0-9_]*)(?:\s+(\S+))?\s*; CPU:\$([0-9A-F]{4}) ")
LABEL = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*):\s*$")


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("project")
    ap.add_argument("--depth", type=int, default=4)
    args = ap.parse_args()

    d = Path(args.project)
    lines = (d / "prg.s").read_text().split("\n")

    # every label and the address of the instruction that follows it
    name_at = {}
    addr_of_name = {}
    cur = None
    first = True
    for line in lines:
        m = LABEL.match(line.strip())
        if m:
            cur = m.group(1)
            first = True
            continue
        a = AUDIT.match(line)
        if a and cur:
            addr = int(a.group(3), 16)
            if first:
                name_at.setdefault(addr, cur)
                addr_of_name.setdefault(cur, addr)
                first = False

    # the calls each routine makes, and the routine it falls through to
    calls = {}
    falls = {}
    cur = None
    cur_addr = None
    last_mn = None
    for line in lines:
        m = LABEL.match(line.strip())
        if m:
            # the routine before this one falls through into it, unless it
            # ended with a jump or a return
            if cur is not None and last_mn not in ("rts", "rti", "jmp", "brk"):
                falls[cur] = m.group(1)
            cur = m.group(1)
            calls.setdefault(cur, [])
            cur_addr = None
            last_mn = None
            continue
        a = AUDIT.match(line)
        if not a:
            continue
        mn, op, addr = a.group(1), a.group(2), int(a.group(3), 16)
        last_mn = mn
        if cur is None:
            continue
        if mn in ("jsr", "jmp"):
            target = None
            if op and re.fullmatch(r"\$([0-9A-F]{4})", op):
                target = int(op[1:], 16)
            elif op and op in addr_of_name:
                target = addr_of_name[op]
            if target is not None:
                calls[cur].append((target, name_at.get(target, f"${target:04X}")))

    # the vectors
    vec = {}
    for line in lines:
        m = re.search(r"; Vectors: NMI=\$([0-9A-F]{4})\s+RESET=\$([0-9A-F]{4})"
                      r"\s+IRQ=\$([0-9A-F]{4})", line)
        if m:
            vec = {"nmi": int(m.group(1), 16), "reset": int(m.group(2), 16),
                   "irq": int(m.group(3), 16)}
            break

    def nm(a):
        return name_at.get(a, f"${a:04X}")

    out = []
    out.append(f"## Power-on flow analysis\n")
    out.append("Derived from the disassembly by `_tools/derive_poweron.py`; the "
               "call graph is walked from the reset entry.\n")

    if "reset" in vec:
        r = vec["reset"]
        out.append(f"**RESET** is at `${r:04X}` (`{nm(r)}`).  The CPU comes "
                   "here out of reset with the PPU and APU in their power-on "
                   "state and the zero page and RAM uninitialised.\n")
        out.append("The reset path calls:\n")
        out.append("| Address | Routine |")
        out.append("|---------|---------|")
        seen = set()
        frontier = [(r, 0)]
        while frontier:
            addr, depth = frontier.pop(0)
            if depth > args.depth or addr in seen:
                continue
            seen.add(addr)
            for tgt, tname in calls.get(nm(addr), []):
                out.append(f"| `${tgt:04X}` | `{tname}` |")
                if depth + 1 <= args.depth:
                    frontier.append((tgt, depth + 1))
            # the routine it falls through into
            f = falls.get(nm(addr))
            if f and f in addr_of_name and depth + 1 <= args.depth:
                tgt = addr_of_name[f]
                out.append(f"| `${tgt:04X}` | `{nm(tgt)}` (falls through) |")
                frontier.append((tgt, depth + 1))
        out.append("")

    if "nmi" in vec:
        n = vec["nmi"]
        out.append(f"**NMI** is at `${n:04X}` (`{nm(n)}`).  It is the per-frame "
                   "handler; what it calls:\n")
        out.append("| Address | Routine |")
        out.append("|---------|---------|")
        for tgt, tname in calls.get(nm(n), []):
            out.append(f"| `${tgt:04X}` | `{tname}` |")
        out.append("")

    if "irq" in vec:
        i = vec["irq"]
        same = ("the same address as RESET, so the IRQ vector reaches the reset "
                "path" if i == vec.get("reset") else
                "a different address from RESET")
        out.append(f"**IRQ** is at `${i:04X}` (`{nm(i)}`) - {same}.\n")

    print("\n".join(out))


if __name__ == "__main__":
    main()