#!/usr/bin/env python3
"""Turn PC-relative branch operands into labels.

`cdl_reclass.py` emits a branch whose target it cannot prove to be an
instruction boundary as an explicit relative offset (`beq *+$05`), because
ca65 cannot resolve a bare number as a branch target.  That is safe, but it
hides the target from the reader and from the Pass 2 naming.

This tool resolves those offsets: it computes the target address (ca65's `*`
is the address of the current instruction), gives the target a label if it
does not have one, and rewrites the operand as `beq Label_XXXX`.

The bytes cannot change - the offset and the label denote the same address -
so `make` still has to print BYTE-IDENTICAL.

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

import argparse
import re
import sys
from pathlib import Path

RAW_RE = re.compile(r"^(\s*(?:b[a-z]{2}|bra)\s+)\*([+-]?)\$([0-9A-Fa-f]{1,4})\s*$")
CPU_RE = re.compile(r"; CPU:\$([0-9A-Fa-f]{4}) PRG:")
LABEL_RE = re.compile(r"^([A-Za-z_][A-Za-z0-9_]*):\s*$")


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("dir", help="project directory")
    ap.add_argument("--write", action="store_true", help="apply (default: dry run)")
    args = ap.parse_args()

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

    # which address does each source line carry?
    addr_of_line = {}
    for i, line in enumerate(lines):
        m = CPU_RE.search(line)
        if m:
            addr_of_line[i] = int(m.group(1), 16)

    # label name -> address
    label_at = {}
    for i, line in enumerate(lines):
        m = LABEL_RE.match(line)
        if m and i in addr_of_line:
            label_at.setdefault(int(m.group(1), 16), m.group(1))

    # the lines that are data rows rather than instructions
    data_lines = set()
    for i, line in enumerate(lines):
        if line.strip().startswith((".byte", ".word", ".dword", ".incbin")):
            data_lines.add(i)

    # find the raw branches and their targets
    fixes = []          # (line index, old operand, new operand, target address)
    for i, line in enumerate(lines):
        m = RAW_RE.match(line.split(";")[0].rstrip())
        if not m:
            continue
        addr = addr_of_line.get(i)
        if addr is None:
            continue
        delta = int(m.group(3), 16)
        target = addr - delta if m.group(2) == '-' else addr + delta
        fixes.append((i, addr, m.group(0).strip(), target))

    if not fixes:
        print("no raw relative branches found")
        return 0

    # the line each target address sits on
    target_line = {}
    for i, a in addr_of_line.items():
        target_line.setdefault(a, i)

    # every target must land on an *instruction*, or there is nothing to
    # label: a target that sits in a `.byte` row is data, and putting a
    # label on it would claim the data is code
    missing = [t for _i, _a, _o, t in fixes if t not in target_line]
    if missing:
        print("target(s) with no instruction to label: " +
              ", ".join(f"${t:04X}" for t in sorted(set(missing))))
        return 1

    # a target in a data block is left as the raw relative offset the
    # reclassifier emitted, which is the repo's documented form
    data_targets = {t for _i, _a, _o, t in fixes
                    if target_line[t] in data_lines}
    if data_targets:
        print("target(s) in a data block, left as a raw relative offset: " +
              ", ".join(f"${t:04X}" for t in sorted(data_targets)))
        fixes = [f for f in fixes if f[3] not in data_targets]

    # the name each target gets: the label it already has, or a new one
    name_of = {}
    for _i, _a, _o, t in fixes:
        if t not in name_of:
            name_of[t] = label_at.get(t, f"Label_{t:04X}")
    unlabelled = sorted(t for t in name_of if t not in label_at)

    # rewrite the branches
    out = list(lines)
    for i, _addr, old, target in fixes:
        code = out[i].split(";")[0].rstrip()
        code = RAW_RE.sub(lambda m: f"{m.group(1)}{name_of[target]}", code)
        out[i] = code + out[i][len(out[i].split(";")[0]):]

    # define the new labels, just above the target's instruction
    for target in unlabelled:
        i = target_line[target]
        out.insert(i, f"{name_of[target]}:")
        # the insertion shifts the line numbers of everything after it
        for k in range(len(fixes)):
            if fixes[k][0] > i:
                fixes[k] = (fixes[k][0] + 1, fixes[k][1], fixes[k][2], fixes[k][3])
        for a in list(target_line):
            if a > i:
                target_line[a] += 1

    print(f"{len(fixes)} raw relative branch(es) resolved, "
          f"{len(unlabelled)} label(s) added")
    for _i, addr, old, target in fixes:
        print(f"  ${addr:04X}: {old} -> {name_of[target]}")

    if args.write:
        prg.write_text("\n".join(out))
        print("written")
    return 0


if __name__ == "__main__":
    sys.exit(main())