import ghidra.app.script.GhidraScript;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionManager;
import ghidra.program.model.listing.Instruction;
import ghidra.program.model.listing.Program;
import ghidra.program.model.mem.Memory;
import ghidra.app.decompiler.DecompInterface;
import ghidra.app.decompiler.DecompileResults;
import ghidra.util.task.TaskMonitor;
import java.io.*;

public class DecompileFunctions extends GhidraScript {
    
    private DecompInterface decompiler;
    
    @Override
    protected void run() throws Exception {
        decompiler = new DecompInterface();
        decompiler.openProgram(currentProgram);
        
        String[] args = getScriptArgs();
        if (args.length == 0) {
            println("Usage: DecompileFunctions.java <comma-separated-addresses>");
            return;
        }
        
        String[] addresses = args[0].split(",");
        
        println("=== Targeted Decompilation ===");
        for (String addrStr : addresses) {
            addrStr = addrStr.trim();
            if (addrStr.isEmpty()) continue;
            
            long addrVal;
            try {
                addrVal = Long.parseUnsignedLong(addrStr.startsWith("0x") ? addrStr.substring(2) : addrStr, 16);
            } catch (NumberFormatException e) {
                println("Invalid address: " + addrStr);
                continue;
            }
            
            Address addr = currentProgram.getAddressFactory().getAddress("0x" + Long.toHexString(addrVal));
            if (addr == null) {
                println("Invalid address: " + addrStr);
                continue;
            }
            
            decompileFunction(addr);
        }
        
        println("\n=== Decompilation Complete ===");
    }
    
    private void decompileFunction(Address entryAddr) throws Exception {
        FunctionManager fm = currentProgram.getFunctionManager();
        Memory mem = currentProgram.getMemory();
        
        // Manually disassemble from entry point by creating instructions
        Address addr = entryAddr;
        int maxBytes = 2000;
        int bytes = 0;
        
        while (bytes < maxBytes && !monitor.isCancelled()) {
            Instruction instr = currentProgram.getListing().getInstructionAt(addr);
            if (instr != null) {
                addr = instr.getMaxAddress().next();
                bytes += instr.getLength();
                continue;
            }
            
            // Try to create instruction from bytes
            try {
                byte[] buf = new byte[4];
                int read = mem.getBytes(addr, buf);
                if (read <= 0) break;
                
                // Create instruction prototype and try to disassemble
                instr = currentProgram.getListing().createInstruction(addr, 
                    currentProgram.getLanguage().getInstructionPrototype(buf), 
                    new ghidra.program.model.mem.ByteMemBuffer(buf, 0, read),
                    currentProgram.getProgramContext(), 0);
                    
                if (instr != null) {
                    addr = instr.getMaxAddress().next();
                    bytes += instr.getLength();
                } else {
                    addr = addr.next().next();
                    bytes += 2;
                }
            } catch (Exception e) {
                addr = addr.next().next();
                bytes += 2;
            }
        }
        println("  Disassembled " + bytes + " bytes from " + entryAddr);
        
        // Get or create function
        Function func = fm.getFunctionAt(entryAddr);
        if (func == null) {
            func = fm.createFunction("FUNC_" + Long.toHexString(entryAddr.getOffset()).toUpperCase(), 
                                   entryAddr, null, ghidra.program.model.symbol.SourceType.ANALYSIS);
            if (func == null) {
                println("Failed to create function at " + entryAddr);
                return;
            }
            println("Created function: " + func.getName() + " @ " + func.getEntryPoint());
        } else {
            println("Existing function: " + func.getName() + " @ " + func.getEntryPoint());
        }
        
        // Decompile
        println("\n--- Decompiling " + func.getName() + " @ " + func.getEntryPoint() + " ---");
        DecompileResults results = decompiler.decompileFunction(func, 120, monitor);
        
        if (results.decompileCompleted()) {
            String cCode = results.getDecompiledFunction().getC();
            
            // Save to file
            String filename = "/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP/decompiled_" 
                            + func.getName() + ".c";
            try (PrintWriter pw = new PrintWriter(new FileWriter(filename))) {
                pw.println("// Function: " + func.getName());
                pw.println("// Address: " + func.getEntryPoint());
                pw.println("// Stack frame: " + func.getStackFrame().getFrameSize() + " bytes");
                pw.println();
                pw.print(cCode);
            }
            println("  Saved to: " + filename);
            
            // Print first 100 lines
            String[] lines = cCode.split("\n");
            for (int i = 0; i < Math.min(100, lines.length); i++) {
                println("  " + lines[i]);
            }
            if (lines.length > 100) {
                println("  ... (" + (lines.length - 100) + " more lines)");
            }
        } else {
            println("  Decompilation failed: " + results.getErrorMessage());
        }
    }
}