#!/usr/bin/env python3
"""
Pattern-based C code generator for Orbital functions.

Reads ROM bytes, identifies common Thumb instruction patterns, and
generates candidate C code. Compiles each candidate with agbcc and
checks for byte-exact match (relocation-masked).

Usage: python3 tools/pattern_match.py
"""

import json
import os
import struct
import subprocess
import sys
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: bytes) -> bytearray:
    masked = bytearray(data)
    i = 0
    while i < len(masked) - 3:
        hw1 = int.from_bytes(masked[i:i+2], 'little')
        hw2 = int.from_bytes(masked[i+2:i+4], 'little')
        if (hw1 >> 11) == 0b11110 and (hw2 >> 11) == 0b11111:
            hw1 &= 0xF800; hw2 &= 0xF800
            masked[i]=hw1&0xFF; masked[i+1]=hw1>>8
            masked[i+2]=hw2&0xFF; masked[i+3]=hw2>>8
            i += 4; continue
        i += 2
    return masked

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

def try_compile(c_code, addr, size, rom):
    with tempfile.TemporaryDirectory() as td:
        c_path = os.path.join(td, "f.c")
        s_path = os.path.join(td, "f.s")
        o_path = os.path.join(td, "f.o")
        bin_path = os.path.join(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, None
        subprocess.run([AS, "-mcpu=arm7tdmi", s_path, "-o", o_path],
                      capture_output=True, timeout=10)
        if not Path(o_path).exists(): return False, None
        subprocess.run([OBJCOPY, "-O", "binary", "-j", ".text", o_path, bin_path],
                      capture_output=True, timeout=10)
        if not Path(bin_path).exists(): return False, None
        compiled = Path(bin_path).read_bytes()
        offset = addr - ROM_BASE
        rom_bytes = rom[offset:offset+size]
        real = size
        while real > 2 and rom_bytes[real-1] == 0: real -= 1
        if len(compiled) < real: return False, None
        c_trimmed = compiled[:real]
        if mask_bl(c_trimmed) == mask_bl(rom_bytes[:real]):
            return True, c_code
        return False, None

def disasm_thumb(func_bytes):
    """Simple Thumb16 disassembler for pattern matching."""
    lines = []
    i = 0
    while i < len(func_bytes) - 1:
        hw = struct.unpack_from('<H', func_bytes, i)[0]
        if hw == 0x4770: lines.append("bx lr")
        elif (hw & 0xFF00) == 0xB500:
            regs = ["lr"]
            for b in range(7, -1, -1):
                if hw & (1 << b): regs.insert(0, f"r{b}")
            lines.append(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}")
            lines.append(f"pop {{{','.join(regs)}}}")
        elif (hw & 0xF800) == 0x2000:
            lines.append(f"movs r{hw>>8&7}, #{hw&0xFF}")
        elif (hw & 0xF800) == 0x6000:
            lines.append(f"str r{hw&7}, [r{(hw>>3)&7}, #{((hw>>6)&0x1f)*4}]")
        elif (hw & 0xF800) == 0x6800:
            lines.append(f"ldr r{hw&7}, [r{(hw>>3)&7}, #{((hw>>6)&0x1f)*4}]")
        elif (hw & 0xF800) == 0x3000:
            lines.append(f"adds r{hw>>8&7}, #{hw&0xFF}")
        elif (hw & 0xF800) == 0x3800:
            lines.append(f"subs r{hw>>8&7}, #{hw&0xFF}")
        elif (hw & 0xFFC0) == 0x1C00:
            lines.append(f"adds r{hw&7}, r{(hw>>3)&7}, #{(hw>>6)&7}")
        elif (hw & 0xFFC0) == 0x1E00:
            lines.append(f"subs r{hw&7}, r{(hw>>3)&7}, #{(hw>>6)&7}")
        elif (hw & 0xFFC0) == 0x4280:
            lines.append(f"cmp r{hw&7}, r{(hw>>3)&7}")
        elif (hw & 0xF800) == 0x2800:
            lines.append(f"cmp r{hw>>8&7}, #{hw&0xFF}")
        elif (hw & 0xFF00) == 0xDF00:
            lines.append(f"swi #{hw&0xFF}")
        elif (hw & 0xFFC0) == 0x4300:
            lines.append(f"orrs r{hw&7}, r{(hw>>3)&7}")
        elif (hw & 0xFFC0) == 0x4040:
            lines.append(f"eors r{hw&7}, r{(hw>>3)&7}")
        elif (hw & 0xFFC0) == 0x4000:
            lines.append(f"ands r{hw&7}, r{(hw>>3)&7}")
        elif (hw & 0xFFC0) == 0x1800:
            lines.append(f"adds r{hw&7}, r{(hw>>3)&7}, r{(hw>>6)&7}")
        elif (hw & 0xFFC0) == 0x1A00:
            lines.append(f"subs r{hw&7}, r{(hw>>3)&7}, r{(hw>>6)&7}")
        elif (hw & 0xFFC0) == 0x4600:
            dst = (hw & 7) | ((hw >> 4) & 8)
            src = (hw >> 3) & 0xF
            lines.append(f"mov r{dst}, r{src}")
        elif (hw >> 11) == 0x9:
            lines.append(f"ldr r{hw>>8&7}, [pc, #{(hw&0xFF)*4}]")
        else:
            lines.append(f".short 0x{hw:04X}")
        i += 2
    return lines

def generate_candidates(func_bytes, addr, size, rom):
    """Generate candidate C code based on disassembly patterns."""
    lines = disasm_thumb(func_bytes)
    candidates = []
    
    # Pattern: push {lr}; ...; pop {r0}; bx r0  (framed void function)
    if len(lines) >= 3 and lines[0].startswith("push") and "lr" in lines[0]:
        if lines[-1] == "bx lr" or (len(lines)>=2 and "pop" in lines[-2] and "bx lr" in lines[-1]):
            # Try to understand the body
            body_lines = lines[1:-1] if lines[-1] == "bx lr" else lines[1:-2]
            
            # Check if body is all stores to r0+offset (struct field writes)
            all_stores = True
            stores = []
            for l in body_lines:
                if l.startswith("str ") and "[r0," in l:
                    import re
                    m = re.search(r'str r(\d+), \[r0, #(\d+)\]', l)
                    if m:
                        stores.append((int(m.group(1)), int(m.group(2))))
                    else:
                        all_stores = False
                elif l.startswith("movs r") and ", #" in l:
                    import re
                    m = re.search(r'movs r(\d+), #(\d+)', l)
                    if m and int(m.group(2)) == 0:
                        stores.append(("zero", int(m.group(1))))
                    else:
                        all_stores = False
                else:
                    all_stores = False
            
            if all_stores and stores:
                # Generate struct-based function
                max_off = max(s[1] if isinstance(s[1], int) else 0 for s in stores)
                fields = []
                for off in range(0, max_off + 4, 4):
                    fields.append(f"    int f{off};")
                struct_def = "struct S {\n" + "\n".join(fields) + "\n};\n"
                assigns = []
                for s in stores:
                    if s[0] == "zero":
                        assigns.append(f"    p->f{s[1]*4} = 0;")
                    else:
                        assigns.append(f"    p->f{s[1]} = r{s[0]};")
                func = f"void f(struct S *p) {{\n" + "\n".join(assigns) + "\n}\n"
                candidates.append(struct_def + "\n" + func)
    
    # Pattern: movs r0, #N; bx lr (return constant) — already matched by batch_match
    # Pattern: str r1, [r0, #off]; str r2, [r0, #off2]; bx lr (simple setter)
    if lines and lines[-1] == "bx lr" and len(lines) >= 2:
        all_str = True
        stores = []
        for l in lines[:-1]:
            import re
            m = re.match(r'str r(\d+), \[r0, #(\d+)\]', l)
            if m:
                stores.append((int(m.group(1)), int(m.group(2))))
            else:
                all_str = False
        if all_str and stores:
            max_off = max(s[1] for s in stores)
            fields = []
            for off in range(0, max_off + 4, 4):
                fields.append(f"    int f{off};")
            struct_def = "struct S {\n" + "\n".join(fields) + "\n};\n"
            params = []
            assigns = []
            reg_params = {}
            for reg, off in stores:
                if reg not in reg_params:
                    reg_params[reg] = f"a{reg}"
                assigns.append(f"    p->f{off} = {reg_params[reg]};")
            param_str = ", ".join(f"int {v}" for _, v in sorted(reg_params.items()))
            func = f"void f(struct S *p{', ' + param_str if param_str else ''}) {{\n" + "\n".join(assigns) + "\n}\n"
            candidates.append(struct_def + "\n" + func)
    
    # Pattern: ldr r0, [pc, #N]; ...; bx lr (return global)
    # Pattern: ldr r0, [r0, #off]; bx lr (dereference struct pointer)
    
    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 = []
    attempts = 0
    
    for addr, size, name in funcs:
        if not (CODE_START <= addr < CODE_END) or size < 2:
            continue
        
        offset = addr - ROM_BASE
        func_bytes = rom[offset:offset+size]
        
        candidates = generate_candidates(func_bytes, addr, size, rom)
        
        for c_code in candidates:
            attempts += 1
            ok, _ = try_compile(c_code, addr, size, rom)
            if ok:
                matched.append((addr, size, name, c_code))
                break
    
    matched_bytes = sum(s for _, s, _, _ in matched)
    total_bytes = sum(s for a, s, _ in funcs if CODE_START <= a < CODE_END and s >= 2)
    
    print(f"Attempts: {attempts}")
    print(f"Matched:  {len(matched)} functions, {matched_bytes} bytes ({100*matched_bytes/total_bytes:.1f}%)")
    print(f"\nMatched functions:")
    for addr, size, name, code in sorted(matched, key=lambda x: -x[1]):
        print(f"  0x{addr:08X} {size:4d}B  {name}")
        # Save matched C to src/game/
        outdir = ROOT / "src" / "matched"
        outdir.mkdir(exist_ok=True)
        (outdir / f"{name}.c").write_text(code)

if __name__ == "__main__":
    main()