// Headless script: attribute functions to their original .cpp via embedded assert source-paths
// (d:\tesla_bt\bt\<file>.cpp), then decompile the BattleTech game-logic functions and write them
// grouped per source file into decomp/recovered/. Engine (munga\) is summarized, not exported.
import ghidra.app.script.GhidraScript;
import ghidra.app.decompiler.*;
import ghidra.program.model.address.Address;
import ghidra.program.model.data.StringDataInstance;
import ghidra.program.model.listing.*;
import ghidra.program.model.symbol.*;
import java.io.*;
import java.util.*;

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

        // 1) map function -> set of source files it references (via assert path strings)
        Map<Function, Set<String>> funcFiles = new HashMap<>();
        Map<String, Set<Function>> fileFuncs = new TreeMap<>();
        ReferenceManager rm = currentProgram.getReferenceManager();
        FunctionManager fm = currentProgram.getFunctionManager();
        Listing listing = currentProgram.getListing();

        int strCount = 0;
        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;
            strCount++;
            String file = s.substring(s.lastIndexOf('\\') + 1);
            // classify subdir
            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) continue;
                funcFiles.computeIfAbsent(f, k -> new HashSet<>()).add(tag);
                fileFuncs.computeIfAbsent(tag, k -> new LinkedHashSet<>()).add(f);
            }
        }
        println("assert source-path strings: " + strCount);
        println("source files attributed: " + fileFuncs.size());

        // 2) decompile + export functions for BT game modules (bt/ and bt_l4/)
        DecompInterface dec = new DecompInterface();
        dec.openProgram(currentProgram);

        PrintWriter index = new PrintWriter(new FileWriter(new File(outDir, "_index.txt")));
        for (Map.Entry<String, Set<Function>> e : fileFuncs.entrySet()) {
            String tag = e.getKey();
            int n = e.getValue().size();
            index.println(tag + "  -> " + n + " functions");
            boolean isGame = tag.startsWith("bt/") || tag.startsWith("bt_l4/");
            if (!isGame) continue; // engine/HAL we already have as source — skip export

            String safe = tag.replace('/', '_');
            PrintWriter out = new PrintWriter(new FileWriter(new File(outDir, safe + ".c")));
            out.println("// Recovered decompilation of functions attributed to " + tag);
            out.println("// (from BTL4OPT.EXE via embedded assert paths). Pseudocode — reconstruct against the header.\n");
            int ok = 0;
            for (Function f : e.getValue()) {
                if (monitor.isCancelled()) break;
                DecompileResults r = dec.decompileFunction(f, 60, monitor);
                if (r != null && r.decompileCompleted()) {
                    DecompiledFunction df = r.getDecompiledFunction();
                    if (df != null) {
                        out.println("/* ---- @ " + f.getEntryPoint() + " ---- */");
                        out.println(df.getC());
                        out.println();
                        ok++;
                    }
                }
            }
            out.close();
            println("exported " + tag + ": " + ok + "/" + n + " functions");
        }
        index.close();
        dec.dispose();
        println("DONE. Output in " + outDir.getAbsolutePath());
    }
}