#include "ModuleBuilder.h"

#include "Error.h"
#include "Logger.h"
#include "Parser/ValueType.h"

#include "Parser/Expression/ExpressionGrouping.h"
#include "Parser/Expression/ExpressionLiteral.h"
#include "Parser/Expression/ExpressionArrayLiteral.h"
#include "Parser/Expression/ExpressionVariable.h"
#include "Parser/Expression/ExpressionCall.h"
#include "Parser/Expression/ExpressionIfElse.h"
#include "Parser/Expression/ExpressionBinary.h"
#include "Parser/Expression/ExpressionUnary.h"
#include "Parser/Expression/ExpressionBlock.h"

#include "Parser/Statement/StatementFunction.h"
#include "Parser/Statement/StatementRawFunction.h"
#include "Parser/Statement/StatementBlob.h"
#include "Parser/Statement/StatementVariable.h"
#include "Parser/Statement/StatementAssignment.h"
#include "Parser/Statement/StatementReturn.h"
#include "Parser/Statement/StatementExpression.h"
#include "Parser/Statement/StatementRepeat.h"
#include "Parser/Statement/StatementMetaExternFunction.h"
#include "Parser/Statement/StatementBlock.h"

ModuleBuilder::ModuleBuilder(string moduleName, string sourceFileName, vector<shared_ptr<Statement>> statements):
moduleName(moduleName), sourceFileName(sourceFileName), statements(statements) {
    context = make_shared<llvm::LLVMContext>();
    module = make_shared<llvm::Module>(moduleName, *context);
    module->setSourceFileName(sourceFileName);
    builder = make_shared<llvm::IRBuilder<>>(*context);

    typeVoid = llvm::Type::getVoidTy(*context);
    typeBool = llvm::Type::getInt1Ty(*context);
    typeU8 = llvm::Type::getInt8Ty(*context);
    typeU32 = llvm::Type::getInt32Ty(*context);
    typeS8 = llvm::Type::getInt8Ty(*context);
    typeS32 = llvm::Type::getInt32Ty(*context);
    typeR32 = llvm::Type::getFloatTy(*context);
}

shared_ptr<llvm::Module> ModuleBuilder::getModule() {
    scopes.push(Scope());
    for (shared_ptr<Statement> &statement : statements)
        buildStatement(statement);

    // verify module
    string errorMessage;
    llvm::raw_string_ostream llvmErrorMessage(errorMessage);
    if (llvm::verifyModule(*module, &llvmErrorMessage))
        markError(0, 0, errorMessage);

    if (!errors.empty()) {
        for (shared_ptr<Error> &error : errors)
            Logger::print(error);
        exit(1);
    }

    return module;
}

void ModuleBuilder::buildStatement(shared_ptr<Statement> statement) {
    switch (statement->getKind()) {
        case StatementKind::FUNCTION:
            buildFunction(dynamic_pointer_cast<StatementFunction>(statement));
            break;
        case StatementKind::RAW_FUNCTION:
            buildRawFunction(dynamic_pointer_cast<StatementRawFunction>(statement));
            break;
        case StatementKind::BLOB:
            buildBlob(dynamic_pointer_cast<StatementBlob>(statement));
            break;
        case StatementKind::VARIABLE:
            buildVarDeclaration(dynamic_pointer_cast<StatementVariable>(statement));
            break;
        case StatementKind::ASSIGNMENT:
            buildAssignment(dynamic_pointer_cast<StatementAssignment>(statement));
            break;
        case StatementKind::BLOCK:
            buildBlock(dynamic_pointer_cast<StatementBlock>(statement));
            break;
        case StatementKind::RETURN:
            buildReturn(dynamic_pointer_cast<StatementReturn>(statement));
            break;
        case StatementKind::REPEAT:
            buildLoop(dynamic_pointer_cast<StatementRepeat>(statement));
            break;
        case StatementKind::META_EXTERN_FUNCTION:
            buildMetaExternFunction(dynamic_pointer_cast<StatementMetaExternFunction>(statement));
            break;
        case StatementKind::EXPRESSION:
            buildExpression(dynamic_pointer_cast<StatementExpression>(statement));
            return;
        default:
            markError(0, 0, "Unexpected statement");
    }
}

void ModuleBuilder::buildFunction(shared_ptr<StatementFunction> statement) {
    // function types
    llvm::Type *returnType = typeForValueType(statement->getReturnValueType());
    if (returnType == nullptr)
        return;
    vector<llvm::Type *> argTypes;
    for (pair<string, shared_ptr<ValueType>> &arg : statement->getArguments()) {
        llvm::Type *argType = typeForValueType(arg.second);
        if (argType == nullptr)
            return;
        argTypes.push_back(argType);
    }

    // build function declaration
    llvm::FunctionType *funType = llvm::FunctionType::get(returnType, argTypes, false);
    llvm::Function *fun = llvm::Function::Create(funType, llvm::GlobalValue::ExternalLinkage, statement->getName(), module.get());
    if (!setFun(statement->getName(), fun))
        return;

    // define function body
    llvm::BasicBlock *block = llvm::BasicBlock::Create(*context, statement->getName(), fun);
    builder->SetInsertPoint(block);

    scopes.push(Scope());

    // build arguments
    int i=0;
    for (auto &arg : fun->args()) {
        string name = statement->getArguments()[i].first;
        llvm::Type *type = argTypes[i];
        arg.setName(name);

        llvm::AllocaInst *alloca = builder->CreateAlloca(type, nullptr, name);
        if (!setAlloca(name, alloca))
            return;
        builder->CreateStore(&arg, alloca);

        i++;
    }

    // build function body
    buildStatement(statement->getStatementBlock());

    scopes.pop();

    // verify
    string errorMessage;
    llvm::raw_string_ostream llvmErrorMessage(errorMessage);
    if (llvm::verifyFunction(*fun, &llvmErrorMessage))
        markError(0, 0, errorMessage);
}

void ModuleBuilder::buildRawFunction(shared_ptr<StatementRawFunction> statement) {
    // function types
    llvm::Type *returnType = typeForValueType(statement->getReturnValueType());
    vector<llvm::Type *> argTypes;
    for (pair<string, shared_ptr<ValueType>> &arg : statement->getArguments())
        argTypes.push_back(typeForValueType(arg.second));

    // build function declaration & body
    llvm::FunctionType *funType = llvm::FunctionType::get(returnType, argTypes, false);
    llvm::InlineAsm *rawFun = llvm::InlineAsm::get(funType, statement->getRawSource(), statement->getConstraints(), true, false, llvm::InlineAsm::AsmDialect::AD_Intel);
    if (!setRawFun(statement->getName(), rawFun))
        return;
}

void ModuleBuilder::buildBlob(shared_ptr<StatementBlob> statement) {
    llvm::StructType *structType = llvm::StructType::create(*context, statement->getIdentifier());

    // Generate types for body
    vector<string> memberNames;
    vector<llvm::Type *> types;
    for (pair<string, shared_ptr<ValueType>> &variable: statement->getVariables()) {
        memberNames.push_back(variable.first);
        llvm::Type *type = typeForValueType(variable.second);
        if (type == nullptr)
            return;
        types.push_back(type);
    }
    structType->setBody(types, false);
    if (!registerStruct(statement->getIdentifier(), structType, memberNames))
        return;
}

void ModuleBuilder::buildVarDeclaration(shared_ptr<StatementVariable> statement) {
    if (statement->getValueType()->getKind() == ValueTypeKind::DATA) {
        vector<llvm::Value *> values = valuesForExpression(statement->getExpression());

        llvm::ArrayType *type = (llvm::ArrayType *)typeForValueType(statement->getValueType(), values.size());
        llvm::AllocaInst *alloca = builder->CreateAlloca(type, nullptr, statement->getName());
        if (!setAlloca(statement->getName(), alloca))
            return;
        for (int i=0; i < type->getNumElements(); i++) {
            llvm::Value *index[] = {
                builder->getInt32(0),
                builder->getInt32(i)
            };
            llvm::Value *elementPtr = builder->CreateGEP(type, alloca, index, format("{}_{}", statement->getName(), i));

            builder->CreateStore(values[i], elementPtr);
        }
    } else if (statement->getValueType()->getKind() == ValueTypeKind::TYPE) {
        llvm::StructType *type = (llvm::StructType *)typeForValueType(statement->getValueType(), 0);
        llvm::AllocaInst *alloca = builder->CreateAlloca(type, nullptr, statement->getName());
        if (!setAlloca(statement->getName(), alloca))
            return;
    } else {
        llvm::AllocaInst *alloca = builder->CreateAlloca(typeForValueType(statement->getValueType(), 0), nullptr, statement->getName());

        if (!setAlloca(statement->getName(), alloca))
            return;
        
        // set initial value
        if (statement->getExpression() != nullptr) {
            llvm::Value *value = nullptr;
            value = valueForExpression(statement->getExpression());
            if (value == nullptr)
                return;
            builder->CreateStore(value, alloca);
        }
    }
}

void ModuleBuilder::buildAssignment(shared_ptr<StatementAssignment> statement) {
    llvm::AllocaInst *alloca = getAlloca(statement->getIdentifier());
    if (alloca == nullptr)
        return;

    llvm::Value *value = valueForExpression(statement->getValueExpression());

    switch (statement->getAssignmentKind()) {
        case StatementAssignmentKind::SIMPLE: {
            builder->CreateStore(value, alloca);
            break;
        }
        case StatementAssignmentKind::DATA: {
            llvm::Value *indexValue = valueForExpression(statement->getIndexExpression());
            llvm::Value *index[] = {
                builder->getInt32(0),
                indexValue
            };
            llvm::ArrayType *type = (llvm::ArrayType *)alloca->getAllocatedType();
            llvm::Value *elementPtr = builder->CreateGEP(type, alloca, index, format("{}[]", statement->getIdentifier()));

            builder->CreateStore(value, elementPtr);
            break;
        }
        case StatementAssignmentKind::BLOB: {        
            llvm::StructType *structType = (llvm::StructType *)alloca->getAllocatedType();
            string structName = string(structType->getName());
            optional<int> memberIndex = getMemberIndex(structName, statement->getMemberName());
            if (!memberIndex)
                return;
            llvm::Value *index[] = {
                builder->getInt32(0),
                builder->getInt32(*memberIndex)
            };
            llvm::Value *elementPtr = builder->CreateGEP(structType, alloca, index, format("{}.{}", statement->getIdentifier(), statement->getMemberName()));
            builder->CreateStore(value, elementPtr);
            break;
        }
    }
}

void ModuleBuilder::buildBlock(shared_ptr<StatementBlock> statement) {
    for (shared_ptr<Statement> &innerStatement : statement->getStatements())
        buildStatement(innerStatement);
}

void ModuleBuilder::buildReturn(shared_ptr<StatementReturn> statement) {
    if (statement->getExpression() != nullptr) {
        llvm::Value *value = valueForExpression(statement->getExpression());
        builder->CreateRet(value);
    } else {
        builder->CreateRetVoid();
    }
}

void ModuleBuilder::buildLoop(shared_ptr<StatementRepeat> statement) {
    shared_ptr<Statement> initStatement = statement->getInitStatement();
    shared_ptr<StatementBlock> bodyStatement= statement->getBodyBlockStatement();
    shared_ptr<Expression> preExpression = statement->getPreConditionExpression();
    shared_ptr<Expression> postExpression = statement->getPostConditionExpression();

    llvm::Function *fun = builder->GetInsertBlock()->getParent();
    llvm::BasicBlock *preBlock = llvm::BasicBlock::Create(*context, "loopPre", fun);
    llvm::BasicBlock *bodyBlock = llvm::BasicBlock::Create(*context, "loopBody");
    llvm::BasicBlock *afterBlock = llvm::BasicBlock::Create(*context, "loopPost");

    scopes.push(Scope());

    // loop init
    if (initStatement != nullptr)
        buildStatement(statement->getInitStatement());
    builder->CreateBr(preBlock);

    // pre condition
    builder->SetInsertPoint(preBlock);
    if (preExpression != nullptr) {
        llvm::Value *preConditionValue = valueForExpression(preExpression);
        builder->CreateCondBr(preConditionValue, bodyBlock, afterBlock);
    } else {
        builder->CreateBr(bodyBlock);
    }

    // body
    fun->insert(fun->end(), bodyBlock);
    builder->SetInsertPoint(bodyBlock);
    buildBlock(bodyStatement);

    // post condition
    if (postExpression != nullptr) {
        llvm::Value *postConditionValue = valueForExpression(postExpression);
        builder->CreateCondBr(postConditionValue, preBlock, afterBlock);
    } else {
        builder->CreateBr(preBlock);
    }

    // loop post
    fun->insert(fun->end(), afterBlock);
    builder->SetInsertPoint(afterBlock);

    scopes.pop();
}

void ModuleBuilder::buildMetaExternFunction(shared_ptr<StatementMetaExternFunction> statement) {
    // get argument types
    vector<llvm::Type *> types;
    for (pair<string, shared_ptr<ValueType>> &arg : statement->getArguments()) {
        types.push_back(typeForValueType(arg.second));
    }

    // build function declaration
    llvm::Type *returnType = typeForValueType(statement->getReturnValueType());
    if (returnType == nullptr)
        return;
    llvm::FunctionType *funType = llvm::FunctionType::get(returnType, types, false);
    llvm::Function *fun = llvm::Function::Create(funType, llvm::GlobalValue::ExternalLinkage, statement->getName(), module.get());
    if (!setFun(statement->getName(), fun))
        return;
    
    // build arguments
    int i=0;
    for (auto &arg : fun->args()) {
        string name = statement->getArguments()[i].first;
        arg.setName(name);
        i++;
    }
}

void ModuleBuilder::buildExpression(shared_ptr<StatementExpression> statement) {
    valueForExpression(statement->getExpression());
}

llvm::Value *ModuleBuilder::valueForExpression(shared_ptr<Expression> expression) {
    switch (expression->getKind()) {
        case ExpressionKind::LITERAL:
            return valueForLiteral(dynamic_pointer_cast<ExpressionLiteral>(expression));
        case ExpressionKind::GROUPING:
            return valueForExpression(dynamic_pointer_cast<ExpressionGrouping>(expression)->getExpression());
        case ExpressionKind::BINARY:
            return valueForBinary(dynamic_pointer_cast<ExpressionBinary>(expression));
        case ExpressionKind::UNARY:
            return valueForUnary(dynamic_pointer_cast<ExpressionUnary>(expression));
        case ExpressionKind::IF_ELSE:
            return valueForIfElse(dynamic_pointer_cast<ExpressionIfElse>(expression));
        case ExpressionKind::VARIABLE:
            return valueForVariable(dynamic_pointer_cast<ExpressionVariable>(expression));
        case ExpressionKind::CALL:
            return valueForCall(dynamic_pointer_cast<ExpressionCall>(expression));
        default:
            markError(0, 0, "Unexpected expression");
            return nullptr;
    }
}

vector<llvm::Value*> ModuleBuilder::valuesForExpression(shared_ptr<Expression> expression) {
    switch (expression->getKind()) {
        case ExpressionKind::ARRAY_LITERAL:
            return valuesForArrayLiteral(dynamic_pointer_cast<ExpressionArrayLiteral>(expression));
        default:
            markError(0, 0, "Unexpected expression");
            return vector<llvm::Value*>();
    }
}

llvm::Value *ModuleBuilder::valueForLiteral(shared_ptr<ExpressionLiteral> expression) {
    if (expression->getValueType() == nullptr)
        return llvm::UndefValue::get(typeVoid);

    switch (expression->getValueType()->getKind()) {
        case ValueTypeKind::NONE:
            return llvm::UndefValue::get(typeVoid);
        case ValueTypeKind::BOOL:
            return llvm::ConstantInt::get(typeBool, expression->getBoolValue(), true);
        case ValueTypeKind::U8:
            return llvm::ConstantInt::get(typeU8, expression->getU8Value(), true);
        case ValueTypeKind::U32:
            return llvm::ConstantInt::get(typeU32, expression->getU32Value(), true);
        case ValueTypeKind::S8:
            return llvm::ConstantInt::get(typeS8, expression->getS8Value(), true);
        case ValueTypeKind::S32:
            return llvm::ConstantInt::get(typeS32, expression->getS32Value(), true);
        case ValueTypeKind::R32:
            return llvm::ConstantFP::get(typeR32, expression->getR32Value());
    }
}

vector<llvm::Value*> ModuleBuilder::valuesForArrayLiteral(shared_ptr<ExpressionArrayLiteral> expression) {
    vector<llvm::Value*> values;
    for (shared_ptr<Expression> &expression : expression->getExpressions()) {
        values.push_back(valueForExpression(expression));
    }
    return values;
}

llvm::Value *ModuleBuilder::valueForGrouping(shared_ptr<ExpressionGrouping> expression) {
    return valueForExpression(expression->getExpression());
}

llvm::Value *ModuleBuilder::valueForBinary(shared_ptr<ExpressionBinary> expression) {
    llvm::Value *leftValue = valueForExpression(expression->getLeft());
    llvm::Value *rightValue = valueForExpression(expression->getRight());

    if (leftValue == nullptr || rightValue == nullptr)
        return nullptr;

    llvm::Type *type = leftValue->getType();

    if (type == typeBool) {
        return valueForBinaryBool(expression->getOperation(), leftValue, rightValue);
    } else if (type == typeU8 || type == typeU32) {
        return valueForBinaryUnsignedInteger(expression->getOperation(), leftValue, rightValue);
    } else if (type == typeS8 || type == typeS32) {
        return valueForBinarySignedInteger(expression->getOperation(), leftValue, rightValue);
    } else if (type == typeR32) {
        return valueForBinaryReal(expression->getOperation(), leftValue, rightValue);
    } else { // FIXME (we have missing value types)
        return valueForBinarySignedInteger(expression->getOperation(), leftValue, rightValue);
    }

    markError(0, 0, "Unexpected operation");
    return nullptr;
}

llvm::Value *ModuleBuilder::valueForBinaryBool(ExpressionBinaryOperation operation, llvm::Value *leftValue, llvm::Value *rightValue) {
    switch (operation) {
    case ExpressionBinaryOperation::EQUAL:
        return builder->CreateICmpEQ(leftValue, rightValue);
    case ExpressionBinaryOperation::NOT_EQUAL:
        return builder->CreateICmpNE(leftValue, rightValue);
    default:
        markError(0, 0, "Unexpected operation for boolean operands");
        return nullptr;
    }
}

llvm::Value *ModuleBuilder::valueForBinaryUnsignedInteger(ExpressionBinaryOperation operation, llvm::Value *leftValue, llvm::Value *rightValue) {
    switch (operation) {
        case ExpressionBinaryOperation::EQUAL:
            return builder->CreateICmpEQ(leftValue, rightValue);
        case ExpressionBinaryOperation::NOT_EQUAL:
            return builder->CreateICmpNE(leftValue, rightValue);
        case ExpressionBinaryOperation::LESS:
            return builder->CreateICmpSLT(leftValue, rightValue);
        case ExpressionBinaryOperation::LESS_EQUAL:
            return builder->CreateICmpSLE(leftValue, rightValue);
        case ExpressionBinaryOperation::GREATER:
            return builder->CreateICmpSGT(leftValue, rightValue);
        case ExpressionBinaryOperation::GREATER_EQUAL:
            return builder->CreateICmpSGE(leftValue, rightValue);
        case ExpressionBinaryOperation::ADD:
            return builder->CreateNUWAdd(leftValue, rightValue);
        case ExpressionBinaryOperation::SUB:
            return builder->CreateNUWSub(leftValue, rightValue);
        case ExpressionBinaryOperation::MUL:
            return builder->CreateNUWMul(leftValue, rightValue);
        case ExpressionBinaryOperation::DIV:
            return builder->CreateUDiv(leftValue, rightValue);
        case ExpressionBinaryOperation::MOD:
            return builder->CreateURem(leftValue, rightValue);
    }
}

llvm::Value *ModuleBuilder::valueForBinarySignedInteger(ExpressionBinaryOperation operation, llvm::Value *leftValue, llvm::Value *rightValue) {
    switch (operation) {
        case ExpressionBinaryOperation::EQUAL:
            return builder->CreateICmpEQ(leftValue, rightValue);
        case ExpressionBinaryOperation::NOT_EQUAL:
            return builder->CreateICmpNE(leftValue, rightValue);
        case ExpressionBinaryOperation::LESS:
            return builder->CreateICmpSLT(leftValue, rightValue);
        case ExpressionBinaryOperation::LESS_EQUAL:
            return builder->CreateICmpSLE(leftValue, rightValue);
        case ExpressionBinaryOperation::GREATER:
            return builder->CreateICmpSGT(leftValue, rightValue);
        case ExpressionBinaryOperation::GREATER_EQUAL:
            return builder->CreateICmpSGE(leftValue, rightValue);
        case ExpressionBinaryOperation::ADD:
            return builder->CreateNSWAdd(leftValue, rightValue);
        case ExpressionBinaryOperation::SUB:
            return builder->CreateNSWSub(leftValue, rightValue);
        case ExpressionBinaryOperation::MUL:
            return builder->CreateNSWMul(leftValue, rightValue);
        case ExpressionBinaryOperation::DIV:
            return builder->CreateSDiv(leftValue, rightValue);
        case ExpressionBinaryOperation::MOD:
            return builder->CreateSRem(leftValue, rightValue);
    }
}

llvm::Value *ModuleBuilder::valueForBinaryReal(ExpressionBinaryOperation operation, llvm::Value *leftValue, llvm::Value *rightValue) {
    switch (operation) {
    case ExpressionBinaryOperation::EQUAL:
        return builder->CreateFCmpOEQ(leftValue, rightValue);
    case ExpressionBinaryOperation::NOT_EQUAL:
        return builder->CreateFCmpONE(leftValue, rightValue);
    case ExpressionBinaryOperation::LESS:
        return builder->CreateFCmpOLT(leftValue, rightValue);
    case ExpressionBinaryOperation::LESS_EQUAL:
        return builder->CreateFCmpOLE(leftValue, rightValue);
    case ExpressionBinaryOperation::GREATER:
        return builder->CreateFCmpOGT(leftValue, rightValue);
    case ExpressionBinaryOperation::GREATER_EQUAL:
        return builder->CreateFCmpOGE(leftValue, rightValue);
    case ExpressionBinaryOperation::ADD:
        return builder->CreateNSWAdd(leftValue, rightValue);
    case ExpressionBinaryOperation::SUB:
        return builder->CreateNSWSub(leftValue, rightValue);
    case ExpressionBinaryOperation::MUL:
        return builder->CreateNSWMul(leftValue, rightValue);
    case ExpressionBinaryOperation::DIV:
        return builder->CreateSDiv(leftValue, rightValue);
    case ExpressionBinaryOperation::MOD:
        return builder->CreateSRem(leftValue, rightValue);
    }
}

llvm::Value *ModuleBuilder::valueForUnary(shared_ptr<ExpressionUnary> expression) {
    llvm::Value *value = valueForExpression(expression->getExpression());
    llvm::Type *type = value->getType();

    // do nothing for plus
    if (expression->getOperation() == ExpressionUnaryOperation::PLUS)
        return value;

    if (type == typeU8 || type == typeU32) {
        return builder->CreateNeg(value);
    } else if (type == typeS8 || type == typeS32) {
        return builder->CreateNSWNeg(value);
    } else if (type == typeR32) {
        return builder->CreateFNeg(value);
    }

    markError(0, 0, "Unexpected operation");
    return nullptr;
}

llvm::Value *ModuleBuilder::valueForIfElse(shared_ptr<ExpressionIfElse> expression) {
    shared_ptr<Expression> conditionExpression = expression->getCondition();

    llvm::Function *fun = builder->GetInsertBlock()->getParent();
    llvm::Value *conditionValue = valueForExpression(conditionExpression);

    llvm::BasicBlock *thenBlock = llvm::BasicBlock::Create(*context, "thenBlock", fun);
    llvm::BasicBlock *elseBlock = llvm::BasicBlock::Create(*context, "elseBlock");
    llvm::BasicBlock *mergeBlock = llvm::BasicBlock::Create(*context, "mergeBlock");

    int valuesCount = 1;
    builder->CreateCondBr(conditionValue, thenBlock, elseBlock);

    // Then
    scopes.push(Scope());
    builder->SetInsertPoint(thenBlock);
    buildStatement(expression->getThenBlock()->getStatementBlock());
    llvm::Value *thenValue = valueForExpression(expression->getThenBlock()->getResultStatementExpression()->getExpression());
    builder->CreateBr(mergeBlock);
    thenBlock = builder->GetInsertBlock();
    scopes.pop();

    // Else
    scopes.push(Scope());
    fun->insert(fun->end(), elseBlock);
    builder->SetInsertPoint(elseBlock);
    llvm::Value *elseValue = nullptr;
    if (expression->getElseBlock() != nullptr) {
        valuesCount++;
        buildStatement(expression->getElseBlock()->getStatementBlock());
        elseValue = valueForExpression(expression->getElseBlock()->getResultStatementExpression()->getExpression());
    }
    builder->CreateBr(mergeBlock);
    elseBlock = builder->GetInsertBlock();
    scopes.pop();

    // Merge
    fun->insert(fun->end(), mergeBlock);
    builder->SetInsertPoint(mergeBlock);

    if (elseValue != nullptr && thenValue->getType() != elseValue->getType()) {
        return llvm::UndefValue::get(typeVoid);
    } else {
        llvm::PHINode *phi = builder->CreatePHI(thenValue->getType(), valuesCount, "ifElseResult");
        phi->addIncoming(thenValue, thenBlock);
        if (elseValue != nullptr)
            phi->addIncoming(elseValue, elseBlock);

        return phi;
    }
}

llvm::Value *ModuleBuilder::valueForVariable(shared_ptr<ExpressionVariable> expression) {
    llvm::AllocaInst *alloca = getAlloca(expression->getIdentifier());
    if (alloca == nullptr)
        return nullptr;

    switch (expression->getVariableKind()) {
        case ExpressionVariableKind::SIMPLE: {
            return builder->CreateLoad(alloca->getAllocatedType(), alloca, expression->getIdentifier());
        }
        case ExpressionVariableKind::DATA: {
            llvm::Value *indexValue = valueForExpression(expression->getIndexExpression());
            llvm::Value *index[] = {
                builder->getInt32(0),
                indexValue
            };
            llvm::ArrayType *type = (llvm::ArrayType *)alloca->getAllocatedType();
            llvm::Value *elementPtr = builder->CreateGEP(type, alloca, index, format("{}[]", expression->getIdentifier()));

            return builder->CreateLoad(type->getArrayElementType(), elementPtr);
        }
        case ExpressionVariableKind::BLOB: {
            llvm::StructType *structType = (llvm::StructType *)alloca->getAllocatedType();
            string structName = string(structType->getName());
            optional<int> memberIndex = getMemberIndex(structName, expression->getMemberName());
            if (!memberIndex)
                return nullptr;
            llvm::Value *index[] = {
                builder->getInt32(0),
                builder->getInt32(*memberIndex)
            };
            llvm::Value *elementPtr = builder->CreateGEP(structType, alloca, index, format("{}.{}", expression->getIdentifier(), expression->getMemberName()));
            return builder->CreateLoad(structType->getElementType(*memberIndex), elementPtr);
        }
    }
}

llvm::Value *ModuleBuilder::valueForCall(shared_ptr<ExpressionCall> expression) {
    llvm::Function *fun = getFun(expression->getName());
    if (fun != nullptr) {
        llvm::FunctionType *funType = fun->getFunctionType();
        vector<llvm::Value*> argValues;
        for (shared_ptr<Expression> &argumentExpression : expression->getArgumentExpressions()) {
            llvm::Value *argValue = valueForExpression(argumentExpression);
            argValues.push_back(argValue);
        }
        return builder->CreateCall(funType, fun, llvm::ArrayRef(argValues));
    }

    llvm::InlineAsm *rawFun = getRawFun(expression->getName());
    if (rawFun != nullptr) {
        vector<llvm::Value *>argValues;
        for (shared_ptr<Expression> &argumentExpression : expression->getArgumentExpressions()) {
            llvm::Value *argValue = valueForExpression(argumentExpression);
            argValues.push_back(argValue);
        }
        return builder->CreateCall(rawFun, llvm::ArrayRef(argValues));
    }

    markError(0, 0, format("Function \"{}\" not defined in scope", expression->getName()));
    return nullptr;
}

bool ModuleBuilder::setAlloca(string name, llvm::AllocaInst *alloca) {
    if (scopes.top().allocaMap[name] != nullptr) {
        markError(0, 0, format("Variable \"{}\" already defined", name));
        return false;
    }

    scopes.top().allocaMap[name] = alloca;
    return true;
}

llvm::AllocaInst* ModuleBuilder::getAlloca(string name) {
    stack<Scope> scopes = this->scopes;

    while (!scopes.empty()) {
        llvm::AllocaInst *alloca = scopes.top().allocaMap[name];
        if (alloca != nullptr)
            return alloca;
        scopes.pop();
    }

    markError(0, 0, format("Variable \"{}\" not defined in scope", name));
    return nullptr;
}

bool ModuleBuilder::setFun(string name, llvm::Function *fun) {
    if (scopes.top().funMap[name] != nullptr) {
        markError(0, 0, format("Function \"{}\" already defined in scope", name));
        return false;
    }

    scopes.top().funMap[name] = fun;
    return true;
}

llvm::Function* ModuleBuilder::getFun(string name) {
    stack<Scope> scopes = this->scopes;

    while (!scopes.empty()) {
        llvm::Function *fun = scopes.top().funMap[name];
        if (fun != nullptr)
            return fun;
        scopes.pop();
    }

    return nullptr;
}

bool ModuleBuilder::setRawFun(string name, llvm::InlineAsm *rawFun) {
    if (scopes.top().rawFunMap[name] != nullptr) {
        markError(0, 0, format("Raw function \"{}\" already defined in scope", name));
        return false;
    }
    
    scopes.top().rawFunMap[name] = rawFun;
    return true;
}

llvm::InlineAsm *ModuleBuilder::getRawFun(string name) {
    stack<Scope> scopes = this->scopes;

    while (!scopes.empty()) {
        llvm::InlineAsm *rawFun = scopes.top().rawFunMap[name];
        if (rawFun != nullptr)
            return rawFun;
        scopes.pop();
    }

    return nullptr;
}

bool ModuleBuilder::registerStruct(string structName, llvm::StructType *structType, vector<string> memberNames) {
    if (scopes.top().structTypeMap[structName] != nullptr) {
        markError(0, 0, format("Blob \"{}\" already defined in scope",structName));
        return false;
    }

    scopes.top().structTypeMap[structName] = structType;
    scopes.top().structMembersMap[structName] = memberNames;

    return true;
}

llvm::StructType *ModuleBuilder::getStructType(string structName) {
    stack<Scope> scopes = this->scopes;

    while (!scopes.empty()) {
        llvm::StructType *structType = scopes.top().structTypeMap[structName];
        if (structType != nullptr)
            return structType;
        scopes.pop();
    }

    return nullptr;
}

optional<int> ModuleBuilder::getMemberIndex(string structName, string memberName) {
        stack<Scope> scopes = this->scopes;

    while (!scopes.empty()) {
        if (scopes.top().structMembersMap.contains(structName)) {
            vector<string> memberNames = scopes.top().structMembersMap[structName];
            for (int i=0; i<memberNames.size(); i++) {
                if (memberNames[i].compare(memberName) == 0)
                    return i;
            }
        }
        scopes.pop();
    }

    return {};
}

llvm::Type *ModuleBuilder::typeForValueType(shared_ptr<ValueType> valueType, int count) {
    if (valueType == nullptr) {
        markError(0, 0, "Missing type");
        return nullptr;
    }

    switch (valueType->getKind()) {
        case ValueTypeKind::NONE:
            return typeVoid;
        case ValueTypeKind::BOOL:
            return typeBool;
        case ValueTypeKind::U8:
            return typeU8;
        case ValueTypeKind::U32:
            return typeU32;
        case ValueTypeKind::S8:
            return typeS8;
        case ValueTypeKind::S32:
            return typeS32;
        case ValueTypeKind::R32:
            return typeR32;
        case ValueTypeKind::DATA: {
            if (valueType->getSubType() == nullptr)
                return nullptr;
            if (valueType->getValueArg() > 0)
                count = valueType->getValueArg();
            return llvm::ArrayType::get(typeForValueType(valueType->getSubType(), count), count);
        }
        case ValueTypeKind::TYPE:
            return getStructType(valueType->getTypeName());
    }
}

void ModuleBuilder::markError(int line, int column, string message) {
    errors.push_back(Error::builderError(line, column, message));
}