//
//@author PhilippTheCat
//@category witcher
//@keybinding
//@menupath
//@toolbar

import ghidra.app.cmd.data.rtti.Rtti4Model;
import ghidra.app.cmd.data.rtti.VfTableModel;
import ghidra.app.util.datatype.microsoft.DataValidationOptions;
import ghidra.program.model.address.Address;
import ghidra.program.model.data.InvalidDataTypeException;
import ghidra.program.model.lang.UndefinedValueException;
import ghidra.program.model.listing.Data;
import ghidra.program.model.symbol.Symbol;
import ghidra.program.model.symbol.SymbolIterator;


import java.io.BufferedWriter;
import java.io.FileWriter;
import java.util.ArrayList;
import java.util.List;

public class DumpVTables extends GhidraScript {
    class FunctionInfo {
        public String name;
        public Address address;

    }


    class RTTITypeInfo {
        public String name;
        public Address address;
        List<String> baseClasses;
        List<FunctionInfo> functionInfos = new ArrayList<>();

        @Override
        public String toString() {
            String out = String.format("\t<RTTITypeInfo address=\"0x%x\" name=\"%s\">\n", address.getOffset(), sanitize(name));

            out += "\t\t<BaseClasses>\n";

            for (String baseClass : baseClasses) {
                out += String.format("\t\t\t<BaseClass>%s<BaseClass>\n", sanitize(baseClass));
            }

            out += "\t\t</BaseClasses>\n";

            out += "\t\t<Functions>\n";

            for (FunctionInfo functionInfo : functionInfos) {
                out += String.format("\t\t\t<Function address=\"0x%x\"/>\n", functionInfo.address.getOffset());
            }

            out += "\t\t</Functions>\n";

            out += "\t</RTTITypeInfo>\n";
            return out;
        }
    }

    private String sanitize(String baseClass) {
        return baseClass.replace("<", "&gt;").replace(">", "&gt;");
    }

    @Override
    protected void run() throws Exception {
        BufferedWriter rttiWriter = new BufferedWriter(new FileWriter("C:\\Users\\chatz\\Desktop\\dev\\cpp\\WitcherMod\\nativeRtti.xml"));

        SymbolIterator vftable = currentProgram.getSymbolTable().getSymbols("vftable");

        int vftablesTotal = 0;
        while (vftable.hasNext()) {
            Symbol table = vftable.next();


            monitor.checkCanceled();
            vftablesTotal += 1;


        }

        vftable = currentProgram.getSymbolTable().getSymbols("vftable");

        monitor.initialize(vftablesTotal);
        rttiWriter.write("<RTTITypeInfos>\n");

        int vftables = 0;
        while (vftable.hasNext()) {
            Symbol table = vftable.next();

            RTTITypeInfo rttiTypeInfo = exportVFTable(table);

            rttiWriter.write(rttiTypeInfo.toString());


            monitor.checkCanceled();
            monitor.incrementProgress(1);
            vftables += 1;

            if (vftables % 100 == 0) {
                printf("dumped %d vftables", vftables);
            }
        }
        rttiWriter.write("</RTTITypeInfos>");


        rttiWriter.close();
    }

    private RTTITypeInfo exportVFTable(Symbol table) throws InvalidDataTypeException, UndefinedValueException {

        Address rttiCompletePointer = table.getAddress().add(-8);
        Data dataAt = getDataAt(rttiCompletePointer);

        Class<?> valueClass = dataAt.getValueClass();

        if (valueClass == Address.class) {
            RTTITypeInfo rttiTypeInfo = new RTTITypeInfo();
            Rtti4Model rtti4Model = new Rtti4Model(currentProgram, (Address) dataAt.getValue(), new DataValidationOptions());

            rttiTypeInfo.address = (Address) dataAt.getValue();
            rttiTypeInfo.name = rtti4Model.getRtti0Model().getDescriptorName();
            rttiTypeInfo.baseClasses = rtti4Model.getBaseClassTypes();

            VfTableModel vfTableModel = new VfTableModel(currentProgram, table.getAddress(), new DataValidationOptions());

            for (int i = 0; i < vfTableModel.getElementCount(); i++) {
                Address virtualFunctionPointer = vfTableModel.getVirtualFunctionPointer(i);
                FunctionInfo functionInfo = new FunctionInfo();
                functionInfo.address = virtualFunctionPointer;
                rttiTypeInfo.functionInfos.add(functionInfo);
            }


            return rttiTypeInfo;
        }
        return null;
    }


}