// SPDX-License-Identifier: BSD-3-Clause
//
// BatchDecompile.java - Ghidra post-script: decompile every function whose
// entry address lies in [start, end) and write the C output plus a
// machine-readable index.
//
// Designed for a very large stripped vendor library (~100k functions, 70 MB).
// The caller splits the code address range into chunks and runs one
// analyzeHeadless process per chunk (all with -readOnly against the same
// analyzed project), which parallelizes and makes the run resumable: a chunk
// whose output already exists is skipped.
//
// Usage (analyzeHeadless, on an already analyzed program):
//   -process <program> -readOnly -noanalysis \
//     -postScript BatchDecompile.java <startAddrHex> <endAddrHex> <outdir> [timeoutSecs]
//
// Outputs:
//   <outdir>/buckets/bucket_<start>_<end>.c      decompiled C, functions in order
//   <outdir>/index/bucket_<start>_<end>.tsv      addr  size  name  file  line  status
//
// Decompiled output is vendor-derived analysis output: it stays in the
// version's gitignored run/ tree and is NEVER committed.

import ghidra.app.script.GhidraScript;
import ghidra.app.decompiler.DecompInterface;
import ghidra.app.decompiler.DecompileResults;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.FunctionIterator;
import ghidra.program.model.listing.Listing;

import java.io.BufferedWriter;
import java.io.File;
import java.io.FileWriter;
import java.io.PrintWriter;

public class BatchDecompile extends GhidraScript {

    private static int countLines(String s) {
        int n = 0;
        for (int i = 0; i < s.length(); i++) {
            if (s.charAt(i) == '\n') n++;
        }
        return n;
    }

    @Override
    public void run() throws Exception {
        String[] args = getScriptArgs();
        if (args.length < 3) {
            println("usage: BatchDecompile.java <startAddrHex> <endAddrHex> <outdir> [timeoutSecs]");
            return;
        }
        long start = Long.parseUnsignedLong(args[0].replaceFirst("^0x", ""), 16);
        long end = Long.parseUnsignedLong(args[1].replaceFirst("^0x", ""), 16);
        String outdir = args[2];
        int timeout = args.length > 3 ? Integer.parseInt(args[3]) : 120;

        File buckets = new File(outdir, "buckets");
        File indexDir = new File(outdir, "index");
        buckets.mkdirs();
        indexDir.mkdirs();
        String tag = String.format("%08x_%08x", start, end);
        File outFile = new File(buckets, "bucket_" + tag + ".c");
        File idxFile = new File(indexDir, "bucket_" + tag + ".tsv");

        if (outFile.isFile() && idxFile.isFile()) {
            println("[BatchDecompile] chunk " + tag + " already done, skipping");
            return;
        }

        Listing listing = currentProgram.getListing();
        long imageBase = currentProgram.getImageBase().getOffset();
        DecompInterface dif = new DecompInterface();
        dif.toggleCCode(true);
        dif.toggleSyntaxTree(false);
        dif.openProgram(currentProgram);

        int n = 0, fail = 0, skip = 0, lines = 0;

        try (PrintWriter idx = new PrintWriter(new BufferedWriter(new FileWriter(idxFile)));
             PrintWriter out = new PrintWriter(new BufferedWriter(new FileWriter(outFile)))) {

            String banner = "/* CUDA userspace driver decompilation - Open-IOV project\n"
                    + " * program: " + currentProgram.getName() + "\n"
                    + " * range:   [0x" + Long.toHexString(start) + ", 0x"
                    + Long.toHexString(end) + ")\n"
                    + " * Machine-generated analysis output; not part of the public tree.\n"
                    + " */\n";
            out.print(banner);
            lines += countLines(banner);
            idx.println("addr\tfile_addr\tsize\tname\tfile\tline\tstatus");

            FunctionIterator fit = listing.getFunctions(true);
            while (fit.hasNext() && !monitor.isCancelled()) {
                Function f = fit.next();
                if (f.isExternal()) continue;
                long addr = f.getEntryPoint().getOffset();
                if (addr < start || addr >= end) continue;

                long size = f.getBody().getNumAddresses();
                String header = String.format(
                        "\n/* ------------------------------------------------------------------"
                        + "----------\n"
                        + " * %s\n"
                        + " * addr=0x%08x file_addr=libcuda+0x%x size=%d\n"
                        + " *-------------------------------------------------------------------"
                        + "-------- */\n",
                        f.getName(), addr, addr - imageBase, size);
                out.print(header);
                lines += countLines(header);

                String status = "ok";
                String code;
                if (size > 200000) {
                    status = "skipped-huge";
                    skip++;
                    code = "/* skipped: body too large (" + size + " bytes) */\n";
                } else {
                    DecompileResults res = dif.decompileFunction(f, timeout, monitor);
                    if (res != null && res.decompileCompleted()) {
                        code = res.getDecompiledFunction().getC();
                    } else {
                        status = "failed";
                        fail++;
                        code = "/* decompilation failed */\n";
                    }
                }
                out.print(code);
                out.print("\n");
                lines += countLines(code) + 1;
                idx.printf("0x%08x\t0x%08x\t%d\t%s\t%s\t%d\t%s%n",
                        addr, addr - imageBase, size, f.getName(),
                        outFile.getName(), lines - countLines(code), status);
                n++;
            }
        }

        println(String.format("[BatchDecompile] %s: functions=%d failed=%d skipped=%d -> %s",
                tag, n, fail, skip, outFile.getAbsolutePath()));
    }
}