// SPDX-License-Identifier: BSD-3-Clause
//
// AnalyzeFunctions.java - Ghidra post-script: run the "Decompiler Parameter ID"
// analyzer over a caller-supplied list of functions, using the analyzer's
// AddressSetView entry point instead of a whole-program analysis pass.
//
// This is the targeted type-recovery pass for the CUDA userspace driver: the
// full auto-analysis deliberately disables the parameter-ID analyzer because it
// decompiles every function and is very expensive on a 29k-function stripped
// binary. Re-running it over a curated address set (crash path, engine-record
// path, handle pool, RM alloc, resolved wrapper implementations, the largest
// .text functions) recovers parameter counts/types and local variables where
// the decompiler can infer them, without touching the rest of the program.
//
// Usage:
//   analyzeHeadless <proj_dir> <proj_name> -process <program> -noanalysis \
//     -scriptPath tools/ghidra \
//     -postScript AnalyzeFunctions.java <addresses_file> \
//         [--timeout <secs>] [--out <tsv>] [--list] [--batch <n>]
//
//   <addresses_file>  one entry per line: "0x<file_vaddr> [# comment]".
//                     Addresses are FILE virtual addresses (libcuda+0x...).
//                     Blank lines and lines starting with '#' are ignored.
//   --timeout <secs>  per-function decompiler timeout for the analyzer
//                     (default 60, matching the analyzer's own default).
//   --out <tsv>       write the before/after report to this file as well.
//   --list            only resolve the addresses to functions and print them;
//                     do not run the analyzer (dry run).
//   --batch <n>       run the analyzer in batches of n functions (default 0 =
//                     one batch). Useful to checkpoint very large sets.
//
// The script translates file vaddr -> program address via the program's image
// base (0x100000 for this binary), exactly like DecompileAt.java. If an address
// is not inside a known function it is disassembled and turned into one.
//
// The analyzer class was renamed between Ghidra versions
// (DecompilerParameterIdAnalyzer -> DecompilerFunctionAnalyzer); the script
// looks up both and invokes the AddressSetView-based entry point
// added(AddressSetView, TaskMonitor) when present, falling back to
// added(Program, AddressSetView, TaskMonitor, MessageLog).

import ghidra.app.script.GhidraScript;
import ghidra.app.services.Analyzer;
import ghidra.app.util.importer.MessageLog;
import ghidra.framework.options.Options;
import ghidra.program.model.address.Address;
import ghidra.program.model.address.AddressSet;
import ghidra.program.model.address.AddressSetView;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.Listing;
import ghidra.program.model.listing.Variable;
import ghidra.util.classfinder.ClassSearcher;
import ghidra.util.task.TaskMonitor;

import java.io.BufferedReader;
import java.io.FileReader;
import java.io.FileWriter;
import java.io.PrintWriter;
import java.lang.reflect.InvocationTargetException;
import java.lang.reflect.Method;
import java.util.ArrayList;
import java.util.List;

public class AnalyzeFunctions extends GhidraScript {

    private static final String[] ANALYZER_CLASSES = {
        "ghidra.app.plugin.core.analysis.DecompilerFunctionAnalyzer",
        "ghidra.app.plugin.core.analysis.DecompilerParameterIdAnalyzer",
    };

    private static final String OPTION_TIMEOUT = "Analysis Decompiler Timeout (sec)";

    private static class Target {
        long fileAddr;
        Address entry;
        String name;
        int paramsBefore;
        int localsBefore;
        int paramsAfter = -1;
        int localsAfter = -1;
    }

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();
        if (args.length < 1) {
            println("usage: AnalyzeFunctions.java <addresses_file> [--timeout secs] " +
                    "[--out tsv] [--list] [--batch n]");
            return;
        }

        String addrFile = args[0];
        int timeoutSecs = 60;
        String outFile = null;
        boolean listOnly = false;
        int batchSize = 0;
        for (int i = 1; i < args.length; i++) {
            switch (args[i]) {
                case "--timeout":
                    timeoutSecs = Integer.parseInt(args[++i]);
                    break;
                case "--out":
                    outFile = args[++i];
                    break;
                case "--list":
                    listOnly = true;
                    break;
                case "--batch":
                    batchSize = Integer.parseInt(args[++i]);
                    break;
                default:
                    println("unknown argument: " + args[i]);
                    return;
            }
        }

        long imageBase = currentProgram.getImageBase().getOffset();
        List<Long> fileAddrs = readAddresses(addrFile);
        println("AnalyzeFunctions: " + fileAddrs.size() + " address(es), image base 0x"
                + Long.toHexString(imageBase));

        List<Target> targets = resolveTargets(fileAddrs, imageBase);
        println("AnalyzeFunctions: resolved " + targets.size() + " function(s)");

        if (listOnly) {
            for (Target t : targets) {
                println(String.format("  libcuda+0x%06x  %s  program=0x%x  params=%d locals=%d",
                        t.fileAddr, t.name, t.entry.getOffset(), t.paramsBefore,
                        t.localsBefore));
            }
            return;
        }

        Analyzer analyzer = newAnalyzer();
        if (analyzer == null) {
            println("ERROR: no Decompiler Parameter ID analyzer class found");
            return;
        }
        println("AnalyzeFunctions: analyzer = " + analyzer.getClass().getName()
                + ", timeout = " + timeoutSecs + "s");

        Options opts = currentProgram.getOptions("Analyzers");
        analyzer.registerOptions(opts, currentProgram);
        if (opts.contains(OPTION_TIMEOUT)) {
            opts.setInt(OPTION_TIMEOUT, timeoutSecs);
        }
        analyzer.optionsChanged(opts, currentProgram);

        int txId = currentProgram.startTransaction("AnalyzeFunctions parameter ID");
        boolean commit = false;
        try {
            runAnalyzer(analyzer, targets, batchSize);
            commit = true;
        }
        finally {
            currentProgram.endTransaction(txId, commit);
        }

        snapshotAfter(targets, imageBase);
        report(targets, outFile);
    }

    private List<Long> readAddresses(String path) throws Exception {
        List<Long> out = new ArrayList<>();
        try (BufferedReader r = new BufferedReader(new FileReader(path))) {
            String line;
            while ((line = r.readLine()) != null) {
                line = line.trim();
                if (line.isEmpty() || line.startsWith("#")) {
                    continue;
                }
                String tok = line.split("\\s+")[0];
                if (tok.startsWith("0x") || tok.startsWith("0X")) {
                    tok = tok.substring(2);
                }
                try {
                    out.add(Long.parseUnsignedLong(tok, 16));
                }
                catch (NumberFormatException e) {
                    println("WARN: bad address line: " + line);
                }
            }
        }
        return out;
    }

    private List<Target> resolveTargets(List<Long> fileAddrs, long imageBase) {
        List<Target> targets = new ArrayList<>();
        Listing listing = currentProgram.getListing();
        for (long fileAddr : fileAddrs) {
            Address addr = toAddr(fileAddr + imageBase);
            Function f = getFunctionContaining(addr);
            if (f == null) {
                FunctionIterator it = listing.getFunctions(addr, false);
                if (it.hasNext()) {
                    f = it.next();
                }
            }
            if (f == null) {
                disassemble(addr);
                f = createFunction(addr, null);
            }
            if (f == null) {
                println("WARN: no function at libcuda+0x" + Long.toHexString(fileAddr));
                continue;
            }
            Target t = new Target();
            t.fileAddr = fileAddr;
            t.entry = f.getEntryPoint();
            t.name = f.getName();
            t.paramsBefore = f.getParameterCount();
            t.localsBefore = f.getLocalVariables().length;
            targets.add(t);
        }
        return targets;
    }

    private Analyzer newAnalyzer() throws Exception {
        // Preferred path: the plugin service discovers analyzers via ClassSearcher,
        // which also works when the analyzer lives in a module that is not visible
        // to the script's class loader (Class.forName then fails).
        try {
            ClassSearcher.search(monitor);
            for (Class<? extends Analyzer> c : ClassSearcher.getClasses(Analyzer.class)) {
                Object o = c.getDeclaredConstructor().newInstance();
                Analyzer a = (Analyzer) o;
                if ("Decompiler Parameter ID".equals(a.getName())) {
                    return a;
                }
            }
        }
        catch (Exception e) {
            println("WARN: ClassSearcher lookup failed: " + e);
        }
        // Fallback: direct reflection over the known class names.
        for (String cn : ANALYZER_CLASSES) {
            try {
                Class<?> c = Class.forName(cn);
                Object o = c.getDeclaredConstructor().newInstance();
                if (o instanceof Analyzer) {
                    return (Analyzer) o;
                }
            }
            catch (ClassNotFoundException e) {
                // try next candidate
            }
        }
        return null;
    }

    private void runAnalyzer(Analyzer analyzer, List<Target> targets, int batchSize)
            throws Exception {
        Method addedSet = findMethod(analyzer.getClass(), "added",
                AddressSetView.class, TaskMonitor.class);
        Method addedProgram = findMethod(analyzer.getClass(), "added",
                ghidra.program.model.listing.Program.class, AddressSetView.class,
                TaskMonitor.class, MessageLog.class);
        if (addedSet == null && addedProgram == null) {
            throw new NoSuchMethodException("no supported added(...) entry point on "
                    + analyzer.getClass().getName());
        }

        int n = (batchSize > 0) ? batchSize : targets.size();
        for (int start = 0; start < targets.size(); start += n) {
            int end = Math.min(start + n, targets.size());
            AddressSet set = new AddressSet();
            for (int i = start; i < end; i++) {
                set.addRange(targets.get(i).entry, targets.get(i).entry);
            }
            println(String.format("AnalyzeFunctions: batch %d..%d (%d functions)",
                    start, end - 1, end - start));
            try {
                if (addedSet != null) {
                    addedSet.invoke(analyzer, set, monitor);
                }
                else {
                    addedProgram.invoke(analyzer, currentProgram, set, monitor,
                            new MessageLog());
                }
            }
            catch (InvocationTargetException e) {
                Throwable cause = e.getCause();
                println("WARN: analyzer threw for batch " + start + ".." + (end - 1)
                        + ": " + (cause == null ? e : cause));
            }
        }
    }

    private Method findMethod(Class<?> c, String name, Class<?>... params) {
        for (Class<?> k = c; k != null; k = k.getSuperclass()) {
            try {
                Method m = k.getDeclaredMethod(name, params);
                m.setAccessible(true);
                return m;
            }
            catch (NoSuchMethodException e) {
                // keep walking up the hierarchy
            }
        }
        return null;
    }

    private void snapshotAfter(List<Target> targets, long imageBase) {
        for (Target t : targets) {
            Function f = getFunctionAt(t.entry);
            if (f == null) {
                continue;
            }
            t.paramsAfter = f.getParameterCount();
            Variable[] locals = f.getLocalVariables();
            t.localsAfter = locals.length;
            if (t.name.startsWith("FUN_") && !f.getName().equals(t.name)) {
                t.name = f.getName();
            }
        }
    }

    private void report(List<Target> targets, String outFile) throws Exception {
        StringBuilder sb = new StringBuilder();
        sb.append("# libcuda 570.144 targeted Decompiler Parameter ID pass\n");
        sb.append("# file_vaddr\tname\tparams_before\tparams_after\tlocals_before\tlocals_after\n");
        int improved = 0;
        for (Target t : targets) {
            boolean gain = t.paramsAfter > t.paramsBefore || t.localsAfter > t.localsBefore;
            if (gain) {
                improved++;
            }
            sb.append(String.format("0x%06x\t%s\t%d\t%d\t%d\t%d%s\n",
                    t.fileAddr, t.name, t.paramsBefore, t.paramsAfter,
                    t.localsBefore, t.localsAfter, gain ? "\tGAIN" : ""));
        }
        sb.append(String.format("# %d/%d functions gained parameters or locals\n",
                improved, targets.size()));
        String text = sb.toString();
        println(text);
        if (outFile != null) {
            try (PrintWriter w = new PrintWriter(new FileWriter(outFile))) {
                w.print(text);
            }
        }
    }
}