EGameTools/_IDAScripts/ExportClassH/HeaderGen.py

import os

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

def IsClassGenerable(className: ClassName) -> bool:
    """
    Check if a class has any parsable elements (class vars, vtable functions, regular functions).
    Returns True if the class is generable, False if it should be treated as a namespace.
    """
    parsedVars = ClassParser.GetParsedClassVars(className)
    parsedVtFuncs = ClassParser.GetParsedVTableFuncs(className) if className else []
    parsedFuncs = ClassParser.GetParsedFuncs(className)
    
    return len(parsedVars) > 0 or len(parsedVtFuncs) > 0 or len(parsedFuncs) > 0

def IdentifyClassHierarchy(targetClass: ClassName) -> list[tuple[ClassName, bool]]:
    """
    Given a class name with namespace/nested parts, identify which parts are classes
    and which are namespaces.
    
    Returns a list of tuples (ClassName, isClass) for each part of the hierarchy.
    """
    if not targetClass.namespaces:
        return [(targetClass, IsClassGenerable(targetClass))]
    
    hierarchy = []
    
    # Check each part of the namespace to see if it's a class
    currentNamespace = []
    for part in targetClass.namespaces:
        currentNamespace.append(part)
        partClass = ClassName("::".join(currentNamespace))
        
        isClass = IsClassGenerable(partClass)
        hierarchy.append((partClass, isClass))
    
    # Add the target class at the end
    hierarchy.append((targetClass, IsClassGenerable(targetClass)))
    
    return hierarchy

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

    access: str = f"{classVar.access}:\n\t" if classVar.access else "\t"
    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}"
    return f"{access}{classVarSig};"

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

    access: str = f"{func.access}:\n\t" if func.access else "\t"
    if currentAccess == func.access:
        access = "\t"
    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 = ClassParser.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})"
    return f"{access}{funcSig};"

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.append(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.append(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.append(GenerateClassFuncCode(func))
        
    contentLines.append("#pragma endregion")
    
    # Insert access specifier if needed
    if not firstVarOrFuncAccess:
        contentLines.insert(1, "public:")
    
    return contentLines

def GenerateClassDefinition(targetClass: ClassName, allParsedElements: tuple[list[ParsedClassVar], list[ParsedFunction], list[ParsedFunction]]) -> list[str]:
    """Generate a class definition from a list of methods."""
    parsedVars, parsedVtFuncs, parsedFuncs = allParsedElements
    if not parsedVars and not parsedVtFuncs and not parsedFuncs:
        return []
    
    classContent: list[str] = GenerateClassContent(allParsedElements)
    
    classLines: list[str] = []
    if classContent:
        classLines.extend(classContent)

    classLines.insert(0, f"{targetClass.type if targetClass.type else 'class'} {targetClass.name}")
    if classContent:
        classLines[0] = f"{classLines[0]} {{"
        if targetClass.type == "struct":
            classLines[1] = "public:"
        classLines.append("};")
    else:
        classLines[0] = f"{classLines[0]};"

    return classLines

def GenerateHeaderCode(targetClass: ClassName, allParsedElements: tuple[list[ParsedClassVar], list[ParsedFunction], list[ParsedFunction]]) -> list[str]:
    """Generate header code for a standard class (not nested)."""
    classDefinition = GenerateClassDefinition(targetClass, allParsedElements)
    if not classDefinition:
        return []
    
    # Wrap 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"

        indentedClassDefinition = [f"{indentLevel}{line}" for line in classDefinition]
        namespaceCode.extend(indentedClassDefinition)

        for namespace in reversed(targetClass.namespaces):
            indentLevel = indentLevel[:-1]  # Remove one level of indentation
            namespaceCode.append(f"{indentLevel}}}")
        
        classDefinition = namespaceCode
    
    # Combine all parts of the header
    headerParts = ["#pragma once", r"#include <EGSDK\Imports.h>", ""]
    headerParts.extend(classDefinition)
    
    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.
    """  
    # Get the parsed elements for the target class
    allParsedElements = ClassParser.GetAllParsedClassVarsAndFuncs(targetClass)
    
    # Generate the header code
    headerCodeLines = GenerateHeaderCode(targetClass, allParsedElements)
    if not headerCodeLines:
        print(f"No functions were found for class {targetClass.namespacedName}, therefore will not generate.")
        return
    headerCode: str = "\n".join(headerCodeLines)
    
    WriteHeaderToFile(targetClass, headerCode, f"{targetClass.name}.h")