#!/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.

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:")


def main():
    proj = Path(sys.argv[1])
    mapfile = Path(sys.argv[2])

    # the hand-written routine names
    names = {}
    for raw in mapfile.read_text().splitlines():
        line = raw.split("#", 1)[0].strip()
        if not line:
            continue
        parts = line.split("\t") if "\t" in line else line.split(None, 1)
        names[int(parts[0].lstrip("$"), 16)] = parts[1].strip()

    lines = (proj / "prg.s").read_text().split("\n")

    # every label and its address
    labels = {}
    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:
            labels[cur] = int(m2.group(1), 16)
            cur = None
    addr_of = {a: n for n, a in labels.items()}

    # the branches that reach each label
    incoming = defaultdict(list)
    for i, line in enumerate(lines):
        m = re.match(r"^\s*(b[a-z]{2})\s+(\S+)\s*;", line)
        if not m:
            continue
        op, tgt = m.group(1), m.group(2)
        if not tgt.startswith("Label_"):
            continue
        am = CPU_RE.search(line)
        if not am:
            continue
        incoming[tgt].append((op, int(am.group(1), 16)))

    # the owner of each address: the nearest routine entry below it
    entries = sorted(a for a in names)
    def owner_of(addr):
        best = None
        for e in entries:
            if e <= addr:
                best = e
            else:
                break
        return best

    # name the internal targets
    internal = defaultdict(list)
    for lab, brs in incoming.items():
        a = labels.get(lab)
        if a is None or a in names:
            continue
        internal[a] = brs

    used = defaultdict(set)
    out = {}
    for a in sorted(internal):
        brs = internal[a]
        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(a)
        base = f"{names.get(own, 'Sub')}_{role}"
        n = 2
        cand = base
        while cand in used[own]:
            cand = f"{base}{n}"
            n += 1
        used[own].add(cand)
        out[a] = cand

    # 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 = []
    for raw in mapfile.read_text().splitlines():
        if raw.startswith("#"):
            header.append(raw)
        elif raw.strip():
            break

    with mapfile.open("w") as f:
        for line in header:
            f.write(line + "\n")
        if header and not header[-1].endswith("name map"):
            f.write("\n")
        f.write("# ---- derived internal branch-target names ----\n")
        for a in sorted(merged):
            f.write(f"{a:04X}\t{merged[a]}\n")
    print(f"{len(names)} routine names + {len(out)} derived branch names "
          f"= {len(merged)} entries")


if __name__ == "__main__":
    main()