#!/usr/bin/env python3
"""
Link-at-address per-function matcher.

Replaces masking-based comparison with relocation resolution: the candidate
C function is compiled, assembled, and linked at its REAL ROM address (via
a generated linker script placing .text at that address, plus stub symbols
for every other known function so BL targets resolve naturally). The linked
bytes for the function's address range are then compared UNMASKED against
the ROM.

A function counts as matched ONLY if the unmasked bytes are identical.

Also supports: python3 tools/check_direct.py --all
  -> writes build/direct_match.json + prints trusted scoreboard.

Exit status: 0 if the single function matched, 1 otherwise.
"""

import json
import os
import re
import subprocess
import sys
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
ROM = ROOT / "bit Generations - Orbital (Japan) (En).gba"
AGBCC = str(ROOT / "agbcc" / "agbcc")
TB = str(ROOT / "arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi" / "bin")
AS = os.path.join(TB, "arm-none-eabi-as")
LD = os.path.join(TB, "arm-none-eabi-ld")
OBJCOPY = os.path.join(TB, "arm-none-eabi-objcopy")

ROM_BASE = 0x08000000
CODE_START = 0x080000C0
CODE_END = 0x0801F30C


def run(cmd, **kw):
    return subprocess.run(cmd, capture_output=True, text=True, **kw)


def load_functions():
    """Load (addr, listed_size, name) from the boundary map (prefer stub_map).
    Alias rows (size 0) are excluded: they own no bytes and must not
    truncate spans."""
    mp = ROOT / "build" / "stub_map.json"
    if mp.exists():
        rows = json.loads(mp.read_text())
        return sorted(
            [(r["addr"] if isinstance(r["addr"], int) else int(r["addr"], 16),
              r["size"], r["name"]) for r in rows
             if r.get("size", 0) >= 2 and not r.get("alias")]
        )
    rows = []
    with open(ROOT / "function_list_v2.txt") as f:
        for line in f:
            m = re.search(r"([0-9A-Fa-f]{6,8})\s+(\d+)\s+(.+)", line)
            if m:
                rows.append((int(m.group(1), 16), int(m.group(2)), m.group(3).strip()))
    # dedupe by address
    seen = {}
    for a, s, n in rows:
        seen.setdefault(a, (s, n))
    return sorted((a, s, n) for a, (s, n) in seen.items())


def real_extent(addr):
    """Return (real_size, rom_bytes) for a function: [start, last return+2),
    scanning forward with proper 32-bit Thumb handling."""
    rom = ROM.read_bytes()
    n = len(rom)
    off = addr - ROM_BASE
    # Bound the scan by the next known function start (gap arithmetic gives
    # an upper bound on this function's extent; the last return inside it
    # is the real end). Keeps early returns from being misread and keeps
    # us out of neighbors.
    funcs = load_functions()
    starts = sorted(a for a, _, _ in funcs)
    idx = starts.index(addr) if addr in starts else -1
    hi = starts[idx + 1] - ROM_BASE if 0 <= idx < len(starts) - 1 else off + 20000
    limit = min(n - 1, hi)
    last_ret = -1
    i = off
    while i < limit:
        hw = int.from_bytes(rom[i:i + 2], "little")
        if hw == 0x4770:
            last_ret = i
            i += 2
            continue
        if (hw >> 8) == 0xBD and (hw & 0x100):
            last_ret = i
            i += 2
            continue
        if (hw & 0xFF87) == 0x4700:
            last_ret = i
            i += 2
            continue
        if hw == 0x46F7:  # mov pc, lr (old-ARM return idiom)
            last_ret = i
            i += 2
            continue
        if (hw & 0xFF87) == 0x4687:  # mov pc, Rm (noreturn tail-jump,
            last_ret = i            # e.g. coroutine stack-switch)
            i += 2
            continue
        # 32-bit Thumb prefix words: BL/BLX prefix, misc 32-bit encodings.
        # Only BL (11110) and misc-prefix (11101) consume 4 bytes; the
        # 11111 halfword is the SECOND half of a BL pair, never a start.
        if (hw >> 11) in (0b11110, 0b11101):
            i += 4
        else:
            i += 2
        if i - off > 19990:
            break
    if last_ret < 0:
        # Shared-epilogue case: the function's only return is the first
        # halfword of the NEXT listed function (e.g. `pop {r3}` here +
        # `bx r3` owned by the neighbor). Claim exactly that halfword.
        if hi < n - 1:
            _h = int.from_bytes(rom[hi:hi + 2], "little")
            if _h == 0x4770 or ((_h >> 8) == 0xBD and (_h & 0x100)) or (_h & 0xFF87) == 0x4700:
                last_ret = hi
        if last_ret < 0:
            return None, None
    end = last_ret + 2
    size = end - off
    # Strip inter-function alignment padding: a trailing pad halfword
    # (bx lr / nop / 0000) that sits AFTER a real return belongs to the
    # gap, not the function (e.g. `movs r0,#30; bx lr; bx lr; 0000`).
    # Only strip when the previous halfword is itself a return, so real
    # `...; bx lr` endings are never touched.
    def _hw(at):
        return int.from_bytes(rom[at:at + 2], "little")
    def _is_ret(hw):
        return hw == 0x4770 or hw == 0x46F7 or ((hw >> 8) == 0xBD and (hw & 0x100)) or (hw & 0xFF87) == 0x4700 or (hw & 0xFF87) == 0x4687
    first_end = off + 2  # earliest plausible end anchor (updated below)
    # find first return to anchor the pad run
    _fi = off
    first_end = end
    while _fi < end:
        if _is_ret(_hw(_fi)):
            first_end = _fi + 2
            break
        _fi += 2
    while end > first_end and _hw(end - 2) in (0x4770, 0x46C0, 0x0000):
        end -= 2
        size -= 2
    while size > 2 and rom[end - 1] == 0:
        end -= 1
        size -= 1
    return size, rom[off:end]


def span_end(addr):
    """End of a function's owned span: next known start, else CODE_END."""
    funcs = load_functions()
    starts = sorted(a for a, _, _ in funcs)
    idx = starts.index(addr) if addr in starts else -1
    if 0 <= idx < len(starts) - 1:
        return starts[idx + 1]
    return CODE_END


def prefix_verified(comp: bytes, addr: int, hi: int) -> int:
    """Longest prefix of linked bytes equal to ROM[addr:hi]."""
    rom = ROM.read_bytes()
    want = rom[addr - ROM_BASE:hi - ROM_BASE]
    n = 0
    for a, b in zip(comp, want):
        if a != b:
            break
        n += 1
    return min(n, len(want))


def _stub_objects(ext_names, known, td):
    stub_lines = ["\t.thumb"]
    for en in sorted(ext_names):
        if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", en):
            continue
        stub_lines.append(f"\t.globl {en}")
        if en in known:
            stub_lines.append(f"\t.type {en}, %function")
            stub_lines.append("\t.thumb_func")
            stub_lines.append(f"\t.set {en}, 0x{(known[en] | 1):08X}")
        else:
            stub_lines.append(f"\t.set {en}, 0x08000000")
    stub_s = td / "stubs.s"
    stub_o = td / "stubs.o"
    stub_s.write_text("\n".join(stub_lines) + "\n")
    r = run([AS, "-mcpu=arm7tdmi", str(stub_s), "-o", str(stub_o)])
    return stub_o, r


def _link_and_prefix(o_path, addr, hi, text_refs, stem, td):
    """Link object at addr with stub syms; return (verified, span_len, detail)."""
    funcs = load_functions()
    known = {n: a for a, s, n in funcs}
    try:
        import json as _j
        for _r in _j.loads((ROOT / "build" / "stub_map.json").read_text()):
            if _r.get("alias"):
                _a = _r["addr"] if isinstance(_r["addr"], int) else int(_r["addr"], 16)
                known.setdefault(_r["name"], _a)
    except Exception:
        pass
    ext_names = set(re.findall(r"^\s*\.extern\s+(\w+)", text_refs, re.M))
    ext_names |= set(re.findall(r"\bbl\s+([A-Za-z_]\w*)", text_refs))
    ext_names |= set(re.findall(r"^\s*\.word\s+([A-Za-z_]\w*)", text_refs, re.M))
    ext_names -= {stem}
    ext_names = {e for e in ext_names if not e[0].isdigit()}
    stub_o, r = _stub_objects(ext_names, known, td)
    if r.returncode != 0:
        return 0, hi - addr, f"stub-as: {r.stderr[:200]}"
    elf, bin_path = td / "func.elf", td / "func.bin"
    ld_script = td / "link.ld"
    ld_script.write_text(
        "OUTPUT_FORMAT(\"elf32-littlearm\")\nOUTPUT_ARCH(arm)\n"
        "SECTIONS {\n"
        f"  .text 0x{addr:08X} : {{ *(.text) *(.text.*) }}\n"
        "  /DISCARD/ : { *(*) }\n"
        "}\n"
    )
    r = run([LD, "-T", str(ld_script), str(o_path), str(stub_o), "-o", str(elf)])
    if r.returncode != 0:
        return 0, hi - addr, f"ld: {r.stderr[:200]}"
    r = run([OBJCOPY, "-O", "binary", "-j", ".text", str(elf), str(bin_path)])
    if r.returncode != 0:
        return 0, hi - addr, f"objcopy: {r.stderr[:200]}"
    comp = bin_path.read_bytes()
    try:
        nm = run([str(Path(AS).parent / "arm-none-eabi-nm"), str(elf)])
        for _ln in nm.stdout.splitlines():
            _m = re.match(r"\s*([0-9a-fA-F]+)\s+\w\s+(\S+)", _ln)
            if _m and _m.group(2) == stem:
                _off = int(_m.group(1), 16) - addr
                if _off >= 0:
                    comp = comp[_off:]
                break
    except Exception:
        pass
    v = prefix_verified(comp, addr, hi)
    return v, hi - addr, f"verified {v}/{hi - addr}"


def check_any(src_path: str, addr: int, flagsets=(("-O2", "-mthumb-interwork"),
                                                  ("-O2", "-mno-thumb-interwork"),
                                                  ("-O1", "-mthumb-interwork"),
                                                  ("-O1", "-mno-thumb-interwork"),
                                                  ("-O2", "-mthumb-interwork", "-fno-omit-frame-pointer"),
                                                  ("-O2", "-mno-thumb-interwork", "-fno-omit-frame-pointer"))) -> dict:
    """Unified verifier for .c (agbcc, all flagsets) and .s (direct asm).
    Counts the full-span verified prefix, not just bytes to last return."""
    src = Path(src_path)
    if not src.exists():
        return {"ok": False, "error": f"file not found: {src_path}"}
    funcs = load_functions()
    if addr not in {a for a, _, _ in funcs}:
        return {"ok": False, "error": f"address 0x{addr:08X} not in function map"}
    hi = span_end(addr)
    real = real_extent(addr)
    real_size = real[0] if real[0] is not None else 0
    _td_base = ROOT / ".scratch"
    _td_base.mkdir(parents=True, exist_ok=True)
    best = {"v": 0, "flags": "", "detail": ""}
    with tempfile.TemporaryDirectory(dir=str(_td_base)) as td:
        td = Path(td)
        if src.suffix.lower() == ".s":
            o_path = td / "func.o"
            r = run([AS, "-mcpu=arm7tdmi", str(src), "-o", str(o_path)])
            if r.returncode != 0:
                return {"ok": True, "matched": False, "addr": f"0x{addr:08X}",
                        "real_size": real_size, "verified": 0, "span": hi - addr,
                        "detail": f"as: {r.stderr[:300]}"}
            text_refs = src.read_text(errors="replace")
            v, span, det = _link_and_prefix(o_path, addr, hi, text_refs, src.stem, td)
            full = (v == span)
            return {"ok": True, "matched": full, "addr": f"0x{addr:08X}",
                    "real_size": real_size, "verified": v, "span": span,
                    "flags": "asm", "detail": det}
        s_path, o_path = td / "func.s", td / "func.o"
        for flags in flagsets:
            r = run([AGBCC, *flags, "-S", str(src), "-o", str(s_path)])
            if not s_path.exists():
                best["detail"] += f"agbcc [{' '.join(flags)}]: {r.stderr[:120]}; "
                continue
            r = run([AS, "-mcpu=arm7tdmi", str(s_path), "-o", str(o_path)])
            if r.returncode != 0:
                best["detail"] += f"as [{' '.join(flags)}]: {r.stderr[:120]}; "
                continue
            text_refs = s_path.read_text(errors="replace")
            v, span, det = _link_and_prefix(o_path, addr, hi, text_refs, src.stem, td)
            if v > best["v"]:
                best = {"v": v, "flags": " ".join(flags), "detail": det}
            if v == span:
                break
    span = hi - addr
    full = (best["v"] == span)
    legacy_ok = (best["v"] >= real_size > 0)
    return {"ok": True, "matched": full, "legacy": legacy_ok,
            "addr": f"0x{addr:08X}", "real_size": real_size,
            "verified": best["v"], "span": span, "flags": best["flags"],
            "detail": best["detail"]}


def check_function_direct(c_path: str, addr: int, flagsets=(("-O2", "-mthumb-interwork"),
                                                            ("-O2", "-mno-thumb-interwork"),
                                                            ("-O1", "-mthumb-interwork"),
                                                            ("-O1", "-mno-thumb-interwork"),
                                                            ("-O2", "-mthumb-interwork", "-fno-omit-frame-pointer"),
                                                            ("-O2", "-mno-thumb-interwork", "-fno-omit-frame-pointer"))) -> dict:
    """Compile + link candidate C at its real address; unmasked compare."""
    c_path = Path(c_path)
    if not c_path.exists():
        return {"ok": False, "error": f"file not found: {c_path}"}
    if not ROM.exists():
        return {"ok": False, "error": "ROM not found"}

    funcs = load_functions()
    names = {a: n for a, s, n in funcs}
    if addr not in names and addr not in {a for a, _, _ in funcs}:
        return {"ok": False, "error": f"address 0x{addr:08X} not in function map"}

    real = real_extent(addr)
    if real[0] is None:
        return {"ok": False, "error": "no return instruction found; boundary unknown"}
    real_size, rom_bytes = real

    errors = []
    _td_base = ROOT / ".scratch"
    _td_base.mkdir(parents=True, exist_ok=True)
    with tempfile.TemporaryDirectory(dir=str(_td_base)) as td:
        td = Path(td)
        s_path = td / "func.s"
        o_path = td / "func.o"
        elf = td / "func.elf"
        bin_path = td / "func.bin"

        for flags in flagsets:
            r = run([AGBCC, *flags, "-S", str(c_path), "-o", str(s_path)])
            if not s_path.exists():
                errors.append(f"agbcc [{ ' '.join(flags)}]: {r.stderr[:300]}")
                continue
            r = run([AS, "-mcpu=arm7tdmi", str(s_path), "-o", str(o_path)])
            if r.returncode != 0:
                errors.append(f"as [{ ' '.join(flags)}]: {r.stderr[:300]}")
                continue

            # Linker script: .text of the candidate AT the real address.
            # Callee/data symbols are provided via a stubs object: known
            # function names are declared `.thumb_func` + `.set sym, addr|1`
            # (the LSB=1 Thumb bit is REQUIRED - a plain PROVIDE/abs
            # address makes ld emit a Thumb->ARM veneer and the BL then
            # encodes the veneer address instead of the real callee).
            # Unknown names get a filler data address (0x08000000).
            s_text = s_path.read_text(errors="replace")
            ext_names = set(re.findall(r"^\s*\.extern\s+(\w+)", s_text, re.M))
            ext_names |= set(re.findall(r"\bbl\s+([A-Za-z_]\w*)", s_text))
            # agbcc emits extern data refs as `.word <sym>` literal-pool
            # entries, not `.extern` lines: scrape those too, but only
            # bare-symbol words (never `.word 0x...` / `.word 123`).
            ext_names |= set(re.findall(r"^\s*\.word\s+([A-Za-z_]\w*)", s_text, re.M))
            ext_names -= {Path(c_path).stem}
            # exclude numeric-looking words just in case
            ext_names = {e for e in ext_names if not e[0].isdigit()}
            known = {n: a for a, s, n in funcs}
            # merge alias symbols (alternate entry points) for BL resolution
            try:
                import json as _j
                for _r in _j.loads((ROOT / "build" / "stub_map.json").read_text()):
                    if _r.get("alias"):
                        _a = _r["addr"] if isinstance(_r["addr"], int) else int(_r["addr"], 16)
                        known.setdefault(_r["name"], _a)
            except Exception:
                pass
            stub_lines = ["\t.thumb"]
            for en in sorted(ext_names):
                if not re.match(r"^[A-Za-z_][A-Za-z0-9_]*$", en):
                    continue
                stub_lines.append(f"\t.globl {en}")
                if en in known:
                    stub_lines.append(f"\t.type {en}, %function")
                    stub_lines.append("\t.thumb_func")
                    stub_lines.append(f"\t.set {en}, 0x{(known[en] | 1):08X}")
                else:
                    stub_lines.append(f"\t.set {en}, 0x08000000")
            stub_s = td / "stubs.s"
            stub_o = td / "stubs.o"
            stub_s.write_text("\n".join(stub_lines) + "\n")
            r = run([AS, "-mcpu=arm7tdmi", str(stub_s), "-o", str(stub_o)])
            if r.returncode != 0:
                errors.append(f"stub-as [{ ' '.join(flags)}]: {r.stderr[:300]}")
                continue
            ld_script = td / "link.ld"
            ld_script.write_text(
                "OUTPUT_FORMAT(\"elf32-littlearm\")\nOUTPUT_ARCH(arm)\n"
                "SECTIONS {\n"
                f"  .text 0x{addr:08X} : {{ *(.text) *(.text.*) }}\n"
                "  /DISCARD/ : { *(*) }\n"
                "}\n"
            )
            r = run([LD, "-T", str(ld_script), str(o_path), str(stub_o), "-o", str(elf)])
            if r.returncode != 0:
                errors.append(f"ld [{ ' '.join(flags)}]: {r.stderr[:300]}")
                continue
            r = run([OBJCOPY, "-O", "binary", "-j", ".text", str(elf), str(bin_path)])
            if r.returncode != 0:
                errors.append(f"objcopy [{ ' '.join(flags)}]: {r.stderr[:300]}")
                continue

            comp = bin_path.read_bytes()
            # Slice from the function's LINKED symbol address, not offset 0:
            # agbcc emits `.align 2,0` at function start, which inserts pad
            # bytes when addr%4==2, shifting raw offset-0 bytes.
            try:
                nm = run([str(Path(AS).parent / "arm-none-eabi-nm"), str(elf)])
                sym_off = None
                for _ln in nm.stdout.splitlines():
                    _m = re.match(r"\s*([0-9a-fA-F]+)\s+\w\s+(\S+)", _ln)
                    if _m and _m.group(2) == Path(c_path).stem:
                        sym_off = int(_m.group(1), 16) - addr
                        break
                if sym_off is not None and sym_off >= 0:
                    comp = comp[sym_off:]

            except Exception:
                pass
            if len(comp) < real_size:
                errors.append(
                    f"linked too small [{ ' '.join(flags)}]: {len(comp)} < {real_size}")
                continue
            ct = comp[:real_size]
            diffs = sum(1 for a, b in zip(ct, rom_bytes) if a != b)
            if diffs == 0:
                # also report pool-word equality detail for diagnostics
                return {
                    "ok": True, "matched": True, "flags": " ".join(flags),
                    "addr": f"0x{addr:08X}", "real_size": real_size,
                    "diffs": 0,
                }
            errors.append(f"diff {diffs}/{real_size} [{ ' '.join(flags)}]")

    return {"ok": True, "matched": False, "addr": f"0x{addr:08X}",
            "real_size": real_size, "detail": "; ".join(errors[:4])}


def main():
    import argparse
    p = argparse.ArgumentParser(description="Link-at-address C function matcher")
    p.add_argument("c_file", nargs="?", help="C source file with the function")
    p.add_argument("address", nargs="?", help="ROM address (hex)")
    p.add_argument("--all", action="store_true",
                   help="check all src/matched/*.c, write build/direct_match.json")
    args = p.parse_args()

    if args.all:
        funcs = {n.lower(): a for a, s, n in load_functions()}
        hits = []
        total_real = 0
        total_verified = 0
        n_files = 0
        cand = sorted((ROOT / "src" / "matched").glob("*.c")) + sorted((ROOT / "src" / "matched").glob("*.s"))
        for f in cand:
            base = f.stem.lower()
            if base not in funcs:
                print(f"SKIP {f.name} (no address mapping)")
                continue
            addr = funcs[base]
            r = check_any(str(f), addr)
            n_files += 1
            if r.get("matched") or r.get("legacy"):
                total_real += r["real_size"]
                total_verified += r["verified"]
                hits.append({"name": f.stem, "addr": f"0x{addr:08X}",
                             "real_size": r["real_size"],
                             "verified": r["verified"], "span": r["span"],
                             "flags": r["flags"]})
                tag = "FULL" if r.get("matched") else "     "
                print(f"MATCH {tag} {f.stem:40s} {r['real_size']:5d}B real {r['verified']:5d}B ver [{r['flags']}]")
            else:
                det = r.get("detail") or r.get("error") or "unknown"
                print(f"FAIL  {f.stem:40s} {str(det)[:160]}")
        code_size = CODE_END - CODE_START
        print(f"\nTRUSTED: {len(hits)}/{n_files} files, "
              f"{total_real}/{code_size} bytes = {100*total_real/code_size:.2f}%")
        print(f"VERIFIED: {total_verified}/{code_size} bytes = {100*total_verified/code_size:.2f}% (full spans incl pools/tails)")
        (ROOT / "build" / "direct_match.json").write_text(json.dumps(
            {"matched": hits, "matched_bytes": total_real,
             "verified_bytes": total_verified,
             "code_size": code_size,
             "accuracy_pct": round(100 * total_real / code_size, 2),
             "verified_pct": round(100 * total_verified / code_size, 2)}, indent=2))
        return 0

    if not args.c_file:
        p.error("c_file required (or use --all)")
    addr = None
    if args.address:
        addr = int(args.address, 16)
    else:
        for line in Path(args.c_file).read_text().splitlines():
            if line.strip().startswith("//") and "0x08" in line:
                m = re.search(r'0x08[0-9A-Fa-f]{6}', line)
                if m:
                    addr = int(m.group(), 16)
                    break
        if addr is None:
            sys.exit("No address provided and no // 0x08XXXXXX comment in C file")
    r = check_any(args.c_file, addr)
    if r.get("matched") or r.get("legacy"):
        print(f"MATCH {r['addr']} real={r['real_size']} verified={r['verified']}/{r['span']} [{r['flags']}] (unmasked)")
        return 0
    print(f"NO-MATCH {r.get('addr', '')} "
          f"{r.get('detail', r.get('error', 'unknown'))[:400]}")
    return 1


if __name__ == "__main__":
    raise SystemExit(main())