// Full decompile: every function (sharded), vtable recovery, and assert-based file attribution.
// Outputs to decomp/recovered/: all/part_NNN.c, functions_index.tsv, vtables.tsv
import ghidra.app.script.GhidraScript;
import ghidra.app.decompiler.*;
import ghidra.program.model.address.*;
import ghidra.program.model.listing.*;
import ghidra.program.model.mem.*;
import ghidra.program.model.symbol.*;
import java.io.*;
import java.util.*;

public class ExportAll extends GhidraScript {
    public void run() throws Exception {
        File outDir = new File("C:\\git\\nick-games\\decomp\\recovered");
        outDir.mkdirs();
        File allDir = new File(outDir, "all");
        allDir.mkdirs();

        FunctionManager fm = currentProgram.getFunctionManager();
        ReferenceManager rm = currentProgram.getReferenceManager();
        Listing listing = currentProgram.getListing();
        Memory mem = currentProgram.getMemory();

        // 1) attribute functions to source files via embedded assert paths
        Map<Long, String> funcFile = new HashMap<>();
        DataIterator di = listing.getDefinedData(true);
        while (di.hasNext() && !monitor.isCancelled()) {
            Data d = di.next();
            if (d == null || !d.hasStringValue()) continue;
            Object v = d.getValue();
            if (v == null) continue;
            String s = v.toString().toLowerCase().replace('/', '\\');
            if (!s.contains("\\tesla_bt\\")) continue;
            if (!(s.endsWith(".cpp") || s.endsWith(".hpp"))) continue;
            String file = s.substring(s.lastIndexOf('\\') + 1);
            String dir = s.contains("\\bt_l4\\") ? "bt_l4" : s.contains("\\bt\\") ? "bt"
                       : s.contains("\\munga_l4\\") ? "munga_l4" : "munga";
            String tag = dir + "/" + file;
            for (Reference ref : rm.getReferencesTo(d.getAddress())) {
                Function f = fm.getFunctionContaining(ref.getFromAddress());
                if (f != null) funcFile.put(f.getEntryPoint().getOffset(), tag);
            }
        }
        println("attributed functions: " + funcFile.size());

        // 2) vtable recovery: runs of >=3 consecutive pointers to function entries in data memory
        PrintWriter vt = new PrintWriter(new FileWriter(new File(outDir, "vtables.tsv")));
        int vtCount = 0;
        for (MemoryBlock b : mem.getBlocks()) {
            if (!b.isInitialized() || b.isExecute()) continue;
            Address a = b.getStart(), end = b.getEnd();
            List<Long> run = new ArrayList<>();
            Address runStart = null;
            while (a.compareTo(end) < 0 && !monitor.isCancelled()) {
                long val;
                try { val = mem.getInt(a) & 0xffffffffL; } catch (Exception ex) { break; }
                Function f = null;
                try { f = fm.getFunctionAt(a.getNewAddress(val)); } catch (Exception ex) {}
                if (f != null) { if (run.isEmpty()) runStart = a; run.add(val); }
                else { if (run.size() >= 3) { writeVt(vt, runStart, run); vtCount++; } run.clear(); }
                a = a.add(4);
            }
            if (run.size() >= 3) { writeVt(vt, runStart, run); vtCount++; }
        }
        vt.close();
        println("vtables recovered: " + vtCount);

        // 3) decompile ALL functions, sharded 400/file
        DecompInterface dec = new DecompInterface();
        dec.openProgram(currentProgram);
        PrintWriter idx = new PrintWriter(new FileWriter(new File(outDir, "functions_index.tsv")));
        FunctionIterator it = fm.getFunctions(true);
        int count = 0, shard = 0;
        PrintWriter out = null;
        while (it.hasNext() && !monitor.isCancelled()) {
            Function f = it.next();
            if (count % 400 == 0) {
                if (out != null) out.close();
                out = new PrintWriter(new FileWriter(new File(allDir, String.format("part_%03d.c", shard++))));
            }
            count++;
            long off = f.getEntryPoint().getOffset();
            String tag = funcFile.getOrDefault(off, "?");
            idx.println(Long.toHexString(off) + "\t" + f.getBody().getNumAddresses() + "\t" + tag + "\t" + f.getName());
            try {
                DecompileResults r = dec.decompileFunction(f, 45, monitor);
                out.println("/* @" + f.getEntryPoint() + " file=" + tag + " name=" + f.getName() + " */");
                if (r != null && r.decompileCompleted() && r.getDecompiledFunction() != null)
                    out.println(r.getDecompiledFunction().getC());
                else
                    out.println("// <decompile failed>");
                out.println();
            } catch (Exception ex) {
                out.println("// <exception: " + ex.getMessage() + ">\n");
            }
            if (count % 500 == 0) println("decompiled " + count + " ...");
        }
        if (out != null) out.close();
        idx.close();
        dec.dispose();
        println("DONE. total functions: " + count + " (shards: " + shard + ")");
    }

    void writeVt(PrintWriter vt, Address start, List<Long> run) {
        StringBuilder sb = new StringBuilder(start.toString());
        for (Long v : run) sb.append("\t").append(Long.toHexString(v));
        vt.println(sb.toString());
    }
}