// Creates classes out namespaces with matching structs, and if they have a
// vtable, sets the calling convention of the contained function typdefs to
// __thiscall.
//
// For vtables to be found, they must be defined as structs with names ending
// in "Vtbl" and be pointed to by the first member of a class struct.
//
// @category Data Types

import ghidra.app.script.GhidraScript;
import ghidra.app.util.NamespaceUtils;
import ghidra.program.database.data.DataTypeUtilities;
import ghidra.program.model.data.DataTypeComponent;
import ghidra.program.model.data.FunctionDefinition;
import ghidra.program.model.data.Pointer;
import ghidra.program.model.data.Structure;
import ghidra.program.model.listing.GhidraClass;
import ghidra.program.model.symbol.Namespace;
import ghidra.program.model.symbol.Symbol;


public class ClassFixup extends GhidraScript {
	@Override
	public void run() throws Exception {
		fixInNamespace(currentProgram.getGlobalNamespace());
	}

	private void fixInNamespace(final Namespace parent) throws Exception {
		for (final Symbol s : currentProgram.getSymbolTable()
		                                    .getChildren(parent.getSymbol()))
		if (s.getObject() instanceof Namespace ns) switch (ns.getType()) {
			case Namespace.Type.NAMESPACE:
				if (shouldBeClass(ns)) {
					println("Converting \"" + ns.getName(true) + "\" to class...");
					NamespaceUtils.convertNamespaceToClass(ns);

					// Re-fetch namespace to get its new class version
					ns = getNamespace(parent, ns.getName());
				} else {
					fixInNamespace(ns);
					break;
				}
			case Namespace.Type.CLASS: // fallthrough
				// Any ns that makes it here should be a class
				// (can't bind via pattern match because of the
				// fallthrough)
				final GhidraClass cls = (GhidraClass)ns;

				// Fix up methods if first member is a vtable
				if (
					DataTypeUtilities.findExistingClassStruct(
						currentProgram.getDataTypeManager(),
						cls
					)                        instanceof Structure         struct      &&
				        struct.getDataTypeAt(0)  instanceof DataTypeComponent memberType  &&
				        memberType.getDataType() instanceof Pointer           ptrType     &&
					ptrType   .getDataType() instanceof Structure         vtblT_maybe &&
					vtblT_maybe.getName().endsWith("Vtbl")
				) fixMethods(vtblT_maybe);
				break;

			default: continue;
		}
	}

	private boolean shouldBeClass(final Namespace ns) throws Exception {
	/* Determine if a namespace should be converted to a class
	The heuristic is whether Ghidra can find a struct that it would be
	linked with if it was a class.  The process to do so is admittedly a
	bit byzantine.
	*/
		// Change name so the dummy class won't have a conflict
		final String name = ns.getName();
		ns.getSymbol()
		  .setName("__" + ns.getName(), ns.getSymbol().getSource());

		// Create a dummy class to check for a matching struct without
		// converting the namespace (as classes can't be reverted)
		final GhidraClass cls = createClass(ns.getParentNamespace(), name);

		// Record whether Ghidra could find a matching struct
		final boolean ret = DataTypeUtilities.findExistingClassStruct(
			currentProgram.getDataTypeManager(),
			cls
		) != null;

		// Clean up
		cls.getSymbol().delete();
		ns.getSymbol().setName(name, ns.getSymbol().getSource());

		return ret;
	}

	private void fixMethods(final Structure vtbl) throws Exception {
	/* Set all the calling conventions in a vtable to __thiscall */
		for (final DataTypeComponent member : vtbl.getComponents())
		if (
			member.getDataType() instanceof Pointer            ptr    &&
			ptr   .getDataType() instanceof FunctionDefinition method &&
			!method.getCallingConventionName().equals("__thiscall")
		) {
			println("Fixing calling convention of \"" + method.getDataTypePath().toString() + "\"...");
			method.setCallingConvention("__thiscall");
		}
	}
}