#!/usr/bin/env python3
"""
Per-function C matcher for the Orbital decomp project.

Given a C source file containing a single function, compiles it with agbcc,
extracts the .text bytes, and compares against the ROM at the expected address.

Relocations (BL/B immediates) are masked before comparison.

Usage:
  python3 tools/check_func.py src/game/GameState_Init.c 0x080001CC
  python3 tools/check_func.py --auto src/game/GameState_Init.c
    (auto mode: reads address from .func_addr directive in the C file)
"""

import os
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")
OBJCOPY = os.path.join(TB, "arm-none-eabi-objcopy")

ROM_BASE = 0x08000000

def mask_reloc_words(data: bytes) -> bytearray:
    """Mask out BL/B immediate bits so relocation targets don't affect comparison."""
    masked = bytearray(data)
    i = 0
    n = len(masked)
    while i < n - 3:
        hw1 = int.from_bytes(masked[i:i+2], 'little')
        hw2 = int.from_bytes(masked[i+2:i+4], 'little')
        # Thumb BL: 11110xxx_xxxxxxxx 11111xxx_xxxxxxxx
        if (hw1 >> 11) == 0b11110 and (hw2 >> 11) == 0b11111:
            hw1 &= 0xF800
            hw2 &= 0xF800
            masked[i]   = hw1 & 0xFF
            masked[i+1] = hw1 >> 8
            masked[i+2] = hw2 & 0xFF
            masked[i+3] = hw2 >> 8
            i += 4
            continue
        i += 2
    return masked

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

def check_function(c_path: str, addr: int) -> dict:
    """Compile C function and compare against ROM."""
    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"}

    c_text = c_path.read_text()

    # Parse expected size from the C file (look for .size directive or function_list)
    # For now, caller provides address, we derive size from function_list
    func_size = None
    with open(ROOT / "function_list_v2.txt") as f:
        for line in f:
            parts = line.strip().split(None, 2)
            if parts and int(parts[0], 16) == addr:
                func_size = int(parts[1])
                break

    if func_size is None or func_size < 2:
        return {"ok": False, "error": f"address 0x{addr:08X} not in function_list_v2.txt"}

    rom = ROM.read_bytes()
    offset = addr - ROM_BASE
    rom_bytes = rom[offset:offset + func_size]

    # Strip trailing padding for comparison
    real_size = func_size
    while real_size > 2 and rom_bytes[real_size - 1] == 0:
        real_size -= 1
    rom_trimmed = rom_bytes[:real_size]

    with tempfile.TemporaryDirectory() as tmpdir:
        s_path = os.path.join(tmpdir, "func.s")
        o_path = os.path.join(tmpdir, "func.o")
        bin_path = os.path.join(tmpdir, "func.bin")

        # Compile with agbcc (the 'Invalid option' message on stderr is harmless)
        r = run([AGBCC, "-O2", "-mthumb-interwork", "-S", str(c_path), "-o", s_path])
        if not Path(s_path).exists():
            return {"ok": False, "error": f"agbcc failed:\n{r.stderr[:500]}"}

        # Assemble
        r = run([AS, "-mcpu=arm7tdmi", s_path, "-o", o_path])
        if r.returncode != 0:
            return {"ok": False, "error": f"as failed:\n{r.stderr[:500]}"}

        # Extract .text
        r = run([OBJCOPY, "-O", "binary", "-j", ".text", o_path, bin_path])
        if r.returncode != 0:
            return {"ok": False, "error": f"objcopy failed:\n{r.stderr[:500]}"}

        compiled = Path(bin_path).read_bytes()

        if len(compiled) < real_size:
            return {"ok": False, "error": f"compiled too small: {len(compiled)} < {real_size}"}

        compiled_trimmed = compiled[:real_size]

        # Compare with relocation masking
        rom_masked = mask_reloc_words(rom_trimmed)
        compiled_masked = mask_reloc_words(compiled_trimmed)

        exact_match = compiled_trimmed == rom_trimmed
        masked_match = compiled_masked == rom_masked

        # Count differing bytes
        if exact_match:
            diffs = 0
        else:
            diffs = sum(1 for a, b in zip(compiled_trimmed, rom_trimmed) if a != b)

        return {
            "ok": True,
            "addr": f"0x{addr:08X}",
            "size": func_size,
            "real_size": real_size,
            "compiled_size": len(compiled),
            "exact_match": exact_match,
            "masked_match": masked_match,
            "diffs": diffs,
            "accuracy_pct": round(100 - 100 * diffs / real_size, 1) if real_size else 0,
        }

def main():
    import argparse
    parser = argparse.ArgumentParser(description="Check a C function against ROM")
    parser.add_argument("c_file", help="C source file with the function")
    parser.add_argument("address", nargs="?", help="ROM address (hex, e.g. 0x080001CC)")
    args = parser.parse_args()

    if args.address:
        addr = int(args.address, 16)
    else:
        # Try to parse from .func_addr in the C file
        addr = None
        for line in Path(args.c_file).read_text().splitlines():
            if line.strip().startswith("//") and "0x08" in line:
                import re
                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 found in C file")

    result = check_function(args.c_file, addr)
    if result.get("ok"):
        status = "EXACT" if result["exact_match"] else ("MASKED" if result["masked_match"] else "DIFF")
        print(f"{status} {result['addr']} size={result['real_size']} diffs={result['diffs']} accuracy={result['accuracy_pct']}%")
    else:
        print(f"FAIL: {result.get('error', 'unknown')}", file=sys.stderr)
        return 1
    return 0

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