#!/usr/bin/env python3
"""
GBA Decompilation Matcher v3
Systematically matches C code against original ROM binary.
Strategy: Analyze each function's assembly, find matching C pattern.
"""
import os
import re
import struct
import subprocess
import sys
import tempfile

ROM_PATH = "bit Generations - Orbital (Japan) (En).gba"
ROM_BASE = 0x08000000
AGBCC = "agbcc/agbcc"
GCC = "arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi/bin/arm-none-eabi-gcc"
OBJCOPY = "arm-gnu-toolchain-13.3.rel1-darwin-arm64-arm-none-eabi/bin/arm-none-eabi-objcopy"
CODE_START = 0x080000C0
CODE_END = 0x0801F30C

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

def disasm_thumb(func_bytes, base_addr):
    """Disassemble Thumb16 code"""
    lines = []
    i = 0
    while i < len(func_bytes) - 1:
        hw = struct.unpack_from('<H', func_bytes, i)[0]
        addr = base_addr + i
        text = f".hword 0x{hw:04X}"
        
        if hw == 0x4770: text = "bx lr"
        elif hw == 0x46C0: text = "nop"
        elif (hw & 0xFF00) == 0xB500:
            regs = ["lr"]
            for b in range(7, -1, -1):
                if hw & (1 << b): regs.insert(0, f"r{b}")
            text = f"push {{{', '.join(regs)}}}"
        elif (hw & 0xFF00) == 0xBD00:
            regs = ["pc"]
            for b in range(7, -1, -1):
                if hw & (1 << b): regs.insert(0, f"r{b}")
            text = f"pop {{{', '.join(regs)}}}"
        elif (hw & 0xF800) == 0x2000:
            text = f"movs r{hw>>8&7}, #{hw&0xFF}"
        elif (hw & 0xF800) == 0x3000:
            text = f"adds r{hw>>8&7}, #{hw&0xFF}"
        elif (hw & 0xF800) == 0x3800:
            text = f"subs r{hw>>8&7}, #{hw&0xFF}"
        elif (hw & 0xF800) == 0x2800:
            text = f"cmp r{hw>>8&7}, #{hw&0xFF}"
        elif (hw & 0xF800) == 0x6000:
            text = f"str r{hw&7}, [r{(hw>>3)&7}, #{((hw>>6)&0x1f)*4}]"
        elif (hw & 0xF800) == 0x6800:
            text = f"ldr r{hw&7}, [r{(hw>>3)&7}, #{((hw>>6)&0x1f)*4}]"
        elif (hw & 0xFFC0) == 0x4280:
            text = f"cmp r{hw&7}, r{(hw>>3)&7}"
        elif (hw & 0xFFC0) == 0x1800:
            text = f"adds r{hw&7}, r{(hw>>3)&7}, r{(hw>>6)&7}"
        elif (hw & 0xFFC0) == 0x1A00:
            text = f"subs r{hw&7}, r{(hw>>3)&7}, r{(hw>>6)&7}"
        elif (hw & 0xFFC0) == 0x4300:
            text = f"orrs r{hw&7}, r{(hw>>3)&7}"
        elif (hw & 0xFFC0) == 0x4040:
            text = f"eors r{hw&7}, r{(hw>>3)&7}"
        
        lines.append((addr, hw, text))
        i += 2
    return lines

def try_compile_match(c_code, addr, size, rom):
    """Try to compile C code and match bytes"""
    orig = rom[addr - ROM_BASE:addr - ROM_BASE + size]
    
    with tempfile.NamedTemporaryFile(suffix=".c", mode="w", delete=False) as f:
        f.write(c_code)
        c_path = f.name
    s_path = c_path.replace(".c", ".s")
    o_path = c_path.replace(".c", ".o")
    bin_path = c_path.replace(".c", ".bin")
    
    try:
        subprocess.run([AGBCC, "-O2", c_path], capture_output=True, timeout=30)
        if not os.path.exists(s_path):
            return False
        
        subprocess.run([GCC, "-mthumb", "-mcpu=arm7tdmi", "-c", s_path, "-o", o_path],
                      capture_output=True, timeout=30)
        if not os.path.exists(o_path):
            return False
        
        subprocess.run([OBJCOPY, "-O", "binary", "-j", ".text", o_path, bin_path],
                      capture_output=True, timeout=30)
        if not os.path.exists(bin_path):
            return False
        
        with open(bin_path, "rb") as f:
            binary = f.read()
        
        binary = binary[:size]
        return binary == orig
    except:
        return False
    finally:
        for p in [c_path, s_path, o_path, bin_path]:
            try: os.unlink(p)
            except: pass

def analyze_and_match(func_bytes, addr, size, rom):
    """Analyze function assembly and try to find matching C code"""
    
    # Pattern 1: Empty function (bx lr)
    if size == 2 and func_bytes[0] == 0x70 and func_bytes[1] == 0x47:
        c = "void func(void) {}"
        if try_compile_match(c, addr, size, rom):
            return c, "empty"
    
    # Pattern 2: Return constant (movs r0, #N; bx lr)
    if size == 4:
        hw1 = struct.unpack_from('<H', func_bytes, 0)[0]
        hw2 = struct.unpack_from('<H', func_bytes, 2)[0]
        if (hw1 & 0xF800) == 0x2000 and (hw1 & 0x700) == 0 and hw2 == 0x4770:
            val = hw1 & 0xFF
            c = f"#include <stdint.h>\nint func(void) {{ return {val}; }}"
            if try_compile_match(c, addr, size, rom):
                return c, f"return_{val}"
    
    # Pattern 3: SWI wrapper (swi #N; bx lr)
    if size == 4:
        hw1 = struct.unpack_from('<H', func_bytes, 0)[0]
        hw2 = struct.unpack_from('<H', func_bytes, 2)[0]
        if (hw1 & 0xFF00) == 0xDF00 and hw2 == 0x4770:
            swi = hw1 & 0xFF
            c = f"#include <stdint.h>\nvoid func(void) {{ __asm__ volatile(\"swi #0x{swi:02X}\"); }}"
            if try_compile_match(c, addr, size, rom):
                return c, f"swi_{swi:02X}"
    
    # Pattern 4: Struct clear - movs r1, #0; str r1, [r0, #0]; ... str r1, [r0, #N]; bx lr
    # Count str instructions
    if size >= 6:
        lines = disasm_thumb(func_bytes, addr)
        str_count = 0
        all_str = True
        all_zero = True
        has_movs_r1_0 = False
        for _, hw, text in lines:
            if "movs r1, #0" in text:
                has_movs_r1_0 = True
            elif "str r1, [r0," in text:
                str_count += 1
            elif text == "bx lr":
                pass
            else:
                all_str = False
                break
        
        if has_movs_r1_0 and str_count >= 2 and all_str and all_zero:
            # Try struct clear with N members
            for n in range(str_count, str_count + 4):
                members = "; ".join(f"int m{i}" for i in range(n))
                assigns = "; ".join(f"p->m{i} = 0" for i in range(n))
                c = f"#include <stdint.h>\nstruct S {{ {members}; }};\nvoid func(struct S *p) {{ {assigns}; }}"
                if try_compile_match(c, addr, size, rom):
                    return c, f"struct_clear_{n}"
    
    # Pattern 5: Push/pop with simple body
    # push {lr}; ...; pop {pc}
    if size >= 4:
        hw1 = struct.unpack_from('<H', func_bytes, 0)[0]
        if (hw1 & 0xFF00) == 0xB500:
            # Has push {lr} - try simple function patterns
            pass
    
    return None, None

def main():
    os.makedirs("build/matching", exist_ok=True)
    
    rom = read_rom()
    
    # Read function list
    with open("function_list_v2.txt") as f:
        func_list = []
        for line in f:
            parts = line.strip().split(None, 2)
            if len(parts) >= 2:
                addr = int(parts[0], 16)
                size = int(parts[1])
                name = parts[2] if len(parts) > 2 else ""
                func_list.append((addr, size, name))
    
    # Filter to code section
    code_funcs = [(a, s, n) for a, s, n in func_list if CODE_START <= a < CODE_END]
    code_funcs.sort(key=lambda x: x[0])
    
    print(f"Functions in code section: {len(code_funcs)}")
    
    matched_count = 0
    matched_bytes = 0
    total_tested = 0
    
    for addr, size, name in code_funcs:
        if size < 2 or size > 100:
            continue
        
        offset = addr - ROM_BASE
        func_bytes = rom[offset:offset+size]
        
        c_code, pattern = analyze_and_match(func_bytes, addr, size, rom)
        
        if c_code:
            matched_count += 1
            matched_bytes += size
            print(f"  ✅ {name:40s} 0x{addr:08X} {size:4d}B  ({pattern})")
        
        total_tested += 1
    
    print(f"\n{'='*60}")
    print(f"RESULTS:")
    print(f"  Functions tested: {total_tested}")
    print(f"  Functions matched: {matched_count}")
    print(f"  Bytes matched: {matched_bytes}")
    print(f"  Code section: {CODE_END - CODE_START} bytes")
    print(f"  Accuracy: {matched_bytes * 100 / (CODE_END - CODE_START):.1f}%")

if __name__ == "__main__":
    main()