#!/usr/bin/env python3
"""
Search for RISC-V function prologues in GSP firmware binary.
RISC-V typical prologue:
  addi sp, sp, -N     (0x1141 + offset)
  sd ra, N-8(sp)      (0x80e2 + offset)  
  sd s0, N-16(sp)     (0x84e2 + offset)
  addi s0, sp, N      (0x1144 + offset)
"""

import struct
import json
import sys
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  # Load address from memory_blocks.json

# RISC-V instruction patterns (compressed and uncompressed)
# addi sp, sp, -imm (opcode 0x13, funct3 0, rd=sp=2, rs1=sp=2)
# sd ra, offset(sp) (opcode 0x23, funct3 3, rs1=sp=2, rs2=ra=1)
# sd s0, offset(sp) (opcode 0x23, funct3 3, rs1=sp=2, rs2=s0=8)
# addi s0, sp, imm (opcode 0x13, funct3 0, rd=s0=8, rs1=sp=2)

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

def find_prologues(data, base_addr):
    """Find potential function prologues in the binary."""
    functions = []
    
    # Search for standard RISC-V function prologue pattern
    # Pattern: addi sp, sp, -N; sd ra, N-8(sp); sd s0, N-16(sp); addi s0, sp, N
    # We'll look for the first two instructions as a signature
    
    for i in range(0, len(data) - 8, 2):  # 16-bit aligned
        # Check for 32-bit instructions (not compressed)
        if i + 8 <= len(data):
            inst1 = struct.unpack('<I', data[i:i+4])[0]
            inst2 = struct.unpack('<I', data[i+4:i+8])[0]
            
            # Check for addi sp, sp, -imm (opcode=0x13, rd=2, rs1=2, funct3=0)
            # 0x00000013 = addi zero, zero, 0 (nop)
            # addi sp, sp, -N: rd=2, rs1=2, funct3=0, opcode=0x13
            # Encoding: imm[11:0] | rs1 | funct3 | rd | opcode
            if (inst1 & 0x707F) == 0x0013:  # opcode=0x13, funct3=0
                rd = (inst1 >> 7) & 0x1F
                rs1 = (inst1 >> 15) & 0x1F
                if rd == 2 and rs1 == 2:  # sp = x2
                    # Potential addi sp, sp, -N
                    imm = (inst1 >> 20) & 0xFFF
                    if imm >= 0x800:  # Negative immediate (sign extended)
                        imm = imm - 0x1000
                    
                    # Check next instruction: sd ra, offset(sp)
                    if (inst2 & 0x707F) == 0x23:  # opcode=0x23 (store)
                        funct3 = (inst2 >> 12) & 0x7
                        rs1_2 = (inst2 >> 15) & 0x1F
                        rs2 = (inst2 >> 20) & 0x1F
                        if funct3 == 3 and rs1_2 == 2 and rs2 == 1:  # sd ra, offset(sp)
                            # Found potential prologue
                            addr = base_addr + i
                            functions.append({
                                "address": f"0x{addr:08x}",
                                "offset": i,
                                "stack_size": -imm,
                                "type": "standard_prologue"
                            })
    
    return functions

def find_all_functions(data, base_addr):
    """Find all potential functions using multiple patterns."""
    functions = []
    
    # Pattern 1: Standard prologue with addi sp, sp, -N
    for i in range(0, len(data) - 8, 2):
        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
            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:
                        imm = imm - 0x1000
                    
                    # sd ra, offset(sp) - check multiple possible offsets
                    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:
                            addr = base_addr + i
                            functions.append({
                                "address": f"0x{addr:08x}",
                                "offset": i,
                                "stack_size": -imm,
                                "pattern": "addi_sp_sd_ra"
                            })
    
    # Pattern 2: Compressed instructions (16-bit)
    # c.addi sp, sp, -N (opcode 0x6101 + offset)
    # c.sd ra, offset(sp) (opcode 0xE006 + offset)
    for i in range(0, len(data) - 4, 2):
        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'=sp=2, imm != 0)
            # Format: 0x6101 | (imm[5] << 12) | (imm[4:0] << 7) | (rs1'/rd' << 7)
            if (inst1 & 0xE003) == 0x6001:  # c.addi
                rd_prime = (inst1 >> 7) & 0x7
                if rd_prime == 2:  # sp
                    # c.sd ra, offset(sp) - ra=x1, sp=x2
                    # Format: 0xE000 | (offset[8:6] << 10) | (rs2' << 7) | (rs1' << 7) | offset[5:3]
                    if (inst2 & 0xE003) == 0xE000:  # c.sd
                        rs1_prime = (inst2 >> 7) & 0x7
                        rs2_prime = (inst2 >> 10) & 0x7
                        if rs1_prime == 2 and rs2_prime == 1:  # sp, ra
                            addr = base_addr + i
                            functions.append({
                                "address": f"0x{addr:08x}",
                                "offset": i,
                                "pattern": "compressed_prologue"
                            })
    
    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...")
    functions = find_all_functions(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)
    
    print(f"Unique functions: {len(unique)}")
    
    # Save results
    with open(f"{OUTPUT_DIR}/found_functions.json", "w") as f:
        json.dump(unique, f, indent=2)
    
    # Also create a simple address list
    with open(f"{OUTPUT_DIR}/function_addresses.txt", "w") as f:
        for func in sorted(unique, key=lambda x: x["offset"]):
            f.write(f"{func['address']}  stack_size={func.get('stack_size', '?')}  pattern={func['pattern']}\n")
    
    print(f"Results saved to {OUTPUT_DIR}/found_functions.json")
    print(f"Address list saved to {OUTPUT_DIR}/function_addresses.txt")
    
    # Show first 50
    for func in sorted(unique, key=lambda x: x["offset"])[:50]:
        print(f"  {func['address']}  stack_size={func.get('stack_size', '?')}  pattern={func['pattern']}")

if __name__ == "__main__":
    main()