#!/usr/bin/env python3
"""
Fixed search for RISC-V function prologues in GSP firmware ELF.
Handles the actual prologue patterns found.
"""

import struct
import json
import os

FWIMAGE_PATH = os.path.join(os.path.dirname(os.path.abspath(__file__)), "gsp-570.144.fwimage.bin")
OUTPUT_DIR = os.path.dirname(os.path.abspath(__file__))
BASE_ADDR = 0x04000000

# Executable sections from memory_blocks.json (virtual addresses)
EXEC_SECTIONS = [
    (0x04005000, 0x04006413),  # .ga100_text
    (0x04007000, 0x04008413),  # .tu10x_text
    (0x04009000, 0x04009073),  # .bootstrap
    (0x0400a000, 0x0400cfff),  # .ga100_resident_text
    (0x0400e000, 0x04010fff),  # .tu10x_resident_text
    (0x04018000, 0x0401cfff),  # .section_task_init_elf_text_instance
    (0x04c58000, 0x05b21fff),  # .section_task_rm_elf_text_instance (MAIN - 15.5MB)
]

def read_binary(path):
    with open(path, 'rb') as f:
        return f.read()

def va_to_offset(va):
    """Convert virtual address to file offset."""
    return va - BASE_ADDR

def is_in_exec_section(va):
    """Check if virtual address is in an executable section."""
    for start, end in EXEC_SECTIONS:
        if start <= va <= end:
            return True
    return False

def find_prologues(data, base_addr):
    functions = []
    
    # Pattern 1: Standard prologue with s0 frame pointer
    # addi sp, sp, -N; sd ra, N-8(sp); addi s0, sp, N; ... ; ld ra, N-8(sp); addi sp, sp, N; ret
    # The key signature: addi sp, sp, -N followed by sd ra, N-8(sp)
    for i in range(0, len(data) - 8, 2):
        va = base_addr + i
        if not is_in_exec_section(va):
            continue
            
        if i + 8 <= len(data):
            inst1 = struct.unpack('<I', data[i:i+4])[0]
            inst2 = struct.unpack('<I', data[i+4:i+8])[0]
            
            # addi sp, sp, -imm (opcode=0x13, funct3=0, rd=2, rs1=2)
            if (inst1 & 0x707F) == 0x0013:
                rd = (inst1 >> 7) & 0x1F
                rs1 = (inst1 >> 15) & 0x1F
                if rd == 2 and rs1 == 2:
                    imm = (inst1 >> 20) & 0xFFF
                    if imm >= 0x800:  # Negative immediate
                        stack_size = -(imm - 0x1000)
                    else:
                        continue  # Not a stack allocation
                    
                    # sd ra, offset(sp) (opcode=0x23, funct3=3, rs1=2, rs2=1)
                    if (inst2 & 0x707F) == 0x23:
                        funct3 = (inst2 >> 12) & 0x7
                        rs1_2 = (inst2 >> 15) & 0x1F
                        rs2 = (inst2 >> 20) & 0x1F
                        if funct3 == 3 and rs1_2 == 2 and rs2 == 1:
                            functions.append({
                                "address": f"0x{va:08x}",
                                "offset": i,
                                "stack_size": stack_size,
                                "pattern": "std_prologue_ra"
                            })
                            continue
    
    # Pattern 2: Compressed prologue
    # c.addi sp, sp, -imm; c.sd ra, offset(sp)
    for i in range(0, len(data) - 4, 2):
        va = base_addr + i
        if not is_in_exec_section(va):
            continue
            
        if i + 4 <= len(data):
            inst1 = struct.unpack('<H', data[i:i+2])[0]
            inst2 = struct.unpack('<H', data[i+2:i+4])[0]
            
            # c.addi sp, sp, imm (rd'=2)
            if (inst1 & 0xE003) == 0x6001:
                rd_prime = (inst1 >> 7) & 0x7
                if rd_prime == 2:
                    # c.sd ra, offset(sp) (rs1'=2, rs2'=1)
                    if (inst2 & 0xE003) == 0xE000:
                        rs1_prime = (inst2 >> 7) & 0x7
                        rs2_prime = (inst2 >> 10) & 0x7
                        if rs1_prime == 2 and rs2_prime == 1:
                            functions.append({
                                "address": f"0x{va:08x}",
                                "offset": i,
                                "pattern": "compressed_prologue"
                            })
                            continue
    
    # Pattern 3: Full frame with s0
    # addi sp, sp, -N; sd ra, N-8(sp); sd s0, N-16(sp); addi s0, sp, N
    for i in range(0, len(data) - 16, 2):
        va = base_addr + i
        if not is_in_exec_section(va):
            continue
            
        if i + 16 <= len(data):
            inst1 = struct.unpack('<I', data[i:i+4])[0]
            inst2 = struct.unpack('<I', data[i+4:i+8])[0]
            inst3 = struct.unpack('<I', data[i+8:i+12])[0]
            inst4 = struct.unpack('<I', data[i+12:i+16])[0]
            
            # addi sp, sp, -N
            if (inst1 & 0x707F) == 0x0013:
                rd = (inst1 >> 7) & 0x1F
                rs1 = (inst1 >> 15) & 0x1F
                if rd == 2 and rs1 == 2:
                    imm = (inst1 >> 20) & 0xFFF
                    if imm >= 0x800:
                        stack_size = -(imm - 0x1000)
                    else:
                        continue
                    
                    # sd ra, N-8(sp)
                    if (inst2 & 0x707F) == 0x23:
                        funct3 = (inst2 >> 12) & 0x7
                        rs1_2 = (inst2 >> 15) & 0x1F
                        rs2 = (inst2 >> 20) & 0x1F
                        if funct3 == 3 and rs1_2 == 2 and rs2 == 1:
                            # sd s0, N-16(sp)
                            if (inst3 & 0x707F) == 0x23:
                                funct3_3 = (inst3 >> 12) & 0x7
                                rs1_3 = (inst3 >> 15) & 0x1F
                                rs2_3 = (inst3 >> 20) & 0x1F
                                if funct3_3 == 3 and rs1_3 == 2 and rs2_3 == 8:
                                    # addi s0, sp, N
                                    if (inst4 & 0x707F) == 0x0013:
                                        rd4 = (inst4 >> 7) & 0x1F
                                        rs1_4 = (inst4 >> 15) & 0x1F
                                        if rd4 == 8 and rs1_4 == 2:
                                            functions.append({
                                                "address": f"0x{va:08x}",
                                                "offset": i,
                                                "stack_size": stack_size,
                                                "pattern": "full_prologue_s0"
                                            })
                                            continue
    
    return functions

def main():
    print(f"Reading firmware binary: {FWIMAGE_PATH}")
    data = read_binary(FWIMAGE_PATH)
    print(f"Binary size: {len(data)} bytes")
    
    print("Searching for function prologues in executable sections...")
    functions = find_prologues(data, BASE_ADDR)
    
    print(f"Found {len(functions)} potential functions")
    
    # Deduplicate by address
    seen = set()
    unique = []
    for f in functions:
        if f["address"] not in seen:
            seen.add(f["address"])
            unique.append(f)
    
    # Sort by address
    unique.sort(key=lambda x: x["offset"])
    
    print(f"Unique functions: {len(unique)}")
    
    # Save results
    with open(f"{OUTPUT_DIR}/found_prologues_v2.json", "w") as f:
        json.dump(unique, f, indent=2)
    
    with open(f"{OUTPUT_DIR}/prologue_addresses_v2.txt", "w") as f:
        for func in unique:
            ss = func.get("stack_size", "?")
            f.write(f"{func['address']}  stack_size={ss}  pattern={func['pattern']}\n")
    
    print(f"Results saved to {OUTPUT_DIR}/found_prologues_v2.json")
    
    # Group by pattern
    by_pattern = {}
    for func in unique:
        p = func["pattern"]
        if p not in by_pattern:
            by_pattern[p] = []
        by_pattern[p].append(func)
    
    for pattern, funcs in by_pattern.items():
        print(f"\n{pattern}: {len(funcs)} functions")
        for func in funcs[:30]:
            ss = func.get("stack_size", "?")
            print(f"  {func['address']}  stack_size={ss}")
        if len(funcs) > 30:
            print(f"  ... and {len(funcs) - 30} more")

if __name__ == "__main__":
    main()