#!/usr/bin/env python3
"""
Binary Accuracy Matching Tool for GBA ROM decompilation.
Compiles C functions with agbcc and compares byte-by-byte with original ROM.
"""
import os
import sys
import subprocess
import struct
import re
import tempfile
import hashlib

ROM_PATH = "bit Generations - Orbital (Japan) (En).gba"
AGBCC = "agbcc/agbcc"
AS = "arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi/bin/arm-none-eabi-as"
OBJCOPY = "arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi/bin/arm-none-eabi-objcopy"
OBJDUMP = "arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi/bin/arm-none-eabi-objdump"

# ROM base address
ROM_BASE = 0x08000000

def read_rom():
    with open(ROM_PATH, "rb") as f:
        return f.read()

def rom_offset(addr):
    """Convert ROM address to file offset"""
    return addr - ROM_BASE

def extract_function_bytes(rom, addr, size):
    """Extract bytes for a function from the ROM"""
    offset = rom_offset(addr)
    return rom[offset:offset+size]

def compile_c_to_asm(c_source, flags="-O2"):
    """Compile C to assembly using agbcc"""
    with tempfile.NamedTemporaryFile(suffix=".c", mode="w", delete=False) as f:
        f.write(c_source)
        c_path = f.name
    s_path = c_path.replace(".c", ".s")
    
    try:
        result = subprocess.run(
            [AGBCC, flags, c_path],
            capture_output=True, text=True, timeout=30
        )
        if result.returncode != 0:
            return None, result.stderr
        
        with open(s_path) as f:
            asm = f.read()
        return asm, None
    except Exception as e:
        return None, str(e)
    finally:
        os.unlink(c_path)
        if os.path.exists(s_path):
            os.unlink(s_path)

def assemble_asm(asm_source, base_addr=0):
    """Assemble Thumb code"""
    with tempfile.NamedTemporaryFile(suffix=".s", mode="w", delete=False) as f:
        # Add proper Thumb directives
        f.write(".syntax unified\n")
        f.write(".thumb\n")
        f.write(".text\n")
        f.write(".align 2\n")
        f.write(asm_source + "\n")
        s_path = f.name
    
    o_path = s_path.replace(".s", ".o")
    bin_path = s_path.replace(".s", ".bin")
    
    try:
        # Assemble
        result = subprocess.run(
            [AS, "-mcpu=arm7tdmi", "-mthumb-interwork", s_path, "-o", o_path],
            capture_output=True, text=True, timeout=30
        )
        if result.returncode != 0:
            return None, result.stderr
        
        # Extract binary
        result = subprocess.run(
            [OBJCOPY, "-O", "binary", o_path, bin_path],
            capture_output=True, text=True, timeout=30
        )
        if result.returncode != 0:
            return None, result.stderr
        
        with open(bin_path, "rb") as f:
            binary = f.read()
        return binary, None
    except Exception as e:
        return None, str(e)
    finally:
        for path in [s_path, o_path, bin_path]:
            if os.path.exists(path):
                os.unlink(path)

def bytes_match(our_bytes, orig_bytes):
    """Check if two byte sequences match"""
    if len(our_bytes) != len(orig_bytes):
        return False
    return our_bytes == orig_bytes

def match_rate(our_bytes, orig_bytes):
    """Calculate match rate for two byte sequences"""
    if len(orig_bytes) == 0:
        return 1.0 if len(our_bytes) == 0 else 0.0
    
    min_len = min(len(our_bytes), len(orig_bytes))
    matching = sum(1 for i in range(min_len) if our_bytes[i] == orig_bytes[i])
    
    # Penalize length difference
    len_penalty = abs(len(our_bytes) - len(orig_bytes)) / len(orig_bytes)
    
    return (matching / len(orig_bytes)) * (1 - min(len_penalty, 1.0))

def try_compile_and_match(c_source, rom, addr, size, flags="-O2"):
    """Try to compile C code and match with ROM"""
    orig_bytes = extract_function_bytes(rom, addr, size)
    
    # Compile C to assembly
    asm, err = compile_c_to_asm(c_source, flags)
    if asm is None:
        return None, None, f"Compile error: {err}"
    
    # Extract just the function body (skip directives)
    # Find the function label
    lines = asm.split('\n')
    func_lines = []
    in_func = False
    for line in lines:
        stripped = line.strip()
        if stripped.endswith(':') and not stripped.startswith('.'):
            in_func = True
        if in_func:
            func_lines.append(line)
        if stripped.startswith('.Lfe') or stripped.startswith('.size'):
            break
    
    func_asm = '\n'.join(func_lines)
    if not func_asm:
        return None, None, "No function body found"
    
    # Assemble
    binary, err = assemble_asm(func_asm)
    if binary is None:
        return None, None, f"Assemble error: {err}"
    
    # Compare
    if bytes_match(binary, orig_bytes):
        return binary, orig_bytes, "MATCH"
    else:
        rate = match_rate(binary, orig_bytes)
        return binary, orig_bytes, f"PARTIAL ({rate*100:.1f}%)"

def scan_functions(rom, start_addr, end_addr):
    """
    Scan ROM for function-like patterns.
    Functions start at word-aligned addresses and end with bx lr or pop {..., pc}.
    """
    functions = []
    
    addr = start_addr
    while addr < end_addr:
        offset = rom_offset(addr)
        if offset + 4 > len(rom):
            break
        
        # Read Thumb instruction pair
        hw1 = struct.unpack_from('<H', rom, offset)[0]
        
        # Check for common function prologues
        is_func_start = False
        
        # push {r4-r7, lr} or similar
        if (hw1 & 0xFF00) == 0xB500:  # push {..., lr}
            is_func_start = True
        elif (hw1 & 0xFF00) == 0xB400:  # push {...}
            # Check if next instruction is also push or if lr is included
            if hw1 & 0x1000:  # lr bit set
                is_func_start = True
        elif (hw1 & 0xFF00) == 0x4600:  # mov (high register)
            is_func_start = True
        elif (hw1 & 0xF800) == 0xB000:  # sub sp, #imm
            is_func_start = True
        
        if is_func_start:
            # Try to find function end (bx lr or pop {..., pc})
            func_size = 4
            found_end = False
            for i in range(1, 256):  # Max 256 halfwords = 512 bytes
                check_offset = rom_offset(addr + i * 2)
                if check_offset + 2 > len(rom):
                    break
                hw = struct.unpack_from('<H', rom, check_offset)[0]
                
                # bx lr
                if hw == 0x4770:
                    func_size = (i + 1) * 2
                    found_end = True
                    break
                
                # pop {..., pc}
                if (hw & 0xFF00) == 0xBD00:
                    func_size = (i + 1) * 2
                    found_end = True
                    break
                
                # pop {...} without pc, continue
                if (hw & 0xFF00) == 0xBC00 and not (hw & 0x0100):
                    continue
                
                # Check for next push (nested function)
                if i > 2 and (hw & 0xFF00) == 0xB500:
                    func_size = i * 2
                    found_end = True
                    break
            
            if found_end and func_size >= 4:
                functions.append((addr, func_size))
                addr += func_size
                continue
        
        addr += 2
    
    return functions

def main():
    print("=" * 60)
    print("GBA ROM Binary Accuracy Matcher")
    print("=" * 60)
    
    rom = read_rom()
    print(f"ROM size: {len(rom)} bytes ({len(rom)/1024:.1f} KB)")
    
    # Read function list if available
    func_list_file = "function_list.txt"
    functions = []
    
    if os.path.exists(func_list_file):
        print(f"Reading function list from {func_list_file}")
        with open(func_list_file) as f:
            for line in f:
                line = line.strip()
                if not line or line.startswith('#'):
                    continue
                parts = line.split()
                if len(parts) >= 2:
                    addr = int(parts[0], 16)
                    size = int(parts[1])
                    name = parts[2] if len(parts) > 2 else f"func_{addr:08X}"
                    functions.append((addr, size, name))
    else:
        print("Scanning ROM for functions...")
        # Focus on the code section (0xC0 to ~0x1EE76)
        raw_funcs = scan_functions(rom, 0x080000C0, 0x0801F30C)
        functions = [(a, s, f"func_{a:08X}") for a, s in raw_funcs]
        print(f"Found {len(functions)} potential functions")
        
        # Save function list for future use
        with open(func_list_file, 'w') as f:
            for addr, size, name in functions:
                f.write(f"{addr:08X} {size:4d} {name}\n")
        print(f"Saved function list to {func_list_file}")
    
    print(f"\nProcessing {len(functions)} functions...")
    
    # Statistics
    total_bytes = 0
    matched_bytes = 0
    matched_funcs = 0
    partial_funcs = 0
    
    results = []
    
    for i, (addr, size, name) in enumerate(functions):
        orig_bytes = extract_function_bytes(rom, addr, size)
        total_bytes += size
        
        # Try to match with simple C patterns
        c_source = f"""
/* Function at 0x{addr:08X}, size {size} */
void {name}(void) {{
    /* placeholder */
}}
"""
        binary, orig, status = try_compile_and_match(c_source, rom, addr, size)
        
        if status == "MATCH":
            matched_bytes += size
            matched_funcs += 1
            results.append((addr, size, name, "MATCH"))
            print(f"  [{i+1:4d}/{len(functions)}] 0x{addr:08X} ({size:3d}B) {name}: ✅ MATCH")
        elif "PARTIAL" in status:
            partial_funcs += 1
            rate = match_rate(binary, orig) if binary and orig else 0
            matched_bytes += int(size * rate)
            results.append((addr, size, name, status))
            print(f"  [{i+1:4d}/{len(functions)}] 0x{addr:08X} ({size:3d}B) {name}: {status}")
        else:
            results.append((addr, size, name, status))
            # Only print failures for now to reduce noise
    
    # Summary
    print("\n" + "=" * 60)
    print("RESULTS SUMMARY")
    print("=" * 60)
    print(f"Total functions: {len(functions)}")
    print(f"Matched functions: {matched_funcs} ({matched_funcs/len(functions)*100:.1f}%)")
    print(f"Partial functions: {partial_funcs}")
    print(f"Total bytes: {total_bytes}")
    print(f"Matched bytes: {matched_bytes} ({matched_bytes/total_bytes*100:.1f}%)")
    print(f"Overall accuracy: {matched_bytes/total_bytes*100:.1f}%")
    
    # Save results
    with open("build/matching_results.txt", 'w') as f:
        f.write("Address    Size  Name                     Status\n")
        f.write("-" * 60 + "\n")
        for addr, size, name, status in results:
            f.write(f"0x{addr:08X} {size:4d}  {name:24s} {status}\n")
    
    print(f"\nDetailed results saved to build/matching_results.txt")

if __name__ == "__main__":
    os.makedirs("build", exist_ok=True)
    main()