#!/usr/bin/env python3
"""
Extended batch matcher for Orbital.

Handles more function patterns: empty, return constant, SWI wrapper,
simple struct stores, clear struct, and basic framed functions.

Usage: python3 tools/batch_match_v2.py
"""

import json
import os
import re
import struct
import subprocess
import tempfile
from pathlib import Path

ROOT = Path(__file__).resolve().parent.parent
ROM_PATH = ROOT / "bit Generations - Orbital (Japan) (En).gba"
FUNC_LIST = ROOT / "function_list_v2.txt"
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
CODE_START = 0x080000C0
CODE_END = 0x0801F30C

def mask_bl(data):
    m = bytearray(data)
    i = 0
    while i < len(m)-3:
        h1 = int.from_bytes(m[i:i+2],'little')
        h2 = int.from_bytes(m[i+2:i+4],'little')
        if (h1>>11)==0b11110 and (h2>>11)==0b11111:
            h1&=0xF800; h2&=0xF800
            m[i]=h1&0xFF; m[i+1]=h1>>8; m[i+2]=h2&0xFF; m[i+3]=h2>>8
            i+=4; continue
        i+=2
    return m

def last_ret_offset(data):
    i = 0; last = -1
    while i < len(data)-1:
        hw = int.from_bytes(data[i:i+2],'little')
        if hw == 0x4770 or ((hw>>8)==0xBD and (hw&0x100)): last = i
        if (hw>>11) in (0b11101, 0b11110, 0b11111): i += 4
        else: i += 2
    return last

def try_compile(c_code, addr, size, rom):
    with tempfile.TemporaryDirectory() as td:
        c_path = f'{td}/f.c'
        s_path = f'{td}/f.s'
        o_path = f'{td}/f.o'
        b_path = f'{td}/f.bin'
        Path(c_path).write_text(c_code)
        subprocess.run([AGBCC,'-O2','-mthumb-interwork','-S',c_path,'-o',s_path],
                      capture_output=True,timeout=10)
        if not Path(s_path).exists(): return False, 0
        subprocess.run([AS,'-mcpu=arm7tdmi',s_path,'-o',o_path],
                      capture_output=True,timeout=10)
        if not Path(o_path).exists(): return False, 0
        subprocess.run([OBJCOPY,'-O','binary','-j','.text',o_path,b_path],
                      capture_output=True,timeout=10)
        if not Path(b_path).exists(): return False, 0
        compiled = Path(b_path).read_bytes()
        off = addr - ROM_BASE
        rb = rom[off:off+size]
        lr = last_ret_offset(compiled)
        code_end = (lr + 2) if lr >= 0 else len(compiled)
        lr_r = last_ret_offset(rb)
        rom_end = (lr_r + 2) if lr_r >= 0 else size
        n = min(code_end, rom_end, size)
        if n == 0: return False, 0
        if mask_bl(compiled[:n]) == mask_bl(rb[:n]):
            return True, n
        return False, n

def hw16(rom, addr):
    return struct.unpack_from('<H', rom, addr - ROM_BASE)[0]

def find_real_end(rom, addr, size):
    """Find the actual function end (last return instruction)."""
    off = addr - ROM_BASE
    i = 0; last = -1
    while i < size - 1:
        hw = struct.unpack_from('<H', rom, off + i)[0]
        if hw == 0x4770 or ((hw >> 8) == 0xBD and (hw & 0x100)):
            last = i
        i += 2
    return (last + 2) if last >= 0 else size

def decode_instr(rom, addr):
    """Decode a single Thumb16 instruction."""
    hw = hw16(rom, addr)
    if hw == 0x4770: return "bx lr"
    if (hw & 0xFF00) == 0xB500:
        regs = ["lr"]
        for b in range(7, -1, -1):
            if hw & (1 << b): regs.insert(0, f"r{b}")
        return f"push {{{','.join(regs)}}}"
    if (hw & 0xFF00) == 0xBD00:
        regs = ["pc"]
        for b in range(7, -1, -1):
            if hw & (1 << b): regs.insert(0, f"r{b}")
        return f"pop {{{','.join(regs)}}}"
    if (hw & 0xF800) == 0x2000: return f"movs r{hw>>8&7}, #{hw&0xFF}"
    if (hw & 0xF800) == 0x6000: return f"str r{hw&7}, [r{(hw>>3)&7}, #{((hw>>6)&0x1F)*4}]"
    if (hw & 0xF800) == 0x6800: return f"ldr r{hw&7}, [r{(hw>>3)&7}, #{((hw>>6)&0x1F)*4}]"
    if (hw & 0xF800) == 0x3000: return f"adds r{hw>>8&7}, #{hw&0xFF}"
    if (hw & 0xF800) == 0x3800: return f"subs r{hw>>8&7}, #{hw&0xFF}"
    if (hw & 0xFFC0) == 0x1C00: return f"adds r{hw&7}, r{(hw>>3)&7}, #{(hw>>6)&7}"
    if (hw & 0xFFC0) == 0x1E00: return f"subs r{hw&7}, r{(hw>>3)&7}, #{(hw>>6)&7}"
    if (hw & 0xFFC0) == 0x4280: return f"cmp r{hw&7}, r{(hw>>3)&7}"
    if (hw & 0xF800) == 0x2800: return f"cmp r{hw>>8&7}, #{hw&0xFF}"
    if (hw & 0xFFC0) == 0x4300: return f"orrs r{hw&7}, r{(hw>>3)&7}"
    if (hw & 0xFFC0) == 0x4000: return f"ands r{hw&7}, r{(hw>>3)&7}"
    if (hw & 0xFFC0) == 0x1800: return f"adds r{hw&7}, r{(hw>>3)&7}, r{(hw>>6)&7}"
    if (hw & 0xFFC0) == 0x1A00: return f"subs r{hw&7}, r{(hw>>3)&7}, r{(hw>>6)&7}"
    if (hw & 0xFFC0) == 0x4600:
        dst = (hw & 7) | ((hw >> 4) & 8); src = (hw >> 3) & 0xF
        return f"mov r{dst}, r{src}"
    if (hw >> 11) == 0x9: return f"ldr r{hw>>8&7}, [pc, #{(hw&0xFF)*4}]"
    if (hw & 0xFF00) == 0xDF00: return f"swi #{hw&0xFF}"
    if (hw & 0xFFC0) == 0x4040: return f"eors r{hw&7}, r{(hw>>3)&7}"
    if (hw >> 11) == 0b11110: return "bl"
    if (hw & 0xF800) == 0xE000: return "b"
    if (hw & 0xD000) == 0xD000 and (hw & 0xF000) == 0xD000: return f"bcond"
    return f".short 0x{hw:04X}"

def generate_c(rom, addr, size, name):
    """Generate C code candidates for a function."""
    candidates = []
    real_end = find_real_end(rom, addr, size)
    if real_end < 4: real_end = size
    
    # Decode instructions up to real_end
    instrs = []
    i = 0
    while i < real_end - 1:
        decoded = decode_instr(rom, addr + i)
        instrs.append(decoded)
        if decoded == "bl": i += 4
        else: i += 2
    
    h1 = hw16(rom, addr)
    
    # Pattern: empty function
    if real_end == 2 and h1 == 0x4770:
        candidates.append("void f(void) {}\n")
        return candidates
    
    # Pattern: return constant
    if real_end == 4 and (h1 & 0xF800) == 0x2000:
        h2 = hw16(rom, addr + 2)
        if h2 == 0x4770 and (h1 & 0x700) == 0:
            val = h1 & 0xFF
            candidates.append(f"int f(void) {{ return {val}; }}\n")
            return candidates
    
    # Pattern: SWI wrapper
    if real_end == 4 and (h1 & 0xFF00) == 0xDF00:
        h2 = hw16(rom, addr + 2)
        if h2 == 0x4770:
            swi = h1 & 0xFF
            candidates.append(f"void f(void) {{ __asm__ volatile(\"swi #0x{swi:02X}\"); }}\n")
            return candidates
    
    # Pattern: movs r0, #val; bx lr; (padding bx lr)
    if real_end >= 4 and (h1 & 0xF800) == 0x2000 and (h1 & 0x700) == 0:
        h2 = hw16(rom, addr + 2)
        if h2 == 0x4770:
            val = h1 & 0xFF
            candidates.append(f"int f(void) {{ return {val}; }}\n")
            return candidates
    
    # Pattern: str rN, [r0, #off]; ...; bx lr (all stores to r0)
    if real_end >= 6:
        stores = []
        valid = True
        for j in range(0, real_end - 2, 2):
            hh = hw16(rom, addr + j)
            if (hh & 0xF800) == 0x6000 and ((hh >> 3) & 7) == 0:
                rd = hh & 7
                offv = ((hh >> 6) & 0x1F) * 4
                stores.append((rd, offv))
            else:
                valid = False
                break
        if valid and stores and len(stores) <= 8:
            max_off = max(s[1] for s in stores)
            fields = [f"    int f{o};" for o in range(0, max_off + 4, 4)]
            struct_def = "struct S {\n" + "\n".join(fields) + "\n};\n"
            params = {}
            assigns = []
            for rd, offv in stores:
                if rd not in params: params[rd] = f"p{rd}"
                assigns.append(f"    p->f{offv} = {params[rd]};")
            param_list = ", ".join(f"int {v}" for _, v in sorted(params.items()))
            func = f"void f(struct S *p{', ' + param_list if param_list else ''}) {{\n"
            func += "\n".join(assigns) + "\n}\n"
            candidates.append(struct_def + "\n" + func)
            return candidates
    
    # Pattern: movs rN, #0; str rN, [r0, #off]; ...; bx lr (clear struct)
    if real_end >= 6 and (h1 & 0xF800) == 0x2000 and (h1 & 0xFF) == 0:
        zero_reg = (h1 >> 8) & 7
        stores = []
        valid = True
        for j in range(2, real_end - 2, 2):
            hh = hw16(rom, addr + j)
            if (hh & 0xF800) == 0x6000 and (hh & 7) == zero_reg and ((hh >> 3) & 7) == 0:
                offv = ((hh >> 6) & 0x1F) * 4
                stores.append(offv)
            else:
                valid = False
                break
        if valid and stores and len(stores) <= 12:
            max_off = max(stores)
            fields = [f"    int f{o};" for o in range(0, max_off + 4, 4)]
            struct_def = "struct S {\n" + "\n".join(fields) + "\n};\n"
            assigns = [f"    p->f{offv} = 0;" for offv in stores]
            func = "void f(struct S *p) {\n" + "\n".join(assigns) + "\n}\n"
            candidates.append(struct_def + "\n" + func)
            return candidates
    
    return candidates

def main():
    rom = ROM_PATH.read_bytes()
    funcs = []
    for line in open(FUNC_LIST):
        parts = line.strip().split(None, 2)
        if len(parts) >= 2:
            a = int(parts[0], 16); s = int(parts[1])
            n = parts[2] if len(parts) > 2 else f"sub_{a:X}"
            funcs.append((a, s, n))
    funcs.sort()
    
    matched = []
    total_bytes = 0
    
    for addr, size, name in funcs:
        if not (CODE_START <= addr < CODE_END) or size < 2:
            continue
        total_bytes += size
        
        candidates = generate_c(rom, addr, size, name)
        
        for c_code in candidates:
            ok, n = try_compile(c_code, addr, size, rom)
            if ok:
                matched.append((addr, size, name, n, c_code))
                break
    
    matched_bytes = sum(n for _, _, _, n, _ in matched)
    
    MATCHED_DIR = ROOT / "src" / "matched"
    MATCHED_DIR.mkdir(parents=True, exist_ok=True)
    for addr, size, name, n, c_code in matched:
        (MATCHED_DIR / f"{name}.c").write_text(c_code)
    
    report = {
        "matched_functions": len(matched),
        "matched_bytes": matched_bytes,
        "total_bytes": total_bytes,
        "accuracy_pct": round(100 * matched_bytes / total_bytes, 1) if total_bytes else 0,
    }
    (ROOT / "build" / "batch_match_report.json").write_text(json.dumps(report, indent=2))
    
    print(f"\n{'='*60}")
    print(f"MATCHED: {len(matched)} functions, {matched_bytes} bytes")
    print(f"TOTAL:   {total_bytes} bytes")
    print(f"ACCURACY: {100*matched_bytes/total_bytes:.1f}%")
    print(f"\nMatched by size:")
    for addr, size, name, n, _ in sorted(matched, key=lambda x: -x[3]):
        print(f"  0x{addr:08X} {n:5d}B  {name}")

if __name__ == "__main__":
    main()