EGameTools/_IDAScripts/ExportClassH/HeaderGen.py

import os

from ExportClassH import ClassGen, Utils, Config
from ExportClassH.ClassDefs import ClassName, ParsedFunction, ParsedClassVar

processedClasses: set[str] = set()

currentAccess: str = "public"
def GenerateClassVarCode(classVar: ParsedClassVar) -> list[str]:
    """Generate code for a single class variable."""
    global currentAccess

    access: str = f"{classVar.access}:" if classVar.access else ""
    if currentAccess == classVar.access:
        access = "\t"
    else:
        currentAccess = classVar.access

    if classVar.varType:
        varType: str = Utils.ReplaceIDATypes(classVar.varType.fullName)
        varType = Utils.CleanType(varType)
        if varType:
            varType += " "
        varType = "GAME_IMPORT " + varType
    else:
        varType: str = ""
    
    classVarSig: str = f"{varType}{classVar.varName}"
    
    classVarLines: list[str] = []
    if access:
        classVarLines.append(access)
    if classVarSig:
        classVarLines.append(f"\t{classVarSig};")
    return classVarLines

def GenerateClassFuncCode(func: ParsedFunction, vtFuncIndex: int = 0) -> list[str]:
    """Generate code for a single class method."""
    global currentAccess

    access: str = f"{func.access}:" if func.access else ""
    if currentAccess == func.access:
        access = ""
    else:
        currentAccess = func.access

    const: str = " const" if func.const else ""
    stripped_vfunc: str = " = 0" if func.type == "stripped_vfunc" else ""

    if func.returnType:
        returnType: str = Utils.ReplaceIDATypes(func.returnType.fullName)
        returnType = Utils.CleanType(returnType)
        if returnType:
            if func.type == "basic_vfunc":
                returnType = returnType.removeprefix("virtual").strip()
            else:
                returnType += " "
    else:
        returnType: str = ""
    if func.type != "stripped_vfunc" and func.type != "basic_vfunc":
        returnType = "GAME_IMPORT " + returnType
    
    if func.params:
        paramsList: list[str] = []
        for param in func.params:
            paramsList.append(param.fullName)
        params: str = Utils.ReplaceIDATypes(", ".join(paramsList))
        params = Utils.CleanType(params)
        if params == "void":
            params = ""
    else:
        params: str = ""

    targetParams: str = ""
    if func.type == "basic_vfunc":
        targetParams = ClassGen.ExtractParamNames(params)
        targetParams = ", " + targetParams if targetParams else ""

    funcSig: str = f"{returnType}{func.funcName}({params}){const}{stripped_vfunc}" if func.type != "basic_vfunc" else f"VIRTUAL_CALL({vtFuncIndex}, {returnType}, {func.funcName}, ({params}){targetParams})"

    classFuncLines: list[str] = []
    if access:
        classFuncLines.append(access)
    if funcSig:
        classFuncLines.append(f"\t{funcSig};")
    return classFuncLines

def GenerateClassContent(allParsedElements: tuple[list[ParsedClassVar], list[ParsedFunction], list[ParsedFunction]]) -> list[str]:
    """
    Generate the content to be inserted into an existing header file.
    This is just the class members, not the full header with includes, etc.
    """
    global currentAccess
    
    parsedVars, parsedVtFuncs, parsedFuncs = allParsedElements
    if not parsedVars and not parsedVtFuncs and not parsedFuncs:
        return []
        
    firstVarOrFuncAccess = ""
    currentAccess = ""
    
    # Generate class content (just the members)
    contentLines = [f"#pragma region GENERATED by ExportClassToCPPH.py"]
    
    # Add class variables
    for classVar in parsedVars:
        if not firstVarOrFuncAccess:
            firstVarOrFuncAccess = classVar.access
        contentLines.extend(GenerateClassVarCode(classVar))
    
    # Add newline between sections if both exist
    if parsedVars and (parsedVtFuncs or parsedFuncs):
        contentLines.append("")
    
    # Add vtable functions
    for index, vTableFunc in enumerate(parsedVtFuncs):
        if not firstVarOrFuncAccess:
            firstVarOrFuncAccess = vTableFunc.access
        contentLines.extend(GenerateClassFuncCode(vTableFunc, index))
    
    # Add newline between sections if both exist
    if parsedVtFuncs and parsedFuncs:
        contentLines.append("")
    
    # Add regular functions
    for func in parsedFuncs:
        if not firstVarOrFuncAccess:
            firstVarOrFuncAccess = func.access
        contentLines.extend(GenerateClassFuncCode(func))
        
    contentLines.append("#pragma endregion")
    
    # Insert access specifier if needed
    if not firstVarOrFuncAccess:
        contentLines.insert(1, "public:")
    
    return contentLines

def GenerateClassDefinition(targetClass: ClassName, forwardDeclare: bool = False) -> list[str]:
    """Generate a class definition from a list of methods."""
    classDefLines: list[str] = [f"{targetClass.type} {targetClass.name}{' {' if not forwardDeclare else ';'}"]
    if not forwardDeclare:
        if targetClass.type == "class":
            classDefLines.append("public:")
        classDefLines.append("};")
    return classDefLines

def GenerateHeaderCode(targetClass: ClassName) -> list[str]:
    """Generate header code for a standard class (not nested)."""
    # Reset processed classes for this generation
    global processedClasses, forwardDeclarations
    processedClasses = set()
    forwardDeclarations = set()
    
    # Mark target class as processed to avoid self-dependency issues
    processedClasses.add(targetClass.namespacedClassedName)
    
    # Get all parsed elements for the target class
    allParsedElements = ClassGen.GetAllParsedClassVarsAndFuncs(targetClass)
    
    # Generate the target class definition
    classContent = GenerateClassContent(allParsedElements)
    classDefinition = GenerateClassDefinition(targetClass)

    fullClassDefinition = classDefinition
    if (classContent):
        fullClassDefinition = fullClassDefinition[:len(classDefinition) - 1] + classContent + fullClassDefinition[len(classDefinition) - 1:]

    # Wrap target class in namespace blocks if needed
    if targetClass.namespaces:
        namespaceCode = []
        indentLevel = ""
        
        # Opening namespace blocks with increasing indentation
        for namespace in targetClass.namespaces:
            namespaceCode.append(f"{indentLevel}namespace {namespace} {{")
            indentLevel += "\t"
        
        for cls in targetClass.classes:
            clsType: str = ClassGen.GetClassTypeFromParsedSigs(ClassName(targetClass.namespacedClassedName), allParsedElements)
            namespaceCode.append(f"{indentLevel}{clsType if clsType else 'class'} {cls} {{")
            indentLevel += "\t"

        indentedClassDefinition = [f"{indentLevel if '#pragma' not in line else ''}{line}" for line in fullClassDefinition]
        namespaceCode.extend(indentedClassDefinition)

        for cls in reversed(targetClass.classes):
            indentLevel = indentLevel[:-1]  # Remove one level of indentation
            namespaceCode.append(f"{indentLevel}}}")

        for namespace in reversed(targetClass.namespaces):
            indentLevel = indentLevel[:-1]  # Remove one level of indentation
            namespaceCode.append(f"{indentLevel}}}")
        
        fullClassDefinition = namespaceCode
    
    # Combine all parts of the header
    headerParts = ["#pragma once", r"#include <EGSDK\Imports.h>", ""]
    
    # Add the target class definition
    headerParts.extend(fullClassDefinition)
    
    return headerParts

def WriteHeaderToFile(targetClass: ClassName, headerCode: str, fileName: str) -> bool:
    """Write the generated header code to a file."""
    outputFolderPath: str = Config.HEADER_OUTPUT_PATH

    if targetClass.namespaces:
        # Create folder structure for namespaces
        classFolderPath: str = os.path.join(*targetClass.namespaces)
        outputFolderPath: str = os.path.join(outputFolderPath, classFolderPath)
        
        # Create directory if it doesn't exist
        os.makedirs(outputFolderPath, exist_ok=True)
        # Output file path is inside the namespace folder
        outputFilePath: str = os.path.join(classFolderPath, fileName)
    else:
        # Create directory if it doesn't exist
        os.makedirs(outputFolderPath, exist_ok=True)
        # No namespace, just save in current directory
        outputFilePath: str = fileName
    
    outputFilePath = os.path.join(Config.HEADER_OUTPUT_PATH, outputFilePath)

    try:
        with open(outputFilePath, 'w') as headerFile:
            headerFile.write(headerCode)
        print(f"Header file '{outputFilePath}' created successfully.")
        
        return True
    except Exception as e:
        print(f"Error writing header file '{outputFilePath}': {e}")
        return False
    
def ExportClassHeader(targetClass: ClassName):
    """
    Generate and save a C++ header file for the target class.
    Handles multiple levels of nested classes and also generates dependencies.
    """  
    global processedClasses

    # Skip if we've already processed this class
    if targetClass.namespacedClassedName in processedClasses:
        print(f"Already processed class {targetClass.namespacedClassedName}, skipping.")
        return
    
    # Add to processed set to prevent infinite recursion
    processedClasses.add(targetClass.namespacedClassedName)

    # Generate the header code
    headerCodeLines = GenerateHeaderCode(targetClass)
    if not headerCodeLines:
        print(f"No functions were found for class {targetClass.namespacedClassedName}, therefore will not generate.")
        return
    headerCode: str = "\n".join(headerCodeLines)
    
    WriteHeaderToFile(targetClass, headerCode, f"{targetClass.name}.h")