#!/usr/bin/env python3
"""
GBA Function Matcher v2 - Systematically matches C code against original ROM.
Each function is analyzed, C code is written, compiled with agbcc, and compared.
"""
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"

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

def disasm_thumb16(rom_bytes, base_addr):
    """Disassemble Thumb16 code into readable assembly"""
    lines = []
    i = 0
    while i < len(rom_bytes) - 1:
        hw = struct.unpack_from('<H', rom_bytes, i)[0]
        addr = base_addr + i
        
        decoded = None
        
        # PUSH
        if (hw & 0xFF00) == 0xB500:
            regs = []
            for bit in range(8):
                if hw & (1 << bit):
                    regs.append(f"r{bit}")
            regs.append("lr")
            decoded = f"push {{{', '.join(regs)}}}"
        # POP
        elif (hw & 0xFF00) == 0xBD00:
            regs = []
            for bit in range(8):
                if hw & (1 << bit):
                    regs.append(f"r{bit}")
            regs.append("pc")
            decoded = f"pop {{{', '.join(regs)}}}"
        # MOV Rd, #imm8
        elif (hw & 0xF800) == 0x2000:
            rd = (hw >> 8) & 7
            imm = hw & 0xFF
            decoded = f"movs r{rd}, #{imm}"
        # ADDS Rd, #imm8
        elif (hw & 0xF800) == 0x3000:
            rd = (hw >> 8) & 7
            imm = hw & 0xFF
            decoded = f"adds r{rd}, #{imm}"
        # SUBS Rd, #imm8
        elif (hw & 0xF800) == 0x3800:
            rd = (hw >> 8) & 7
            imm = hw & 0xFF
            decoded = f"subs r{rd}, #{imm}"
        # CMP Rd, #imm8
        elif (hw & 0xF800) == 0x2800:
            rd = (hw >> 8) & 7
            imm = hw & 0xFF
            decoded = f"cmp r{rd}, #{imm}"
        # STR Rd, [Rn, #imm5*4]
        elif (hw & 0xF800) == 0x6000:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = ((hw >> 6) & 0x1F) * 4
            decoded = f"str r{rd}, [r{rn}, #{imm}]"
        # LDR Rd, [Rn, #imm5*4]
        elif (hw & 0xF800) == 0x6800:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = ((hw >> 6) & 0x1F) * 4
            decoded = f"ldr r{rd}, [r{rn}, #{imm}]"
        # STRB Rd, [Rn, #imm5]
        elif (hw & 0xF800) == 0x7000:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = (hw >> 6) & 0x1F
            decoded = f"strb r{rd}, [r{rn}, #{imm}]"
        # LDRB Rd, [Rn, #imm5]
        elif (hw & 0xF800) == 0x7800:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = (hw >> 6) & 0x1F
            decoded = f"ldrb r{rd}, [r{rn}, #{imm}]"
        # STRH Rd, [Rn, #imm5*2]
        elif (hw & 0xF800) == 0x8000:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = ((hw >> 6) & 0x1F) * 2
            decoded = f"strh r{rd}, [r{rn}, #{imm}]"
        # LDRH Rd, [Rn, #imm5*2]
        elif (hw & 0xF800) == 0x8800:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = ((hw >> 6) & 0x1F) * 2
            decoded = f"ldrh r{rd}, [r{rn}, #{imm}]"
        # STR Rd, [Rn, Rs]
        elif (hw & 0xFFC0) == 0x5000:
            rd = hw & 7
            rn = (hw >> 3) & 7
            rs = (hw >> 6) & 7
            decoded = f"str r{rd}, [r{rn}, r{rs}]"
        # LDR Rd, [Rn, Rs]
        elif (hw & 0xFFC0) == 0x5800:
            rd = hw & 7
            rn = (hw >> 3) & 7
            rs = (hw >> 6) & 7
            decoded = f"ldr r{rd}, [r{rn}, r{rs}]"
        # ADD Rd, Rn, Rs (low registers)
        elif (hw & 0xFFC0) == 0x1800:
            rd = hw & 7
            rn = (hw >> 3) & 7
            rs = (hw >> 6) & 7
            decoded = f"adds r{rd}, r{rn}, r{rs}"
        # SUB Rd, Rn, Rs (low registers)
        elif (hw & 0xFFC0) == 0x1A00:
            rd = hw & 7
            rn = (hw >> 3) & 7
            rs = (hw >> 6) & 7
            decoded = f"subs r{rd}, r{rn}, r{rs}"
        # ADD Rd, Rn, #imm3
        elif (hw & 0xFFC0) == 0x1C00:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = (hw >> 6) & 7
            decoded = f"adds r{rd}, r{rn}, #{imm}"
        # SUB Rd, Rn, #imm3
        elif (hw & 0xFFC0) == 0x1E00:
            rd = hw & 7
            rn = (hw >> 3) & 7
            imm = (hw >> 6) & 7
            decoded = f"subs r{rd}, r{rn}, #{imm}"
        # CMP Rn, Rs
        elif (hw & 0xFFC0) == 0x4280:
            rn = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"cmp r{rn}, r{rs}"
        # ANDS Rd, Rs (TST)
        elif (hw & 0xFFC0) == 0x4200:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"tst r{rd}, r{rs}"
        # EORS Rd, Rs
        elif (hw & 0xFFC0) == 0x4040:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"eors r{rd}, r{rs}"
        # ORRS Rd, Rs
        elif (hw & 0xFFC0) == 0x4300:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"orrs r{rd}, r{rs}"
        # BICS Rd, Rs
        elif (hw & 0xFFC0) == 0x4380:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"bics r{rd}, r{rs}"
        # MVNS Rd, Rs
        elif (hw & 0xFFC0) == 0x43C0:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"mvns r{rd}, r{rs}"
        # MULS Rd, Rs
        elif (hw & 0xFFC0) == 0x4340:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"muls r{rd}, r{rs}"
        # LSL Rd, Rs
        elif (hw & 0xFFC0) == 0x4080:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"lsls r{rd}, r{rs}"
        # LSR Rd, Rs
        elif (hw & 0xFFC0) == 0x40C0:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"lsrs r{rd}, r{rs}"
        # ASR Rd, Rs
        elif (hw & 0xFFC0) == 0x4100:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"asrs r{rd}, r{rs}"
        # ADCS Rd, Rs
        elif (hw & 0xFFC0) == 0x4140:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"adcs r{rd}, r{rs}"
        # SBCS Rd, Rs
        elif (hw & 0xFFC0) == 0x4180:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"sbcs r{rd}, r{rs}"
        # ROR Rd, Rs
        elif (hw & 0xFFC0) == 0x41C0:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"rors r{rd}, r{rs}"
        # CMN Rd, Rs
        elif (hw & 0xFFC0) == 0x42C0:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"cmn r{rd}, r{rs}"
        # NEG Rd, Rs (RSBS Rd, Rs, #0)
        elif (hw & 0xFFC0) == 0x4240:
            rd = hw & 7
            rs = (hw >> 3) & 7
            decoded = f"negs r{rd}, r{rs}"
        # MOV Rd, Rs (high registers)
        elif (hw & 0xFFC0) == 0x4600:
            rd = ((hw >> 4) & 8) | (hw & 7)
            rs = (hw >> 3) & 0xF
            decoded = f"mov r{rd}, r{rs}"
        # ADD Rd, Rs (high registers)  
        elif (hw & 0xFFC0) == 0x4400:
            rd = ((hw >> 4) & 8) | (hw & 7)
            rs = (hw >> 3) & 0xF
            decoded = f"add r{rd}, r{rs}"
        # BX Rs
        elif hw == 0x4770:
            decoded = "bx lr"
        # BLX Rs
        elif (hw & 0xFFC0) == 0x4780:
            rs = (hw >> 3) & 0xF
            decoded = f"blx r{rs}"
        # NOP (MOV R8, R8)
        elif hw == 0x46C0:
            decoded = "nop"
        # B<cond> offset
        elif (hw & 0xF000) == 0xD000:
            cond = (hw >> 8) & 0xF
            offset_val = hw & 0xFF
            if offset_val & 0x80:
                offset_val -= 256
            target = addr + 4 + offset_val * 2
            conds = ["eq","ne","cs","cc","mi","pl","vs","vc",
                    "hi","ls","ge","lt","gt","le","al"]
            c = conds[cond] if cond < len(conds) else f"x{cond}"
            decoded = f"b{c} 0x{target:08X}"
        # B offset (unconditional)
        elif (hw & 0xF800) == 0xE000:
            offset_val = hw & 0x7FF
            if offset_val & 0x400:
                offset_val -= 2048
            target = addr + 4 + offset_val * 2
            decoded = f"b 0x{target:08X}"
        # BL prefix (high)
        elif (hw & 0xF800) == 0xF000:
            decoded = f".hword 0x{hw:04X}  @ BL prefix"
        # BL suffix (low) 
        elif (hw & 0xF800) == 0xF800:
            decoded = f".hword 0x{hw:04X}  @ BL suffix"
        # LSL/LSR/ASR with immediate
        elif (hw & 0xF000) == 0x0000:
            rd = hw & 7
            rs = (hw >> 3) & 7
            imm = (hw >> 6) & 0x1F
            op = (hw >> 11) & 3
            ops = ["lsls", "lsrs", "asrs"]
            if imm == 0 and op < 2:
                decoded = f"{ops[op]} r{rd}, r{rs}"
            else:
                decoded = f"{ops[op]} r{rd}, r{rs}, #{imm}"
        # ADD Rd, PC, #imm8 (word-aligned)
        elif (hw & 0xFF00) == 0xA000:
            rd = (hw >> 8) & 7
            imm = (hw & 0xFF) * 4
            decoded = f"add r{rd}, pc, #{imm}"
        # ADD Rd, SP, #imm8 (word-aligned)
        elif (hw & 0xFF00) == 0xA800:
            rd = (hw >> 8) & 7
            imm = (hw & 0xFF) * 4
            decoded = f"add r{rd}, sp, #{imm}"
        # SUB SP, #imm7
        elif (hw & 0xFF00) == 0xB080:
            imm = (hw & 0x7F) * 4
            decoded = f"sub sp, #{imm}"
        # SWI #imm8
        elif (hw & 0xFF00) == 0xDF00:
            imm = hw & 0xFF
            decoded = f"swi #{imm}"
        # SXTH
        elif hw == 0xB200:
            decoded = "sxth r0, r0"
        # SXTB
        elif hw == 0xB240:
            decoded = "sxtb r0, r0"
        # UXTH
        elif hw == 0xB280:
            decoded = "uxth r0, r0"
        # UXTB
        elif hw == 0xB2C0:
            decoded = "uxtb r0, r0"
        
        if decoded is None:
            decoded = f".hword 0x{hw:04X}"
        
        lines.append((addr, hw, decoded))
        i += 2
    
    if len(rom_bytes) % 2 == 1:
        lines.append((base_addr + len(rom_bytes) - 1, rom_bytes[-1], f".byte 0x{rom_bytes[-1]:02X}"))
    
    return lines

def compile_and_compare(c_code, addr, size):
    """Compile C with agbcc, assemble with gcc, compare bytes"""
    rom = read_rom()
    offset = addr - ROM_BASE
    orig = rom[offset:offset+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 "COMPILE_FAIL", None, None
        
        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 "ASM_FAIL", None, None
        
        subprocess.run([OBJCOPY, "-O", "binary", "-j", ".text", o_path, bin_path],
                      capture_output=True, timeout=30)
        if not os.path.exists(bin_path):
            return "OBJCOPY_FAIL", None, None
        
        with open(bin_path, "rb") as f:
            binary = f.read()
        
        if len(binary) > size:
            binary = binary[:size]
        
        if binary == orig:
            return "MATCH", binary, orig
        else:
            matching = sum(1 for a, b in zip(binary, orig) if a == b)
            rate = matching / size if size > 0 else 0
            return f"RATE:{rate:.2f}", binary, orig
    except Exception as e:
        return f"ERROR:{e}", None, None
    finally:
        for p in [c_path, s_path, o_path, bin_path]:
            try: os.unlink(p)
            except: pass

def main():
    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 f"func_{addr:08X}"
                func_list.append((addr, size, name))
    
    # Focus on small functions (<=20 bytes) - more likely to match
    small_funcs = [(a, s, n) for a, s, n in func_list 
                   if 0x080000C0 <= a < 0x0801F30C and 4 <= s <= 20]
    
    print(f"Total functions: {len(func_list)}")
    print(f"Small functions (4-20 bytes): {len(small_funcs)}")
    
    # Disassemble and try to match first 20 small functions
    print("\n=== Analyzing small functions ===\n")
    
    matched_count = 0
    matched_bytes = 0
    total_tested = 0
    total_bytes = 0
    
    for addr, size, name in small_funcs[:50]:
        offset = addr - ROM_BASE
        func_bytes = rom[offset:offset+size]
        
        # Disassemble
        lines = disasm_thumb16(func_bytes, addr)
        
        # Print disassembly
        print(f"--- {name} @ 0x{addr:08X} ({size}B) ---")
        for a, hw, text in lines:
            print(f"  0x{a:08X}: {hw:04X}  {text}")
        
        total_tested += 1
        total_bytes += size
    
    print(f"\nAnalyzed {total_tested} functions, {total_bytes} bytes")

if __name__ == "__main__":
    main()