import ghidra.app.script.GhidraScript;
import ghidra.program.model.listing.FunctionManager;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.address.Address;
import ghidra.program.model.listing.Program;
import ghidra.program.model.mem.Memory;
import ghidra.program.model.mem.MemoryBlock;
import ghidra.util.task.TaskMonitor;
import java.io.FileWriter;
import java.io.IOException;
import java.util.ArrayList;
import java.util.List;

public class ExtractFunctions extends GhidraScript {

    @Override
    protected void run() throws Exception {
        Program program = currentProgram;
        FunctionManager functionManager = program.getFunctionManager();
        Memory memory = program.getMemory();

        String outputDir = "/home/rmcore/Git/integdev_gpu_drv/open-iov/firmware/570-GSP";

        // 1. Extract memory blocks
        List<MemoryBlockInfo> memBlocks = new ArrayList<>();
        for (MemoryBlock block : memory.getBlocks()) {
            memBlocks.add(new MemoryBlockInfo(
                block.getName(),
                block.getStart().toString(),
                block.getEnd().toString(),
                block.getSize(),
                block.isExecute(),
                block.isRead(),
                block.isWrite()
            ));
        }
        writeJson(outputDir + "/memory_blocks.json", memBlocks);

        // 2. Extract all functions (analysis must have run first)
        List<FunctionInfo> functions = new ArrayList<>();
        FunctionIterator funcIter = functionManager.getFunctions(true);
        while (funcIter.hasNext()) {
            Function func = funcIter.next();
            functions.add(new FunctionInfo(
                func.getName(),
                func.getEntryPoint().toString(),
                func.getBody().getNumAddresses()
            ));
        }
        writeJson(outputDir + "/functions.json", functions);

        // 3. Extract interesting strings
        String[] keywords = {"vgpu", "vmiop", "rpc", "hal", "bar", "migrate", "init", "boot", "kgsp", "wpr", "firmware", "version", "error", "fail", "assert", "panic", "nvidia", "rmapi", "gsp", "falcon", "mailbox", "interrupt", "dma", "iommu", "page", "mmu", "pde", "pte"};
        List<StringInfo> strings = new ArrayList<>();
        var listing = program.getListing();
        var dataIter = listing.getDefinedData(true);
        while (dataIter.hasNext()) {
            var data = dataIter.next();
            if (data.getDataType().getName().equals("string")) {
                String value = data.getValue().toString();
                if (value.length() > 4) {
                    String lower = value.toLowerCase();
                    for (String kw : keywords) {
                        if (lower.contains(kw)) {
                            strings.add(new StringInfo(data.getAddress().toString(), value));
                            break;
                        }
                    }
                }
            }
        }
        writeJson(outputDir + "/interesting_strings.json", strings);

        println("Extracted " + functions.size() + " functions, " + memBlocks.size() + " memory blocks, " + strings.size() + " interesting strings");
    }

    private void writeJson(String path, List<?> list) throws IOException {
        try (FileWriter fw = new FileWriter(path)) {
            fw.write("[\n");
            for (int i = 0; i < list.size(); i++) {
                fw.write("  " + toJson(list.get(i)));
                if (i < list.size() - 1) fw.write(",");
                fw.write("\n");
            }
            fw.write("]\n");
        }
    }

    private String toJson(Object obj) {
        if (obj instanceof MemoryBlockInfo) {
            MemoryBlockInfo m = (MemoryBlockInfo) obj;
            return String.format("{\"name\":\"%s\",\"start\":\"%s\",\"end\":\"%s\",\"size\":%d,\"execute\":%b,\"read\":%b,\"write\":%b}",
                m.name, m.start, m.end, m.size, m.execute, m.read, m.write);
        } else if (obj instanceof FunctionInfo) {
            FunctionInfo f = (FunctionInfo) obj;
            return String.format("{\"name\":\"%s\",\"entry\":\"%s\",\"size\":%d}",
                f.name, f.entry, f.size);
        } else if (obj instanceof StringInfo) {
            StringInfo s = (StringInfo) obj;
            return String.format("{\"address\":\"%s\",\"value\":\"%s\"}",
                s.address, s.value.replace("\\", "\\\\").replace("\"", "\\\""));
        }
        return "{}";
    }

    static class MemoryBlockInfo {
        String name, start, end;
        long size;
        boolean execute, read, write;
        MemoryBlockInfo(String n, String s, String e, long sz, boolean ex, boolean rd, boolean wr) {
            name = n; start = s; end = e; size = sz; execute = ex; read = rd; write = wr;
        }
    }

    static class FunctionInfo {
        String name, entry;
        long size;
        FunctionInfo(String n, String e, long sz) { name = n; entry = e; size = sz; }
    }

    static class StringInfo {
        String address, value;
        StringInfo(String a, String v) { address = a; value = v; }
    }
}