#!/usr/bin/env python3
"""
GBA Binary Accuracy Pipeline
Strategy: 
1. Extract each function's original bytes from the ROM
2. Try to compile matching C with agbcc → assemble with gcc → compare bytes
3. For functions where C doesn't match, try hand-written matching assembly
4. Report byte-level accuracy across the entire ROM
"""
import os
import re
import struct
import subprocess
import sys
import tempfile
import hashlib

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 rom_offset(addr):
    return addr - ROM_BASE

def compile_c_to_asm(c_code):
    with tempfile.NamedTemporaryFile(suffix=".c", mode="w", delete=False) as f:
        f.write(c_code)
        c_path = f.name
    try:
        result = subprocess.run(
            [AGBCC, "-O2", c_path],
            capture_output=True, text=True, timeout=30
        )
        s_path = c_path.replace(".c", ".s")
        if os.path.exists(s_path):
            with open(s_path) as f:
                asm = f.read()
            os.unlink(s_path)
            return asm
        return None
    except:
        return None
    finally:
        if os.path.exists(c_path):
            os.unlink(c_path)

def assemble_with_gcc(asm_text, func_name="test_func"):
    """Assemble using gcc -c (which supports Thumb16)"""
    with tempfile.NamedTemporaryFile(suffix=".s", mode="w", delete=False) as f:
        f.write(asm_text)
        s_path = f.name
    
    o_path = s_path.replace(".s", ".o")
    bin_path = s_path.replace(".s", ".bin")
    
    try:
        result = subprocess.run(
            [GCC, "-mthumb", "-mcpu=arm7tdmi", "-c", s_path, "-o", o_path],
            capture_output=True, text=True, timeout=30
        )
        if result.returncode != 0:
            return None
        
        result = subprocess.run(
            [OBJCOPY, "-O", "binary", "-j", ".text", o_path, bin_path],
            capture_output=True, text=True, timeout=30
        )
        if result.returncode != 0:
            return None
        
        with open(bin_path, "rb") as f:
            return f.read()
    except:
        return None
    finally:
        for p in [s_path, o_path, bin_path]:
            if os.path.exists(p):
                os.unlink(p)

def extract_function_from_agbcc_asm(asm_text, func_name):
    """Extract just the function body from agbcc assembly output"""
    lines = asm_text.split('\n')
    body_lines = []
    in_func = False
    
    for line in lines:
        stripped = line.strip()
        # Start at function label
        if stripped.startswith(f"{func_name}:"):
            in_func = True
            continue
        # End at size directive
        if in_func and (stripped.startswith('.Lfe') or stripped.startswith('.size')):
            break
        if in_func and stripped:
            # Keep only actual instructions, skip assembler directives
            if not stripped.startswith('.') and not stripped.startswith('@'):
                body_lines.append(line)
    
    return '\n'.join(body_lines)

def match_bytes(our, orig):
    """Check exact byte match"""
    if len(our) != len(orig):
        return False, 0.0
    matching = sum(1 for a, b in zip(our, orig) if a == b)
    return matching == len(orig), matching / len(orig) if orig else 1.0

def main():
    os.makedirs("build", exist_ok=True)
    
    rom = read_rom()
    
    # Read function list from Ghidra
    with open("function_list_full.txt") as f:
        func_list = []
        for line in f:
            line = line.strip()
            if not line:
                continue
            parts = line.split(None, 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))
    
    # Read decompiled C source and extract function bodies
    with open("src/orbital_complete.c") as f:
        c_source = f.read()
    
    # Parse function bodies from orbital_complete.c
    func_c_bodies = {}
    pattern = r'/\*\s*----\s+(\w+)\s+@\s+(0x[0-9A-Fa-f]+)\s+----\s*\*/(.*?)(?=/\*\s*----|\Z)'
    for name, addr, body in re.findall(pattern, c_source, re.DOTALL):
        func_c_bodies[name] = body.strip()
    
    print(f"Functions in list: {len(func_list)}")
    print(f"C functions extracted: {len(func_c_bodies)}")
    
    # Build the matching binary
    total_bytes = sum(size for _, size, _ in func_list)
    matched_exact = 0
    matched_partial = 0
    unmatched = 0
    
    # Track byte-level accuracy across the ENTIRE code section
    # Code section: 0xC0 to ~0x1F30C
    code_start = 0xC0
    code_end = 0x1F30C
    code_size = code_end - code_start
    
    # Create a byte-level match map
    byte_match = bytearray(code_size)  # 0=untested, 1=match, 2=mismatch
    
    results = []
    
    for addr, size, name in func_list:
        if addr < ROM_BASE + code_start or addr >= ROM_BASE + code_end:
            continue
            
        offset = rom_offset(addr)
        orig_bytes = rom[offset:offset+size]
        
        if name in func_c_bodies:
            c_body = func_c_bodies[name]
            
            # Extract only the function implementation (skip comments and type decls)
            # Look for the actual function body starting with {
            func_match = re.search(r'\{(.*)\}', c_body, re.DOTALL)
            if func_match:
                func_body = '{' + func_match.group(1) + '}'
                
                # Build a compilable C file
                c_code = f"""#include <stdint.h>
#include <stddef.h>
#define CONCAT44(a,b) (((uint64_t)(a)<<32)|(uint32_t)(b))
#define CONCAT14(a,b) ((uint32_t)(b)<<8|(uint8_t)(a))
#define INIT_POS_X 0x78000
#define INIT_POS_Y 0x50000
{c_body}
"""
                
                # Try to compile with agbcc
                asm = compile_c_to_asm(c_code)
                
                if asm:
                    # Extract function body
                    func_asm = extract_function_from_agbcc_asm(asm, name)
                    
                    if func_asm.strip():
                        # Assemble with gcc
                        binary = assemble_with_gcc(func_asm, name)
                        
                        if binary:
                            exact, rate = match_bytes(binary, orig_bytes)
                            
                            if exact:
                                matched_exact += size
                                results.append((addr, size, name, "EXACT"))
                                for i in range(size):
                                    byte_match[offset + i - code_start] = 1
                            elif rate > 0.8:
                                matched_partial += int(size * rate)
                                results.append((addr, size, name, f"PARTIAL {rate*100:.0f}%"))
                                for i in range(min(len(binary), size)):
                                    if binary[i] == orig_bytes[i]:
                                        byte_match[offset + i - code_start] = 1
                            else:
                                unmatched += size
                                results.append((addr, size, name, f"MISMATCH {rate*100:.0f}%"))
                        else:
                            unmatched += size
                            results.append((addr, size, name, "ASM_FAIL"))
                    else:
                        unmatched += size
                        results.append((addr, size, name, "NO_BODY"))
                else:
                    unmatched += size
                    results.append((addr, size, name, "COMPILE_FAIL"))
            else:
                unmatched += size
                results.append((addr, size, name, "PARSE_FAIL"))
        else:
            unmatched += size
            results.append((addr, size, name, "NO_C_SOURCE"))
    
    # Count matched bytes from byte_match
    bytes_matched = sum(1 for b in byte_match if b == 1)
    bytes_tested = sum(1 for b in byte_match if b != 0)
    
    # Summary
    print("\n" + "=" * 70)
    print("BINARY ACCURACY RESULTS")
    print("=" * 70)
    print(f"Total functions: {len(func_list)}")
    print(f"  EXACT match:   {sum(1 for _,_,_,s in results if s == 'EXACT'):4d}")
    print(f"  PARTIAL:       {sum(1 for _,_,_,s in results if 'PARTIAL' in s):4d}")
    print(f"  MISMATCH:      {sum(1 for _,_,_,s in results if 'MISMATCH' in s):4d}")
    print(f"  FAIL/NO SRC:   {sum(1 for _,_,_,s in results if s not in ('EXACT',) and 'PARTIAL' not in s and 'MISMATCH' not in s):4d}")
    print(f"\nCode section: {code_size} bytes (0xC0 - 0x1F30C)")
    print(f"Functions cover: {sum(size for _,size,_ in func_list if ROM_BASE + code_start <= addr < ROM_BASE + code_end):d} bytes")
    print(f"Bytes matched:   {bytes_matched}/{bytes_tested} tested, {bytes_matched}/{code_size} total")
    print(f"\nOverall binary accuracy: {bytes_matched/code_size*100:.1f}%")
    
    # Save results
    with open("build/binary_accuracy.txt", "w") as f:
        f.write(f"Binary Accuracy: {bytes_matched/code_size*100:.1f}%\n")
        f.write(f"Matched bytes: {bytes_matched}/{code_size}\n\n")
        for addr, size, name, status in results:
            f.write(f"0x{addr:08X} {size:4d} {name:40s} {status}\n")
    
    print(f"\nDetailed results: build/binary_accuracy.txt")

if __name__ == "__main__":
    main()