codegen.cpp

#include "ast.h"
#include <iostream>
#include <fstream>
#include <string>
#include <vector>
#include <map>
#include <memory>
#include <llvm/IR/Verifier.h>
#include <llvm/IR/LegacyPassManager.h>
#include <llvm/Transforms/Scalar.h>
#include <llvm/Transforms/InstCombine/InstCombine.h>
#include <llvm/Transforms/Utils/Mem2Reg.h>
#include <llvm/Transforms/Scalar/GVN.h>
#include <llvm/Transforms/Utils/BasicBlockUtils.h>
#include <llvm/IR/PassManager.h>
#include <llvm/Passes/PassBuilder.h>
#include <llvm/Support/raw_ostream.h>
#include <llvm/Support/FileSystem.h>
#include <llvm/ExecutionEngine/ExecutionEngine.h>
#include <llvm/ExecutionEngine/GenericValue.h>
#include <llvm/Support/TargetSelect.h>
#include <llvm/Transforms/Utils/Cloning.h>

// Debug flag to enable verbose output
#define DEBUG_OUTPUT 1

// Helper function for debug output
void debugOutput(const std::string& message) {
    #if DEBUG_OUTPUT
    std::cout << "[DEBUG] " << message << std::endl;
    #endif
}

static void ensureTerminatorForBlock(llvm::BasicBlock *BB, llvm::Function *F, llvm::IRBuilder<> &Builder) {
    // If there's already a terminator, do nothing.
    if (BB->getTerminator()) return;

    Builder.SetInsertPoint(BB);
    if (F->getReturnType()->isVoidTy()) {
        Builder.CreateRetVoid();
    } else {
        Builder.CreateRet(llvm::ConstantInt::get(
            llvm::Type::getInt32Ty(F->getContext()), 0));
    }
}

// Helper functions for printf
static llvm::Function* createPrintfFunction(CodeGenContext& context) {
    std::vector<llvm::Type*> args;
    args.push_back(llvm::PointerType::get(context.context, 0));
    
    llvm::FunctionType* printfType = llvm::FunctionType::get(
        llvm::Type::getInt32Ty(context.context), args, true);
    
    llvm::Function* func = llvm::Function::Create(
        printfType, llvm::Function::ExternalLinkage, "printf", context.module.get());
    
    return func;
}

// Helper function to create allocas at the entry block of a function
llvm::AllocaInst* CodeGenContext::createEntryBlockAlloca(llvm::Function* function, 
                                                       const std::string& varName,
                                                       llvm::Type* type) {
    if (!function) {
        std::cerr << "Error: Null function in createEntryBlockAlloca" << std::endl;
        return nullptr;
    }
    
    llvm::IRBuilder<> TmpB(&function->getEntryBlock(), 
                         function->getEntryBlock().begin());
    return TmpB.CreateAlloca(type, nullptr, varName.c_str());
}

// Implementation of code generation for expressions
llvm::Value* IntLiteral::codegen(CodeGenContext& context) {
    return llvm::ConstantInt::get(llvm::Type::getInt32Ty(context.context), value);
}

llvm::Value* BoolLiteral::codegen(CodeGenContext& context) {
    return llvm::ConstantInt::get(llvm::Type::getInt1Ty(context.context), value ? 1 : 0);
}

llvm::Value* StringLiteral::codegen(CodeGenContext& context) {
    llvm::Constant* strConstant = llvm::ConstantDataArray::getString(context.context, value);
    
    // Create a global variable for the string constant
    llvm::GlobalVariable* globalStr = new llvm::GlobalVariable(
        *context.module, strConstant->getType(), true,
        llvm::GlobalValue::PrivateLinkage, strConstant, ".str");
    
    // Get a pointer to the beginning of the string
    llvm::Value* zero = llvm::ConstantInt::get(llvm::Type::getInt32Ty(context.context), 0);
    std::vector<llvm::Value*> indices;
    indices.push_back(zero);
    indices.push_back(zero);
    
    // Use the correct GEP instruction for LLVM 19
    return context.builder.CreateGEP(strConstant->getType(), globalStr, 
                                   llvm::ArrayRef<llvm::Value*>(indices), "strptr");
}

llvm::Value* VarExpr::codegen(CodeGenContext& context) {
    // Look up the variable in the named values map
    auto it = context.namedValues.find(name);
    if (it == context.namedValues.end() || !it->second) {
        std::cerr << "Unknown variable: " << name << std::endl;
        return nullptr;
    }
    
    llvm::AllocaInst* alloca = it->second;
    if (!alloca) {
        std::cerr << "Error: Null alloca for variable: " << name << std::endl;
        return nullptr;
    }
    
    // Use the correct load instruction for LLVM 19
    return context.builder.CreateLoad(alloca->getAllocatedType(), alloca, name.c_str());
}

llvm::Value* BinaryExpr::codegen(CodeGenContext& context) {
    if (!lhs || !rhs) {
        std::cerr << "Error: Null operand in binary expression" << std::endl;
        return nullptr;
    }
    
    llvm::Value* L = lhs->codegen(context);
    llvm::Value* R = rhs->codegen(context);
    
    if (!L || !R) {
        std::cerr << "Error: Failed to generate code for binary expression operands" << std::endl;
        return nullptr;
    }
    
    // Handle logical operations
    if (op == BinaryOp::AND || op == BinaryOp::OR) {
        // Convert operands to boolean if needed
        if (!L->getType()->isIntegerTy(1)) {
            L = context.builder.CreateICmpNE(
                L, llvm::ConstantInt::get(L->getType(), 0), "tobool");
        }
        
        if (!R->getType()->isIntegerTy(1)) {
            R = context.builder.CreateICmpNE(
                R, llvm::ConstantInt::get(R->getType(), 0), "tobool");
        }
        
        if (op == BinaryOp::AND) {
            return context.builder.CreateAnd(L, R, "andtmp");
        } else { // op == BinaryOp::OR
            return context.builder.CreateOr(L, R, "ortmp");
        }
    }
    
    // Check if we're dealing with string concatenation (string + anything)
    bool isStringL = L->getType()->isPointerTy();
    bool isStringR = R->getType()->isPointerTy();
    
    // If either operand is a string and we're doing addition, handle as string concatenation
    if (op == BinaryOp::ADD && (isStringL || isStringR)) {
        // Get or create the printf function for formatting
        llvm::Function* printfFunc = context.module.get()->getFunction("printf");
        if (!printfFunc) {
            printfFunc = createPrintfFunction(context);
        }
        
        // Get or create snprintf function for string formatting
        llvm::Function* snprintfFunc = context.module.get()->getFunction("snprintf");
        if (!snprintfFunc) {
            std::vector<llvm::Type*> args;
            args.push_back(llvm::PointerType::get(context.context, 0)); // char* buffer
            args.push_back(llvm::Type::getInt32Ty(context.context));    // size_t size
            args.push_back(llvm::PointerType::get(context.context, 0)); // const char* format
            
            llvm::FunctionType* snprintfType = llvm::FunctionType::get(
                llvm::Type::getInt32Ty(context.context), args, true);
            
            snprintfFunc = llvm::Function::Create(
                snprintfType, llvm::Function::ExternalLinkage, "snprintf", context.module.get());
        }
        
        // Create a buffer for the result
        llvm::AllocaInst* buffer = context.builder.CreateAlloca(
            llvm::ArrayType::get(llvm::Type::getInt8Ty(context.context), 1024),
            nullptr, "concat_buffer");
        
        std::vector<llvm::Value*> formatArgs;
        
        // Format string depends on the types of operands
        std::string formatStr;
        if (isStringL && isStringR) {
            formatStr = "%s%s"; // string + string
            formatArgs.push_back(L);
            formatArgs.push_back(R);
        } else if (isStringL && R->getType()->isIntegerTy()) {
            formatStr = "%s%d"; // string + int
            formatArgs.push_back(L);
            formatArgs.push_back(R);
        } else if (isStringR && L->getType()->isIntegerTy()) {
            formatStr = "%d%s"; // int + string
            formatArgs.push_back(L);
            formatArgs.push_back(R);
        } else {
            std::cerr << "Unsupported operand types for string concatenation" << std::endl;
            return nullptr;
        }
        
        // Create the format string
        llvm::Value* formatStrVal = context.builder.CreateGlobalStringPtr(formatStr);
        
        // Create the call to snprintf
        std::vector<llvm::Value*> snprintfArgs;
        snprintfArgs.push_back(context.builder.CreateBitCast(buffer, llvm::PointerType::get(context.context, 0)));
        snprintfArgs.push_back(llvm::ConstantInt::get(llvm::Type::getInt32Ty(context.context), 1024));
        snprintfArgs.push_back(formatStrVal);
        snprintfArgs.insert(snprintfArgs.end(), formatArgs.begin(), formatArgs.end());
        
        context.builder.CreateCall(snprintfFunc, snprintfArgs);
        
        // Return the buffer as a string
        return context.builder.CreateBitCast(buffer, llvm::PointerType::get(context.context, 0), "concat_result");
    }
    
    // For non-string operations, ensure both operands have the same type
    if (L->getType() != R->getType()) {
        debugOutput("Type mismatch in binary expression");
        
        // Try to convert types if possible (e.g., int to float or vice versa)
        if (L->getType()->isIntegerTy() && R->getType()->isIntegerTy()) {
            // If both are integers but different bit widths, convert to the larger one
            unsigned LBits = L->getType()->getIntegerBitWidth();
            unsigned RBits = R->getType()->getIntegerBitWidth();
            
            if (LBits < RBits) {
                L = context.builder.CreateIntCast(L, R->getType(), true, "cast");
            } else {
                R = context.builder.CreateIntCast(R, L->getType(), true, "cast");
            }
        } else {
            // For now, we don't handle other type conversions
            return nullptr;
        }
    }
    
    switch (op) {
        case BinaryOp::ADD:
            return context.builder.CreateAdd(L, R, "addtmp");
        case BinaryOp::SUB:
            return context.builder.CreateSub(L, R, "subtmp");
        case BinaryOp::MUL:
            return context.builder.CreateMul(L, R, "multmp");
        case BinaryOp::DIV:
            return context.builder.CreateSDiv(L, R, "divtmp");
        case BinaryOp::EQ:
            return context.builder.CreateICmpEQ(L, R, "eqtmp");
        case BinaryOp::NEQ:
            return context.builder.CreateICmpNE(L, R, "neqtmp");
        case BinaryOp::LT:
            return context.builder.CreateICmpSLT(L, R, "lttmp");
        case BinaryOp::GT:
            return context.builder.CreateICmpSGT(L, R, "gttmp");
        default:
            std::cerr << "Unknown binary operator" << std::endl;
            return nullptr;
    }
}

// Implementation of code generation for statements
llvm::Value* DeclStmt::codegen(CodeGenContext& context) {
    if (!initializer) {
        std::cerr << "Error: Null initializer in declaration statement" << std::endl;
        return nullptr;
    }
    
    // Check if variable already exists in current scope
    auto it = context.namedValues.find(name);
    if (it != context.namedValues.end()) {
        std::cerr << "Error: Variable '" << name << "' already defined in this scope" << std::endl;
        return nullptr;
    }
    
    llvm::Function* function = context.builder.GetInsertBlock()->getParent();
    if (!function) {
        std::cerr << "Error: Null function in declaration statement" << std::endl;
        return nullptr;
    }
    
    // Generate code for the initializer
    llvm::Value* initVal = initializer->codegen(context);
    if (!initVal) {
        std::cerr << "Error: Failed to generate initializer for " << name << std::endl;
        return nullptr;
    }
    
    // Create an alloca for this variable
    llvm::AllocaInst* alloca = context.createEntryBlockAlloca(
        function, name, initVal->getType());
    
    if (!alloca) {
        std::cerr << "Error: Failed to create alloca for " << name << std::endl;
        return nullptr;
    }
    
    // Store the initial value into the alloca
    context.builder.CreateStore(initVal, alloca);
    
    // Add the variable to the symbol table
    context.namedValues[name] = alloca;
    
    return alloca;
}

llvm::Value* AssignStmt::codegen(CodeGenContext& context) {
    if (!value) {
        std::cerr << "Error: Null value in assignment statement" << std::endl;
        return nullptr;
    }
    
    // Look up the variable in the named values map
    auto it = context.namedValues.find(name);
    if (it == context.namedValues.end() || !it->second) {
        std::cerr << "Unknown variable: " << name << std::endl;
        return nullptr;
    }
    
    llvm::AllocaInst* alloca = it->second;
    
    // Generate code for the value being assigned
    llvm::Value* val = value->codegen(context);
    if (!val) {
        std::cerr << "Error: Failed to generate value for assignment to " << name << std::endl;
        return nullptr;
    }
    
    // Store the value into the alloca
    context.builder.CreateStore(val, alloca);
    
    return val;
}

llvm::Value* IfStmt::codegen(CodeGenContext& context) {
    debugOutput("Generating code for if statement");
    
    if (!condition) {
        std::cerr << "Error: Null condition in if statement" << std::endl;
        return nullptr;
    }
    
    llvm::Value* condValue = condition->codegen(context);
    if (!condValue) {
        std::cerr << "Error: Failed to generate condition for if statement" << std::endl;
        return nullptr;
    }
    
    // Convert condition to a boolean (i1)
    if (!condValue->getType()->isIntegerTy(1)) {
        condValue = context.builder.CreateICmpNE(
            condValue, llvm::ConstantInt::get(condValue->getType(), 0), "ifcond");
    }
    
    llvm::Function* function = context.builder.GetInsertBlock()->getParent();
    if (!function) {
        std::cerr << "Error: Null function in if statement" << std::endl;
        return nullptr;
    }
    
    // Save the current block to continue from
    llvm::BasicBlock* currentBlock = context.builder.GetInsertBlock();
    
    // Create blocks for the then, else, and merge cases
    llvm::BasicBlock* thenBB = llvm::BasicBlock::Create(context.context, "then", function);
    llvm::BasicBlock* elseBB = elseBlock ? 
        llvm::BasicBlock::Create(context.context, "else") : nullptr;
    llvm::BasicBlock* mergeBB = llvm::BasicBlock::Create(context.context, "ifcont");
    
    // Create the conditional branch
    context.builder.CreateCondBr(condValue, thenBB, elseBB ? elseBB : mergeBB);
    
    // Emit the 'then' block
    context.builder.SetInsertPoint(thenBB);
    debugOutput("Setting insertion point to then block");
    
    if (thenBlock) {
        for (auto stmt : *thenBlock) {
            if (stmt) {
                stmt->codegen(context);
                
                // If we just generated a return statement, don't continue generating code
                if (dynamic_cast<ReturnStmt*>(stmt) && !context.builder.GetInsertBlock()->getTerminator()) {
                    break;
                }
            }
        }
    }
    
    // Create a branch to the merge block if there isn't already a terminator
    if (!context.builder.GetInsertBlock()->getTerminator()) {
        debugOutput("Adding branch from then block to merge block");
        context.builder.CreateBr(mergeBB);
    }
    
    // Emit the 'else' block if present
    if (elseBlock) {
        // Add the else block to the function
        elseBB->insertInto(function);
        context.builder.SetInsertPoint(elseBB);
        debugOutput("Setting insertion point to else block");
        
        for (auto stmt : *elseBlock) {
            if (stmt) {
                stmt->codegen(context);
                
                // If we just generated a return statement, don't continue generating code
                if (dynamic_cast<ReturnStmt*>(stmt) && !context.builder.GetInsertBlock()->getTerminator()) {
                    break;
                }
            }
        }
        
        // Create a branch to the merge block if there isn't already a terminator
        if (!context.builder.GetInsertBlock()->getTerminator()) {
            debugOutput("Adding branch from else block to merge block");
            context.builder.CreateBr(mergeBB);
        }
    }
    
    // Add the merge block to the function
    mergeBB->insertInto(function);
    
    // Set insertion point to merge block
    context.builder.SetInsertPoint(mergeBB);
    debugOutput("Setting insertion point to merge block");
    
    return llvm::Constant::getNullValue(llvm::Type::getInt32Ty(context.context));
}

llvm::Value* LoopStmt::codegen(CodeGenContext& context) {
    debugOutput("Generating code for loop statement");

    if (!count) {
        std::cerr << "Error: Null count in loop statement" << std::endl;
        return nullptr;
    }

    llvm::Function* function = context.builder.GetInsertBlock()->getParent();
    if (!function) {
        std::cerr << "Error: Null function in loop statement" << std::endl;
        return nullptr;
    }

    // Save the current insertion point (preheader)
    llvm::BasicBlock* preheaderBB = context.builder.GetInsertBlock();

    static int loopCounter = 0;
    loopCounter++;
    llvm::BasicBlock* loopBB  = llvm::BasicBlock::Create(context.context,
                                "loop" + std::to_string(loopCounter), function);
    llvm::BasicBlock* afterBB = llvm::BasicBlock::Create(context.context,
                                "afterloop" + std::to_string(loopCounter));

    // Generate the loop count
    llvm::Value* countValue = count->codegen(context);
    if (!countValue) {
        std::cerr << "Error: Failed to generate count for loop statement" << std::endl;
        return nullptr;
    }

    // Create and initialize counter variable to 0
    llvm::AllocaInst* counter = context.createEntryBlockAlloca(
        function, "loopcounter", llvm::Type::getInt32Ty(context.context));
    context.builder.CreateStore(
        llvm::ConstantInt::get(llvm::Type::getInt32Ty(context.context), 0), counter);

    // Branch from preheader to loop block
    if (!preheaderBB->getTerminator()) {
        context.builder.CreateBr(loopBB);
    }

    // Start insertion in loop body
    context.builder.SetInsertPoint(loopBB);
    debugOutput("Setting insertion point to loop block");

    // Generate the loop body statements
    if (body) {
        for (auto stmt : *body) {
            if (stmt) {
                stmt->codegen(context);
                
                // If we just generated a return statement, don't continue generating code
                if (dynamic_cast<ReturnStmt*>(stmt) && !context.builder.GetInsertBlock()->getTerminator()) {
                    break;
                }
            }
        }
    }

    // If we have a terminator (like a return), skip the rest of the loop logic
    if (context.builder.GetInsertBlock()->getTerminator()) {
        // Add the after block to the function and set insertion point
        afterBB->insertInto(function);
        context.builder.SetInsertPoint(afterBB);
        return llvm::Constant::getNullValue(llvm::Type::getInt32Ty(context.context));
    }

    // Increment the counter
    llvm::Value* currentCount = context.builder.CreateLoad(
            counter->getAllocatedType(), counter, "currentcount");
    llvm::Value* nextCount = context.builder.CreateAdd(
            currentCount, llvm::ConstantInt::get(
            llvm::Type::getInt32Ty(context.context), 1), "nextcount");
    context.builder.CreateStore(nextCount, counter);

    // Check loop condition: nextCount < countValue
    llvm::Value* endCond = context.builder.CreateICmpSLT(
            nextCount, countValue, "loopcond");

    // Conditional branch: back to loop or to after block
    context.builder.CreateCondBr(endCond, loopBB, afterBB);

    // Append the after block and set insertion point
    afterBB->insertInto(function);
    context.builder.SetInsertPoint(afterBB);
    debugOutput("Setting insertion point to afterloop block");

    return llvm::Constant::getNullValue(
        llvm::Type::getInt32Ty(context.context));
}

llvm::Value* PrintStmt::codegen(CodeGenContext& context) {
    if (!value) {
        std::cerr << "Error: Null value in print statement" << std::endl;
        return nullptr;
    }
    
    // Get or create the printf function
    llvm::Function* printfFunc = context.module.get()->getFunction("printf");
    if (!printfFunc) {
        printfFunc = createPrintfFunction(context);
    }
    
    if (!printfFunc) {
        std::cerr << "Error: Failed to create printf function" << std::endl;
        return nullptr;
    }
    
    // Special handling for BinaryExpr with string concatenation
    if (auto binExpr = dynamic_cast<BinaryExpr*>(value)) {
        if (binExpr->op == BinaryOp::ADD) {
            // Check if either operand is a string
            bool isStringLiteral = dynamic_cast<StringLiteral*>(binExpr->lhs) != nullptr || 
                                  dynamic_cast<StringLiteral*>(binExpr->rhs) != nullptr;
            
            if (isStringLiteral) {
                // Handle string concatenation in print statement
                std::vector<llvm::Value*> printfArgs;
                std::string formatStr;
                
                // Generate left operand
                llvm::Value* lhsVal = binExpr->lhs->codegen(context);
                if (!lhsVal) return nullptr;
                
                // Generate right operand
                llvm::Value* rhsVal = binExpr->rhs->codegen(context);
                if (!rhsVal) return nullptr;
                
                // Determine format string based on operand types
                if (lhsVal->getType()->isPointerTy() && rhsVal->getType()->isIntegerTy()) {
                    // String + Int
                    formatStr = "%s%d\n";
                    printfArgs.push_back(context.builder.CreateGlobalStringPtr(formatStr));
                    printfArgs.push_back(lhsVal);
                    printfArgs.push_back(rhsVal);
                } else if (lhsVal->getType()->isIntegerTy() && rhsVal->getType()->isPointerTy()) {
                    // Int + String
                    formatStr = "%d%s\n";
                    printfArgs.push_back(context.builder.CreateGlobalStringPtr(formatStr));
                    printfArgs.push_back(lhsVal);
                    printfArgs.push_back(rhsVal);
                } else if (lhsVal->getType()->isPointerTy() && rhsVal->getType()->isPointerTy()) {
                    // String + String
                    formatStr = "%s%s\n";
                    printfArgs.push_back(context.builder.CreateGlobalStringPtr(formatStr));
                    printfArgs.push_back(lhsVal);
                    printfArgs.push_back(rhsVal);
                } else {
                    // Fall back to normal binary expression
                    llvm::Value* val = binExpr->codegen(context);
                    if (!val) return nullptr;
                    
                    if (val->getType()->isIntegerTy(32)) {
                        formatStr = "%d\n";
                        printfArgs.push_back(context.builder.CreateGlobalStringPtr(formatStr));
                        printfArgs.push_back(val);
                    } else if (val->getType()->isPointerTy()) {
                        formatStr = "%s\n";
                        printfArgs.push_back(context.builder.CreateGlobalStringPtr(formatStr));
                        printfArgs.push_back(val);
                    } else {
                        std::cerr << "Unsupported type for print statement" << std::endl;
                        return nullptr;
                    }
                }
                
                return context.builder.CreateCall(printfFunc, printfArgs, "printfcall");
            }
        }
    }
    
    // Normal handling for non-concatenation expressions
    llvm::Value* val = value->codegen(context);
    if (!val) {
        std::cerr << "Error: Failed to generate value for print statement" << std::endl;
        return nullptr;
    }
    
    // Choose format string based on the type of the value
    std::string formatStr;
    std::vector<llvm::Value*> args;
    
    llvm::Type* valType = val->getType();
    if (valType->isIntegerTy(32)) {
        formatStr = "%d\n";
        args.push_back(val);
    } else if (valType->isIntegerTy(1)) {
        // Boolean - print as "true" or "false"
        formatStr = "%s\n";
        llvm::Value* trueStr = context.builder.CreateGlobalStringPtr("true");
        llvm::Value* falseStr = context.builder.CreateGlobalStringPtr("false");
        args.push_back(context.builder.CreateSelect(val, trueStr, falseStr));
    } else if (valType->isPointerTy()) {
        // For pointers (including string pointers)
        formatStr = "%s\n";
        args.push_back(val);
    } else {
        std::cerr << "Unsupported type for print statement" << std::endl;
        return nullptr;
    }
    
    // Create the format string
    llvm::Value* formatStrVal = context.builder.CreateGlobalStringPtr(formatStr);
    
    // Create the call to printf
    std::vector<llvm::Value*> printfArgs;
    printfArgs.push_back(formatStrVal);
    printfArgs.insert(printfArgs.end(), args.begin(), args.end());
    
    return context.builder.CreateCall(printfFunc, printfArgs, "printfcall");
}

// Implementation for the new ReturnStmt class
llvm::Value* ReturnStmt::codegen(CodeGenContext& context) {
    debugOutput("Generating code for return statement");
    
    if (!value) {
        std::cerr << "Error: Null value in return statement" << std::endl;
        return nullptr;
    }
    
    // Generate code for the return value
    llvm::Value* retVal = value->codegen(context);
    if (!retVal) {
        std::cerr << "Error: Failed to generate return value" << std::endl;
        return nullptr;
    }
    
    // Get the current function
    llvm::Function* function = context.builder.GetInsertBlock()->getParent();
    if (!function) {
        std::cerr << "Error: Null function in return statement" << std::endl;
        return nullptr;
    }
    
    // Check if the function return type matches the return value type
    llvm::Type* funcRetType = function->getReturnType();
    llvm::Type* valType = retVal->getType();
    
    // If the function is declared as void but we're returning a value, change the function type
    if (funcRetType->isVoidTy() && !valType->isVoidTy()) {
        // This is a bit complex and would require modifying the function signature
        // For simplicity, we'll just convert the return value to match the function type
        debugOutput("Warning: Returning a value from a void function");
        context.builder.CreateRetVoid();
        return retVal;
    }
    
    // If the types don't match, try to convert
    if (funcRetType != valType) {
        if (funcRetType->isIntegerTy() && valType->isIntegerTy()) {
            // Convert between integer types
            retVal = context.builder.CreateIntCast(retVal, funcRetType, true, "cast_ret");
        } else {
            std::cerr << "Error: Return value type doesn't match function return type" << std::endl;
            return nullptr;
        }
    }
    
    // Create the return instruction
    context.builder.CreateRet(retVal);
    
    return retVal;
}

// Implementation for FunctionDef with parameters
llvm::Value* FunctionDef::codegen(CodeGenContext& context) {
    debugOutput("Generating code for function: " + name);

    if (!body) {
        std::cerr << "Error: Null body in function definition" << std::endl;
        return nullptr;
    }

    // Check if the function contains a return statement to determine return type
    bool hasReturnStmt = false;
    llvm::Type* returnType = llvm::Type::getVoidTy(context.context);
    
    // Scan the function body for return statements
    for (auto stmt : *body) {
        if (auto retStmt = dynamic_cast<ReturnStmt*>(stmt)) {
            hasReturnStmt = true;
            // For simplicity, we'll assume all functions return int32 if they have a return statement
            returnType = llvm::Type::getInt32Ty(context.context);
            break;
        }
    }

    // Create function type with parameters
    std::vector<llvm::Type*> argTypes(params.size(), llvm::Type::getInt32Ty(context.context));
    llvm::FunctionType* funcType = llvm::FunctionType::get(returnType, argTypes, false);

    // Create the function
    llvm::Function* function = llvm::Function::Create(
        funcType, llvm::Function::ExternalLinkage, name, context.module.get());

    if (!function) {
        std::cerr << "Error: Failed to create function: " << name << std::endl;
        return nullptr;
    }

    // Set names for all arguments
    unsigned idx = 0;
    for (auto &arg : function->args()) {
        if (idx < params.size()) {
            arg.setName(params[idx++]);
        }
    }

    // Add the function to our function table before generating body
    // This allows for recursive calls
    context.functions[name] = function;

    // Create a basic block for the function entry
    llvm::BasicBlock* entryBB = llvm::BasicBlock::Create(context.context, "entry", function);
    
    // Save the current insertion point
    llvm::BasicBlock* savedBlock = context.builder.GetInsertBlock();
    llvm::Function* savedFunction = savedBlock ? savedBlock->getParent() : nullptr;
    
    // Set insertion point to the new function's entry block
    context.builder.SetInsertPoint(entryBB);
    debugOutput("Setting insertion point to entry block of function: " + name);
    
    // Save the current named values
    std::map<std::string, llvm::AllocaInst*> oldNamedValues = context.namedValues;
    context.namedValues.clear(); // Clear the named values for the new function scope

    // Create allocas for all parameters
    idx = 0;
    for (auto &arg : function->args()) {
        // Create an alloca for this argument
        llvm::AllocaInst* alloca = context.createEntryBlockAlloca(
            function, arg.getName().str(), arg.getType());
        
        // Store the initial value into the alloca
        context.builder.CreateStore(&arg, alloca);
        
        // Add arguments to variable symbol table
        context.namedValues[arg.getName().str()] = alloca;
    }

    // Generate code for the function body
    for (auto stmt : *body) {
        if (stmt) {
            stmt->codegen(context);
            
            // If we just generated a return statement, don't continue generating code
            if (dynamic_cast<ReturnStmt*>(stmt)) {
                break;
            }
        }
    }

    // Ensure the function has a terminator
    if (!context.builder.GetInsertBlock()->getTerminator()) {
        if (returnType->isVoidTy()) {
            context.builder.CreateRetVoid();
        } else {
            // Default return value for non-void functions
            context.builder.CreateRet(llvm::ConstantInt::get(returnType, 0));
        }
    }

    // Verify the function
    std::string verifyErrors;
    llvm::raw_string_ostream errStream(verifyErrors);
    bool hasErrors = llvm::verifyFunction(*function, &errStream);
    
    if (hasErrors) {
        std::cerr << "Function verification failed for " << name << ":" << std::endl;
        std::cerr << verifyErrors << std::endl;
        function->eraseFromParent();
        return nullptr;
    }

    // Restore the previous named values
    context.namedValues = oldNamedValues;
    
    // Restore the previous insertion point
    if (savedBlock) {
        context.builder.SetInsertPoint(savedBlock);
    }

    return function;
}

// Updated FunctionCall implementation to handle arguments
llvm::Value* FunctionCall::codegen(CodeGenContext& context) {
    // Look up the function in the function table
    auto it = context.functions.find(name);
    if (it == context.functions.end() || !it->second) {
        std::cerr << "Unknown function: " << name << std::endl;
        return nullptr;
    }
    
    llvm::Function* function = it->second;
    if (!function) {
        std::cerr << "Error: Null function pointer for: " << name << std::endl;
        return nullptr;
    }
    
    // Check argument count
    if (function->arg_size() != (args ? args->size() : 0)) {
        std::cerr << "Error: Function " << name << " called with wrong number of arguments. ";
        std::cerr << "Expected " << function->arg_size() << ", got " << (args ? args->size() : 0) << std::endl;
        return nullptr;
    }
    
    // Process arguments if present
    std::vector<llvm::Value*> argValues;
    if (args) {
        for (auto arg : *args) {
            if (arg) {
                llvm::Value* argValue = arg->codegen(context);
                if (argValue) {
                    argValues.push_back(argValue);
                } else {
                    std::cerr << "Error: Failed to generate code for function argument" << std::endl;
                    return nullptr;
                }
            }
        }
    }
    
    // Create the function call
    llvm::Value* callResult = context.builder.CreateCall(function, argValues);
    
    // If the function returns void, return a dummy value
    if (function->getReturnType()->isVoidTy()) {
        return llvm::Constant::getNullValue(llvm::Type::getInt32Ty(context.context));
    }
    
    return callResult;
}

static void finalizeFunction(llvm::Function *F) {
    if (F->isDeclaration()) return;  // skip external declarations

    // A single IRBuilder we can reuse for all blocks
    llvm::IRBuilder<> Builder(F->getContext());

    for (llvm::BasicBlock &BB : *F) {
        // Check if the block has a terminator
        if (!BB.getTerminator()) {
            Builder.SetInsertPoint(&BB);
            
            // Add appropriate terminator based on function return type
            if (F->getReturnType()->isVoidTy()) {
                Builder.CreateRetVoid();
            } else {
                Builder.CreateRet(llvm::ConstantInt::get(
                    llvm::Type::getInt32Ty(F->getContext()), 0));
            }
        }
    }
}

llvm::Value* Program::codegen(CodeGenContext& context) {
    debugOutput("Generating code for main program");
    
    // Create the main function
    std::vector<llvm::Type*> argTypes;
    llvm::FunctionType* mainFuncType = llvm::FunctionType::get(
        llvm::Type::getInt32Ty(context.context), argTypes, false);
    
    llvm::Function* mainFunc = llvm::Function::Create(
        mainFuncType, llvm::Function::ExternalLinkage, "main", context.module.get());
    
    if (!mainFunc) {
        std::cerr << "Error: Failed to create main function" << std::endl;
        return nullptr;
    }
    
    // Create a basic block for the entry point
    llvm::BasicBlock* entryBB = llvm::BasicBlock::Create(context.context, "entry", mainFunc);
    context.builder.SetInsertPoint(entryBB);
    debugOutput("Setting insertion point to entry block of main function");
    
    // First, process all function definitions to make them available for calls
    if (statements) {
        for (auto stmt : *statements) {
            if (auto funcDef = dynamic_cast<FunctionDef*>(stmt)) {
                funcDef->codegen(context);
            }
        }
    }
    
    // Make sure we're still in the main function's entry block
    context.builder.SetInsertPoint(entryBB);
    
    // Then process all non-function-definition statements
    if (statements) {
        for (auto stmt : *statements) {
            if (!dynamic_cast<FunctionDef*>(stmt)) {
                stmt->codegen(context);
                
                // If we just generated a return statement, don't continue generating code
                if (dynamic_cast<ReturnStmt*>(stmt) && !context.builder.GetInsertBlock()->getTerminator()) {
                    break;
                }
            }
        }
    }
    
    // Add a return 0 to the main function if it doesn't already have a terminator
    if (!context.builder.GetInsertBlock()->getTerminator()) {
        context.builder.CreateRet(llvm::ConstantInt::get(llvm::Type::getInt32Ty(context.context), 0));
    }

    return mainFunc;
}

// Generate code for the entire program
void CodeGenContext::generateCode(Program* program) {
    if (!program) {
        std::cerr << "Error: Null program in generateCode" << std::endl;
        return;
    }
    std::cout << "Starting code generation..." << std::endl;
    llvm::Value* result = program->codegen(*this);
    std::cout << "Code generation " << (result ? "successful" : "failed") << std::endl;
    
    // Verify the entire module
    std::string verifyErrors;
    llvm::raw_string_ostream errStream(verifyErrors);
    bool hasErrors = llvm::verifyModule(*module, &errStream);
    
    if (hasErrors) {
        std::cerr << "Module verification failed:" << std::endl;
        std::cerr << verifyErrors << std::endl;
    } else {
        debugOutput("Module verified successfully");
    }
}

// Print the generated IR
void CodeGenContext::printIR() {
    if (!module) {
        std::cerr << "Error: No module to print" << std::endl;
        return;
    }
    
    std::cout << "Printing generated LLVM IR:" << std::endl;
    module->print(llvm::outs(), nullptr);
}

// Write the generated IR to a file
void CodeGenContext::writeIR(const std::string& filename) {
    std::error_code EC;
    llvm::raw_fd_ostream dest(filename, EC, llvm::sys::fs::OF_None);
    
    if (EC) {
        std::cerr << "Could not open file: " << EC.message() << std::endl;
        return;
    }
    
    module->print(dest, nullptr);
}

// Run optimization passes on the generated code
void CodeGenContext::optimize() {
    llvm::LoopAnalysisManager LAM;
    llvm::FunctionAnalysisManager FAM;
    llvm::CGSCCAnalysisManager CGAM;
    llvm::ModuleAnalysisManager MAM;
    
    // Create pass builder and register analysis passes
    llvm::PassBuilder PB;
    PB.registerModuleAnalyses(MAM);
    PB.registerCGSCCAnalyses(CGAM);
    PB.registerFunctionAnalyses(FAM);
    PB.registerLoopAnalyses(LAM);
    PB.crossRegisterProxies(LAM, FAM, CGAM, MAM);
    
    // Create optimization pipeline
    llvm::ModulePassManager MPM = PB.buildPerModuleDefaultPipeline(llvm::OptimizationLevel::O2);
    
    // Run the optimizations
    MPM.run(*module, MAM);
}

// Run the JIT compiler
void CodeGenContext::runJIT() {
    std::cout << "Initializing JIT execution..." << std::endl;
    
    // Initialize LLVM targets
    llvm::InitializeNativeTarget();
    llvm::InitializeNativeTargetAsmPrinter();
    llvm::InitializeNativeTargetAsmParser();
    
    std::string error;
    
    // Verify the module before execution
    std::string verifyErrors;
    llvm::raw_string_ostream errStream(verifyErrors);
    
    if (llvm::verifyModule(*module, &errStream)) {
        std::cerr << "Error: Module verification failed. Cannot execute." << std::endl;
        std::cerr << "Verification errors:" << std::endl;
        std::cerr << verifyErrors << std::endl;
        return;
    }
    
    std::cout << "Module verified successfully. Creating execution engine..." << std::endl;
    
    // Clone the module to avoid moving it out of the context
    std::unique_ptr<llvm::Module> moduleCopy;
    try {
        moduleCopy = llvm::CloneModule(*module);
    } catch (const std::exception& e) {
        std::cerr << "Error cloning module: " << e.what() << std::endl;
        return;
    }
    
    if (!moduleCopy) {
        std::cerr << "Error: Failed to clone module for JIT execution" << std::endl;
        return;
    }
    
    // Create execution engine with proper error handling
    llvm::EngineBuilder engineBuilder(std::move(moduleCopy));
    engineBuilder.setErrorStr(&error);
    engineBuilder.setEngineKind(llvm::EngineKind::JIT);
    // Use CodeGenOptLevel instead of CodeGenOpt::Default
    engineBuilder.setOptLevel(llvm::CodeGenOptLevel::Default);
    
    auto executionEngine = engineBuilder.create();
    
    if (!executionEngine) {
        std::cerr << "Could not create execution engine: " << error << std::endl;
        return;
    }
    
    std::cout << "Execution engine created. Looking for main function..." << std::endl;
    
    // Find the main function
    llvm::Function* mainFunc = executionEngine->FindFunctionNamed("main");
    if (!mainFunc) {
        std::cerr << "Could not find main function" << std::endl;
        return;
    }
    
    std::cout << "Found main function. Running program..." << std::endl;
    
    try {
        // Run the main function
        std::vector<llvm::GenericValue> noargs;
        auto result = executionEngine->runFunction(mainFunc, noargs);
        
        std::cout << "Program exited with code: " << result.IntVal.getZExtValue() << std::endl;
    } catch (const std::exception& e) {
        std::cerr << "Error during execution: " << e.what() << std::endl;
    }
    
    // Clean up
    delete executionEngine;
    std::cout << "JIT execution completed." << std::endl;
}