// Test script
import java.io.File;
import java.io.PrintWriter;
import java.net.URL;
import java.util.ArrayList;
import java.util.Dictionary;
import java.util.Hashtable;
import java.util.Iterator;
import java.util.List;

import ghidra.app.decompiler.DecompInterface;
import ghidra.app.decompiler.DecompileResults;
import ghidra.app.script.GhidraScript;
import ghidra.app.script.GhidraScriptUtil;
import ghidra.program.model.address.GlobalSymbol;
import ghidra.program.model.data.DataType;
import ghidra.program.model.listing.Function;
import ghidra.program.model.listing.Variable;
import ghidra.program.model.pcode.HighSymbol;
import ghidra.program.model.symbol.SourceType;
import ghidra.program.model.symbol.Symbol;

public class DecompileC extends GhidraScript {
	private static final String OUTPUT_DIR = "game_re";
	private static final int TIMEOUT = 10000;

	// Auto rename invalid symbols
	private static final boolean AUTO_RENAME_SYMBOLS = true;

	@Override
	public void run() throws Exception {
		if (currentProgram == null) {
			return;
		}

		// Make sure to create OUTPUT_PATH
		File rootDir = new File(sourceFile.getAbsolutePath()).getParentFile().getParentFile();
		File outputDir = new File(rootDir, OUTPUT_DIR);

		if (!outputDir.exists()) {
			throw new Exception("Output directory does not exist: " + outputDir.getCanonicalPath());
		}

		File dirDecompAuto = new File(outputDir, "gh_auto");
		File dirDecompFix = new File(outputDir, "gh_fix");

		println("Output path: " + outputDir.getCanonicalPath());

		DecompInterface decomp = new DecompInterface();
		decomp.openProgram(currentProgram);

		List<Function> functions = new ArrayList<>();
		List<Variable> uniqueVars = new ArrayList<>();

		Iterator<Function> functionsIt = currentProgram.getFunctionManager().getFunctions(true).iterator();
		while (functionsIt.hasNext()) {
			Function function = functionsIt.next();
			String comment = function.getComment();
			if (comment != null && comment.contains("TODO")) {
				println("Function: " + function.getName() + " - " + comment);
			}
			functions.add(function);

			for (Variable var : function.getAllVariables()) {
				uniqueVars.add(var);
			}
		}

		Hashtable<String, HighSymbol> globalSymbols = new Hashtable<>();

		File functionList = new File(outputDir, "functions.txt");
		PrintWriter writer = new PrintWriter(functionList, "UTF-8");
		for (Function function : functions) {
			String fileName = sanitizeFunctionName(function.getName()) + ".c";

			File f1 = new File(dirDecompFix, fileName);
			if (f1.exists()) {
				println("Func " + function.getName() + " skipped (gh_fix)");
				continue;
			}

			File f0 = new File(dirDecompAuto, fileName);
			if (f0.exists()) {
				f0.delete();
			}

			DecompileResults decompRes = decomp.decompileFunction(function, TIMEOUT, monitor);
			PrintWriter writer2 = new PrintWriter(f0, "UTF-8");
			writer2.println(decompRes.getDecompiledFunction().getC());
			writer2.close();

			Iterator<HighSymbol> smyIt = decompRes.getHighFunction().getGlobalSymbolMap().getSymbols();
			while (smyIt.hasNext()) {
				HighSymbol gsym = smyIt.next();
				if (globalSymbols.containsKey(gsym.getName()))
					continue;
				println("GLOBAL: " + gsym.getName());
				String sanitizedName = sanitizeFunctionName(gsym.getName());
				if (!sanitizedName.equals(gsym.getName())) {
					if (AUTO_RENAME_SYMBOLS) {
						println("Renaming global symbol: " + gsym.getName() + " -> " + sanitizedName);
						Symbol symbol = gsym.getSymbol();// currentProgram.getSymbolTable().getSymbol(gsym.getName());
						if (symbol != null) {
							try {
								symbol.setName(sanitizedName, SourceType.USER_DEFINED);
								println("Renamed global symbol: " + gsym.getName() + " -> " + sanitizedName);
							} catch (Exception e) {
								println("Error renaming symbol: " + e.getMessage());
							}
						} else {
							println("Could not find symbol to rename: " + gsym.getName());
						}
					}
					println("Invalid global symbol name: " + gsym.getName() + " - " + function.getName());
				}
				globalSymbols.put(gsym.getName(), gsym);
			}
		}
		writer.close();

		File uniqueVarsList = new File(outputDir, "unique_vars.txt");
		writer = new PrintWriter(uniqueVarsList, "UTF-8");
		for (Variable var : uniqueVars) {
			writer.println(var.getName());
		}
		writer.close();

		File globalSymbolsListH = new File(outputDir, "global.h");
		File globalSymbolsListC = new File(outputDir, "global.c");
		PrintWriter hwriter = new PrintWriter(globalSymbolsListH, "UTF-8");
		PrintWriter cwriter = new PrintWriter(globalSymbolsListC, "UTF-8");
		for (HighSymbol sym : globalSymbols.values()) {
			DataType dt = sym.getDataType();
			hwriter.println("extern " + dt.toString() + " " + sym.getName() + ";");
			cwriter.println(sym.getDataType().toString() + " " + sym.getName() + " {}");
		}
		hwriter.close();
		cwriter.close();
	}

	String sanitizeFunctionName(String name) {
		return name.replaceAll("[^a-zA-Z0-9_]", "_");
	}
}