#!/usr/bin/env python3
"""Derive the internal branch-target names for a Pass 2 map.

A routine entry gets a semantic name written by hand (see routinemap.tsv).
The internal branch targets inside it are named mechanically from the branch
that reaches them: the role says what has to be true for the branch to be
taken, and the owner is the routine the target belongs to.

  beq -> OnZero      bne -> OnNonzero
  bcc -> OnCarryClear  bcs -> OnCarrySet
  bmi -> OnNegative    bpl -> OnNonnegative
  bvc -> OnOverflowClear  bvs -> OnOverflowSet
  a backwards branch (the target is below the branch) -> Loop

Names are made unique inside their owner by appending 2, 3, ... in address
order, so the map is deterministic.

A banked cartridge (MMC3 and friends) has one copy of every label per bank,
so the map may name a bank's label with a `<prefix>:<addr>` key and
`--bank-prefix` says which prefix belongs to which part of the file.

Usage:
    python3 _tools/derive_branch_names.py <project dir> <routinemap.tsv>
"""

import re
import sys
from collections import defaultdict
from pathlib import Path

ROLE = {
    "beq": "OnZero", "bne": "OnNonzero",
    "bcc": "OnCarryClear", "bcs": "OnCarrySet",
    "bmi": "OnNegative", "bpl": "OnNonnegative",
    "bvc": "OnOverflowClear", "bvs": "OnOverflowSet",
}

LABEL_RE = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*):\s*$")
CPU_RE = re.compile(r"; CPU:\$([0-9A-F]{4}) PRG:([0-9A-F]{4})(?: BANK:([0-9A-F]+))?")
# a name that is only the address: `SubC800F`, `Sub_800F`, `Label_800F`
ADDRESS_DERIVED = re.compile(
    r"^(?:[A-Za-z][A-Za-z0-9]*_)?(?:SubC|Sub|Label)_?[0-9A-F]{4}$")


def address_derived(name):
    """True when the name is only the address it sits at."""
    return bool(ADDRESS_DERIVED.match(name))


def map_key(line):
    """The (prefix, address) key of a map line, or None."""
    parts = line.split("\t")
    if len(parts) < 2:
        return None
    key = parts[0].strip()
    if ":" in key:
        prefix, _, addr = key.partition(":")
    else:
        prefix, addr = "", key
    try:
        return (prefix, int(addr, 16))
    except ValueError:
        return None

BANKED_LABEL_RE = re.compile(
    r"^(?:(?P<pre>[A-Za-z][A-Za-z0-9]*)_)?"
    r"(?:SubC|Sub|Label)_?(?P<addr>[0-9A-F]{4})$")


def read_map(lines):
    """{(prefix, address): name}"""
    names = {}
    for raw in lines:
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        parts = line.split("\t") if "\t" in line else line.split(None, 1)
        key = parts[0].strip()
        prefix = ""
        if ":" in key:
            prefix, _, key = key.partition(":")
            prefix = prefix.strip()
        names[(prefix, int(key.strip().lstrip("$"), 16))] = parts[1].strip()
    return names


def read_bank_prefixes(spec):
    out = []
    for tok in spec.split(","):
        tok = tok.strip()
        if not tok:
            continue
        off, _, pre = tok.partition(":")
        out.append((int(off, 16), pre.strip()))
    return sorted(out)


def prefix_at(bank_prefixes, file_offset):
    pre = ""
    for off, p in bank_prefixes:
        if file_offset >= off:
            pre = p
        else:
            break
    return pre


def file_offset_of(match):
    """File offset of an audit comment.

    A banked image larger than 64 KB cannot fit its file offset in the
    four-digit PRG field, so banked_dis.py records the bank base in a
    BANK: field and leaves PRG as the offset within the bank.
    """
    if match.lastindex and match.lastindex >= 3 and match.group(3):
        return int(match.group(3), 16) + int(match.group(2), 16)
    return int(match.group(2), 16)


def main():
    proj = Path(sys.argv[1])
    mapfile = Path(sys.argv[2])
    bank_prefixes = read_bank_prefixes(sys.argv[3] if len(sys.argv) > 3 else "")
    outfile = Path(sys.argv[4]) if len(sys.argv) > 4 else mapfile

    # The file holds the hand-written routine names and, below a marker, the
    # names this tool derived.  Only the hand-written part is input: reading
    # the derived part back would make the tool accumulate names over
    # repeated runs.
    raw_lines = mapfile.read_text().splitlines()
    marker = "# ---- derived internal branch-target names ----"
    hand = []
    for line in raw_lines:
        if line.startswith(marker):
            break
        hand.append(line)

    # the hand-written routine names
    names = read_map(hand)

    lines = []
    # a banked project has one source file per bank, and the branch graph
    # of a bank is its own: the same CPU address holds different bytes in
    # each.  The file is chosen with --file.
    for f in (sys.argv[5].split(",") if len(sys.argv) > 5 else ["prg.s"]):
        p = proj / f
        if p.is_file():
            lines.extend(p.read_text().split("\n"))

    # every label and its address
    labels = {}
    addr_of = {}
    cur = None
    for i, line in enumerate(lines):
        m = LABEL_RE.match(line)
        if m:
            cur = m.group(1)
            continue
        m2 = CPU_RE.search(line)
        if m2 and cur:
            lm = BANKED_LABEL_RE.match(cur)
            if lm:
                key = (lm.group("pre") or "", int(lm.group("addr"), 16))
            else:
                key = ("", int(m2.group(1), 16))
            labels[key] = int(m2.group(1), 16)
            addr_of[key] = cur
            cur = None

    # every address that carries an instruction
    code_at = set()
    for line in lines:
        m = CPU_RE.search(line)
        if m:
            code_at.add(int(m.group(1), 16))

    # the branches that reach each label
    incoming = defaultdict(list)
    # the calls that reach each label: a `jsr`/`jmp` target that no hand
    # written name covers keeps the address-derived form the repo accepts
    # (`Sub_XXXX`).
    called = set()
    for i, line in enumerate(lines):
        m = re.match(r"^\s*(b[a-z]{2})\s+(\S+)\s*;", line)
        if m:
            op, tgt = m.group(1), m.group(2)
            lm = BANKED_LABEL_RE.match(tgt)
            if not lm:
                continue
            am = CPU_RE.search(line)
            if not am:
                continue
            prefix = prefix_at(bank_prefixes, file_offset_of(am))
            incoming[(prefix, int(lm.group("addr"), 16))].append(
                (op, int(am.group(1), 16)))
            continue
        m = re.match(r"^\s*(jsr|jmp)\s+(\S+)\s*;", line)
        if m:
            lm = BANKED_LABEL_RE.match(m.group(2))
            if lm:
                called.add((lm.group("pre") or "", int(lm.group("addr"), 16)))
                continue
            am = CPU_RE.search(line)
            rm = re.match(r"^\$([0-9A-F]{4})$", m.group(2))
            if rm and am:
                prefix = prefix_at(bank_prefixes, file_offset_of(am))
                addr = int(rm.group(1), 16)
                # A call into RAM ($0000-$07FF) can never be given a label
                # in prg.s - there is no code there - so it is not named.
                if addr < 0x0800:
                    continue
                # A call into unmapped space, or into a region the tracer
                # never decoded, is a false decode: there is no instruction
                # there to label, so the operand stays numeric and the
                # project records it as residue.
                if addr not in code_at:
                    continue
                # A call into $E000-$FFFF always lands in the fixed bank,
                # whichever bank the caller runs in.
                if 0xE000 <= addr <= 0xFFFF:
                    prefix = ""
                called.add((prefix, addr))

    # the owner of each address: the nearest routine entry below it
    entries = defaultdict(list)
    for (prefix, a) in names:
        entries[prefix].append(a)
    for prefix in entries:
        entries[prefix].sort()

    def owner_of(prefix, addr):
        best = None
        for e in entries.get(prefix, []):
            if e <= addr:
                best = e
            else:
                break
        return best

    # name the internal targets
    #
    # A target that already carries a hand-written name keeps it.  A target
    # whose name is only address-derived does not: `PROMPTS.md` asks every
    # branch target to have a name that reflects its role in program flow, and
    # `SubC800F` says nothing about the flow.  Those get the role name the
    # branches reaching them imply.
    internal = defaultdict(list)
    for key, brs in incoming.items():
        if key in names and not address_derived(names[key]):
            continue
        internal[key] = brs

    used = defaultdict(set)
    out = {}
    for key in sorted(internal):
        prefix, a = key
        brs = internal[key]
        roles = []
        for op, baddr in brs:
            if a <= baddr:
                roles.append(ROLE.get(op, "Branch"))
            else:
                roles.append("Loop")
        role = sorted(roles, key=lambda r: (roles.count(r), r))[-1]
        own = owner_of(prefix, a)
        base = f"{prefix + '_' if prefix else ''}{names.get((prefix, own), 'Sub')}_{role}"
        n = 2
        cand = base
        while cand in used[prefix]:
            cand = f"{base}{n}"
            n += 1
        used[prefix].add(cand)
        out[key] = cand

    # A `jsr`/`jmp` target that no hand-written name covers keeps the
    # address-derived form: the repo accepts `Sub_XXXX` for an entry point
    # whose role the dump does not reveal.
    for key in sorted(called):
        if key in names or key in out:
            continue
        # a call into another bank is named in that bank's own file
        if key not in labels:
            continue
        prefix, a = key
        name = f"{prefix + '_' if prefix else ''}Sub_{a:04X}"
        if name in used[prefix]:
            continue
        used[prefix].add(name)
        out[key] = name

    # merge: the hand-written routine names plus the derived ones
    merged = dict(names)
    merged.update(out)

    # keep the header the hand-written map already carries, so the file
    # documents itself in the project's own terms
    header = [l for l in hand if l.startswith("#")]

    with outfile.open("w") as f:
        for line in hand:
            if not line.startswith("#"):
                # a hand-written name for an address that also got a derived
                # role name is replaced by it: the map holds one name per
                # address.
                key = map_key(line)
                if key is not None and key in out:
                    continue
                f.write(line + "\n")
        f.write("\n")
        f.write(marker + "\n")
        for key in sorted(merged):
            prefix, a = key
            f.write(f"{prefix + ':' if prefix else ''}{a:04X}\t{merged[key]}\n")
    print(f"{len(names)} routine names + {len(out)} derived branch names "
          f"= {len(merged)} entries -> {outfile}")


if __name__ == "__main__":
    main()