// Applies Visual C++ 7.0 name mangling to the symbols within the selected
// address range (or the whole program if nothing is selected).
//
// Be aware that the mangling implementation is only partial.
//
// @category Symbol

import ghidra.app.script.GhidraScript;
import ghidra.program.model.listing.Data;
import ghidra.program.model.listing.Function;
import ghidra.program.model.data.DataType;
import ghidra.program.model.data.Enum;
import ghidra.program.model.data.IntegerDataType;
import ghidra.program.model.data.VoidDataType;
import ghidra.program.model.symbol.Reference;
import ghidra.program.model.symbol.SourceType;
import ghidra.program.model.symbol.Symbol;
import ghidra.program.model.symbol.SymbolIterator;

import java.util.Arrays;
import java.util.Collections;
import java.util.List;


public class MSVC7Mangle extends GhidraScript{
	@Override
	public void run() throws Exception {
		final SymbolIterator iter = currentProgram.getSymbolTable()
		                                          .getPrimarySymbolIterator(currentSelection, true);

		while (iter.hasNext() && !monitor.isCancelled()) {
			final Symbol s = iter.next();

			switch (s.getObject()) {
				case Function f -> demangleFn(f);
				case Data     d -> demangleData(s);
				default         -> {}
			}
		}
	}

	private void demangleFn(final Function f) throws Exception {
		// Gather everything needed for mangling
		final List<String> name  = Arrays.asList(f.getName(true)
		                                          .split("::"));
		Collections.reverse(name);
		final String       callc = f.getCallingConventionName();
		final DataType     ret   = f.getReturnType();
		final DataType[]   args  = Arrays.stream(f.getSignature(true)
		                                          .getArguments())
		                                 .map(x -> x.getDataType())
		                                 .toArray(DataType[]::new);

		// Construct mangled name
		final String mangled =
			"?" + String.join("@", name) + "@@" +
			switch (callc) {
				case "__cdecl"    -> "YA";
				case "__thiscall" -> isVirtual(f) ? "UAE" :
				                                    "QAE";
				case "__fastcall" -> "";    // TODO
				default           -> throw new Exception(
					"Need to specify calling convention"
				);
			} +
			mangleType(ret) +
			mangleArgs(args) +
			"Z";

		f.setName(mangled, SourceType.USER_DEFINED);
	}

	private void demangleData(final Symbol s) {
		// TODO
		printf("TODO: data symbol \"%s\"\n", s.getName(true));
	}

	private boolean isVirtual(final Function f) {
	/* Attempt to determine whether a method is virtual
	We essentially try to figure out if any references are from a vtable.
	*/
		final Reference[] refs = getReferencesTo(f.getEntryPoint());

		// TODO

		return false;
	}

	private String mangleType(final DataType t) throws Exception {
	/* Mangle a data type in a function name */
		return switch(t) {
			case Enum            e -> "W4" + e.getName() + "@@";
			case IntegerDataType x -> "H";
			case VoidDataType    x -> "X";
			default                -> throw new Exception(
				"Unhandled data type \"" + t.toString() + "\""
			);
		};
	}

	private String mangleArgs(final DataType[] args) throws Exception {
	/* Mangle the arguments for a function */
		if (args.length == 0) return "X";
		else {
			String encoded = "";
			for (int i = 0; i < args.length; i++)
				encoded += mangleType(args[i]);
			return encoded + "@";
		}
	}
}