// Headless post-analysis script: decompiles every function in the current
// program and writes C pseudo-code grouped by namespace (class) into
// <outDir>/<program>/src/<Namespace>.c, plus a functions.tsv index and a
// strings.tsv dump.
//
// Usage (via analyzeHeadless): -postScript ExportDecompiled.java <outDir> [<baseDir>]
//@category Export

import java.io.*;
import java.nio.charset.StandardCharsets;
import java.nio.file.*;
import java.util.*;
import java.util.concurrent.ConcurrentHashMap;

import generic.concurrent.GThreadPool;
import ghidra.app.decompiler.*;
import ghidra.app.decompiler.parallel.*;
import ghidra.app.script.GhidraScript;
import ghidra.program.model.data.StringDataInstance;
import ghidra.program.model.listing.*;
import ghidra.program.model.symbol.Namespace;
import ghidra.program.util.DefinedDataIterator;
import ghidra.util.task.TaskMonitor;

public class ExportDecompiled extends GhidraScript {

	private static final int TIMEOUT_SECS = 120;

	@Override
	protected void run() throws Exception {
		String[] args = getScriptArgs();
		Path out = Paths.get(args.length > 0 ? args[0] : "decomp")
				.resolve(outputName().replaceAll("[^A-Za-z0-9._/-]", "_"));
		Path src = out.resolve("src");
		Files.createDirectories(src);

		Map<String, List<String[]>> byNs = new ConcurrentHashMap<>();
		List<Function> funcs = new ArrayList<>();
		for (Function f : currentProgram.getFunctionManager().getFunctions(true)) {
			if (!f.isThunk() && !f.isExternal()) {
				funcs.add(f);
			}
		}
		println("Decompiling " + funcs.size() + " functions");

		DecompileConfigurer cfg = d -> {
			DecompileOptions opts = new DecompileOptions();
			opts.grabFromProgram(currentProgram);
			d.setOptions(opts);
			d.toggleCCode(true);
			d.toggleSyntaxTree(false);
			d.setSimplificationStyle("decompile");
		};
		DecompilerCallback<Void> cb = new DecompilerCallback<>(currentProgram, cfg) {
			@Override
			public Void process(DecompileResults r, TaskMonitor m) throws Exception {
				Function f = r.getFunction();
				String code = r.decompileCompleted()
						? r.getDecompiledFunction().getC()
						: "/* decompilation failed: " + r.getErrorMessage() + " */\n";
				String header = "// " + f.getEntryPoint() + "  " + f.getName(true) + "\n";
				byNs.computeIfAbsent(nsName(f.getParentNamespace()),
					k -> Collections.synchronizedList(new ArrayList<>()))
						.add(new String[] { f.getEntryPoint().toString(), header + code });
				return null;
			}
		};
		cb.setTimeout(TIMEOUT_SECS);
		// Each worker is a separate native decompiler process; one per core ran the
		// machine out of memory on steamclient.so, so cap it (DECOMP_THREADS, default 8 if unset).
		String threads = System.getenv().getOrDefault("DECOMP_THREADS", "8");
		GThreadPool.getSharedThreadPool("Parallel Decompiler").setMaxThreadCount(Integer.parseInt(threads));
		ParallelDecompiler.decompileFunctions(cb, funcs, monitor);
		cb.dispose();

		try (BufferedWriter idx = Files.newBufferedWriter(out.resolve("functions.tsv"), StandardCharsets.UTF_8)) {
			idx.write("address\tnamespace\tname\tsignature\n");
			for (Function f : funcs) {
				idx.write(f.getEntryPoint() + "\t" + nsName(f.getParentNamespace()) + "\t" + f.getName() + "\t"
						+ f.getSignature().getPrototypeString().replace('\t', ' ') + "\n");
			}
		}

		for (Map.Entry<String, List<String[]>> e : byNs.entrySet()) {
			List<String[]> list = new ArrayList<>(e.getValue());
			list.sort(Comparator.comparing(a -> a[0]));
			String file = e.getKey().replaceAll("[^A-Za-z0-9._-]", "_");
			if (file.length() > 180) {
				file = file.substring(0, 160) + "_" + Integer.toHexString(e.getKey().hashCode());
			}
			try (BufferedWriter w = Files.newBufferedWriter(src.resolve(file + ".c"), StandardCharsets.UTF_8)) {
				for (String[] fn : list) {
					w.write(fn[1]);
					w.write("\n");
				}
			}
		}

		try (BufferedWriter w = Files.newBufferedWriter(out.resolve("strings.tsv"), StandardCharsets.UTF_8)) {
			for (Data d : DefinedDataIterator.byDataInstance(currentProgram, StringDataInstance::isString)) {
				StringDataInstance s = StringDataInstance.getStringDataInstance(d);
				String v = s.getStringValue();
				if (v != null) {
					w.write(d.getAddress() + "\t" + v.replace("\\", "\\\\").replace("\n", "\\n").replace("\t", "\\t") + "\n");
				}
			}
		}
		println("Wrote " + byNs.size() + " namespace files to " + src);
	}

	// <parent dir>_<binary>, e.g. linux64_steamclient.so, so 32/64-bit builds don't collide.
	// With a second script arg (a base dir), mirrors the binary's path relative to it
	// instead, for trees like runtimes where parent dir names repeat.
	private String outputName() {
		Path exe = Paths.get(currentProgram.getExecutablePath());
		String[] args = getScriptArgs();
		if (args.length > 1 && exe.startsWith(Paths.get(args[1]))) {
			return Paths.get(args[1]).relativize(exe).toString();
		}
		Path parent = exe.getParent();
		return (parent == null ? "" : parent.getFileName() + "_") + exe.getFileName();
	}

	private static String nsName(Namespace ns) {
		return (ns == null || ns.isGlobal()) ? "_global" : ns.getName(true);
	}
}