#!/usr/bin/env python3
"""Build a project's routinemap.tsv from the labels its source already has.

The map is the address-to-name table Pass 2 applies.  It is built from the
source rather than from memory, and two things matter:

  * a label's ADDRESS is the address of the instruction that follows it,
    taken from the file.  A label's *name* is not reliable - some of the
    names minted by earlier passes carry a mistyped address suffix - so the
    name is never parsed for an address.

  * a label keeps the name it has, unless it is one of the address-derived
    `Label_XXXX` forms, which take the repository's `SubC<addr>` name.

A data label is not paired with anything: a `.byte` row clears the pending
label, so a data block never claims the address of a later instruction.

Usage:
    python3 _tools/build_routinemap.py <project dir>
"""

import re
import sys
from pathlib import Path

CPU = re.compile(
    r"^    [A-Za-z@][A-Za-z0-9_]*\s+.*?; CPU:\$([0-9A-F]{4}) PRG:")
LABEL = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*):\s*$")
DATA = (".byte", ".word", ".dword", ".incbin", ".res", ".text", ".asciiz",
        ".segment", ".org", ".include", ".setcpu", ".feature")


def main():
    d = Path(sys.argv[1])
    prefix = ""
    if len(sys.argv) > 2 and sys.argv[2] != "-":
        # the labels of this file live in a bank that shares its CPU window
        # with another; every key and every name carries the bank prefix.
        prefix = sys.argv[2]
    src = (d / "prg.s").read_text()
    lines = src.split("\n")
    # a banked project has one source file per bank; the labels span them all
    for f in ("prg.6.s", "prg.07.s", "prg.0.s", "prg.1.s"):
        p = d / f
        if p.is_file():
            lines.extend(p.read_text().split("\n"))

    addr_of = {}
    cur = None
    first = True
    for line in lines:
        s = line.strip()
        m = LABEL.match(s)
        if m:
            cur = m.group(1)
            first = True
            continue
        if s.startswith(DATA):
            cur = None
            continue
        a = CPU.match(line)
        if a and cur:
            if first:
                addr_of.setdefault(int(a.group(1), 16), cur)
                first = False

    items = {}
    for addr, name in addr_of.items():
        m = re.fullmatch(r"((?:[A-Za-z][A-Za-z0-9]*_)?)Label_([0-9A-F]{4})", name)
        if m:
            # the bank prefix, when the label carries one, stays
            base = f"{m.group(1)}SubC{m.group(2)}"
            items[addr] = base
        else:
            items[addr] = name

    # A bank prefix in the label name is the key's prefix too, so the map
    # can tell two banks' labels at the same CPU address apart.
    keyed = {}
    for addr, name in items.items():
        m = re.match(r"^(Bank[0-9A-F]+)_(?:SubC|Sub|Label)_?", name)
        keyed[(m.group(1) if m else prefix, addr)] = name
    items = keyed

    m = re.search(r"; Vectors: NMI=\$([0-9A-F]{4})\s+RESET=\$([0-9A-F]{4})"
                  r"\s+IRQ=\$([0-9A-F]{4})", src)
    if m:
        for addr, name in ((int(m.group(1), 16), "NmiHandler"),
                           (int(m.group(2), 16), "Reset"),
                           (int(m.group(3), 16), "IrqHandler")):
            if addr in items and items[addr] not in ("NmiHandler", "Reset",
                                                     "IrqHandler"):
                items[addr] = name

    out = sorted(items.items(), key=lambda kv: (kv[0][0], kv[0][1]))
    text = "".join(f"{p + ':' if p else ''}{a:04X}\t{n}\n" for (p, a), n in out)
    (d / "routinemap.tsv").write_text(text)
    print(f"{d}/routinemap.tsv: {len(out)} entries")


if __name__ == "__main__":
    main()