#!/usr/bin/env python3
"""Work out the --override list a cdl_reclass run needs.

`cdl_reclass.py` walks a code region linearly from its start.  When the CPU
entered the region somewhere else - a branch target or a call target the
capture never executed - the walk misses that address, and the tool emits the
branch as a raw relative offset and warns.

This reads the warnings the reclassifier printed (or re-derives them from the
source), decodes the instruction at each target, and prints the override
arguments that force those addresses into the code region.

Usage:
    python3 _tools/plan_overrides.py <project dir> <cdl file> [--base 0x8000]
"""

import argparse
import re
import subprocess
import sys
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent))
import nesdis  # noqa: E402

RAW = re.compile(r"^\s*(?:b[a-z]{2}|bra)\s+\*([+-])\$([0-9A-Fa-f]{1,4})\s*$")
ADDR = re.compile(r"; CPU:\$([0-9A-F]{4}) PRG:")


def main():
    ap = argparse.ArgumentParser(description=__doc__.splitlines()[0])
    ap.add_argument("project")
    ap.add_argument("cdl")
    ap.add_argument("--base", default="0x8000")
    args = ap.parse_args()

    d = Path(args.project)
    prg = (d / "prg.bin").read_bytes()
    base = int(args.base, 0)

    # The raw branch targets in the current source.
    targets = []
    for line in (d / "prg.s").read_text().split("\n"):
        code = line.split(";")[0]
        m = RAW.match(code.rstrip())
        if not m:
            continue
        am = ADDR.search(line)
        if not am:
            continue
        off = int(am.group(1), 16) - base
        rel = int(m.group(2), 16)
        if m.group(1) == "-":
            rel = -rel
        # ca65's `*` is the address of the branch itself, so the target is
        # the branch's address plus the signed offset.
        targets.append(off + rel)

    seen, uniq = set(), []
    for t in targets:
        if t not in seen and 0 <= t < len(prg):
            seen.add(t)
            uniq.append(t)

    print(f"# {len(uniq)} raw branch target(s) to force into the code region")
    args_out = []
    for t in uniq:
        b = prg[t]
        ent = nesdis.OPCODES.get(b)
        if ent is None:
            n = 1
            what = f"undocumented ${b:02X}"
        else:
            n = nesdis.MODE_LEN[ent[1]]
            what = f"{ent[0]} {ent[1]}"
        lo = base + t
        hi = base + t + n - 1
        args_out.append(f"0x{lo:X}-0x{hi:X}:C")
        print(f"#  ${lo:04X}: {what}  ->  --override '0x{lo:X}-0x{hi:X}:C'")
    print()
    print(" ".join(f"--override '{a}'" for a in args_out))


if __name__ == "__main__":
    main()