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.app.decompiler.DecompInterface;
import ghidra.app.decompiler.DecompileResults;
import ghidra.util.task.TaskMonitor;
import java.io.File;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.util.Arrays;
import java.util.HashSet;
import java.util.Set;

public class DecompileFunctions extends GhidraScript {
    private static final String OUTPUT_DIR = "/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP/decompiled/";

    @Override
    protected void run() throws Exception {
        File outDir = new File(OUTPUT_DIR);
        if (!outDir.exists()) {
            outDir.mkdirs();
        }

        String[] args = getScriptArgs();
        if (args.length == 0) {
            println("No addresses provided. Usage: <addr1>,<addr2>,...");
            return;
        }

        Set<Address> targetAddresses = new HashSet<>();
        for (String arg : args) {
            String[] addrs = arg.split(",");
            for (String addrStr : addrs) {
                addrStr = addrStr.trim();
                if (addrStr.startsWith("0x")) {
                    addrStr = addrStr.substring(2);
                }
                Address addr = toAddr(addrStr);
                if (addr != null) {
                    targetAddresses.add(addr);
                    println("Target: " + addr);
                }
            }
        }

        FunctionManager fm = currentProgram.getFunctionManager();
        DecompInterface decompiler = new DecompInterface();
        decompiler.openProgram(currentProgram);

        for (Address addr : targetAddresses) {
            Function func = fm.getFunctionAt(addr);
            if (func == null) {
                println("No function at " + addr + ", creating...");
                func = createFunction(addr, null);
            }
            if (func != null) {
                decompileAndSave(func, decompiler);
            } else {
                println("Failed to create function at " + addr);
            }
        }

        decompiler.dispose();
    }

    private void decompileAndSave(Function func, DecompInterface decompiler) throws Exception {
        String name = func.getName();
        Address entry = func.getEntryPoint();
        String safeName = name.replaceAll("[^a-zA-Z0-9_]", "_");
        File outFile = new File(OUTPUT_DIR, safeName + "_" + entry.toString().replace(":", "_") + ".c");

        DecompileResults results = decompiler.decompileFunction(func, 60, TaskMonitor.DUMMY);
        if (results.decompileCompleted()) {
            String code = results.getDecompiledFunction().getC();
            try (PrintWriter pw = new PrintWriter(new FileWriter(outFile))) {
                pw.println("// Function: " + name);
                pw.println("// Address: " + entry);
                pw.println("// Size: " + func.getBody().getNumAddresses() + " bytes");
                pw.println();
                pw.print(code);
            }
            println("Decompiled: " + name + " @ " + entry + " -> " + outFile.getName());
        } else {
            println("Decompilation failed for " + name + " @ " + entry + ": " + results.getErrorMessage());
        }
    }
}