#include "codegen.h"
#include "ast.h"

#include <llvm/IR/Module.h>
#include <llvm/IR/Verifier.h>
#include <memory>
#include <iostream>

namespace AST {
    llvm::Value* IntLiteralExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto ty = builder.builder->getInt32Ty();
        return llvm::ConstantInt::get(ty, this->m_value);
    }

    llvm::Value* ValueReferenceExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto value = scope.values.find(this->m_name);
        if (value != scope.values.end()) {
            return builder.builder->CreateLoad(value->second.ty, value->second.value);
        }
        else {
            throw new std::runtime_error("Value " + this->m_name + " not found");
        }
    }

    llvm::Value* BinaryOperationExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto lhs = this->m_lhs->codegen(builder, scope);
        auto rhs = this->m_rhs->codegen(builder, scope);
        switch (this->m_binop) {
        case BinOp::Assignment:
            builder.builder->CreateStore(rhs, lhs, false);
            return lhs;
        default:
            throw std::runtime_error("invalid binop");
        }
    }

    void ReturnStatement::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        if (!builder.block)
            return;

        builder.builder->SetInsertPoint(builder.block);

        auto value = this->m_expr->codegen(builder, scope);
        builder.builder->CreateRet(value);
    }

    void ExpressionStatement::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        if (!builder.block)
            return;

        builder.builder->SetInsertPoint(builder.block);
        this->m_expr->codegen(builder, scope);
    }

    void InitializationStatement::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        if (!builder.block)
            return;

        builder.builder->SetInsertPoint(builder.block);

        auto ty = this->m_type->codegen(builder, scope);
        auto ptr = builder.builder->CreateAlloca(ty);
        if (this->m_expr.has_value()) {
            auto value = this->m_expr->get()->codegen(builder, scope);
            builder.builder->CreateStore(value, ptr, false);
        }

        scope.values[this->m_name] = codegen::StackValue{ ptr, ty };
    }

    void Function::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        codegen::Scope inner_scope{ scope };

        auto ret_ty = this->m_return_ty->codegen(builder, inner_scope);
        std::vector<llvm::Type*> params{};


        auto fn_ty = llvm::FunctionType::get(ret_ty, params, false);
        auto function = llvm::Function::Create(
            fn_ty,
            llvm::GlobalValue::LinkageTypes::ExternalLinkage,
            this->m_name,
            builder.mod.get()
        );
        auto BB = llvm::BasicBlock::Create(*builder.context, "entry", function, nullptr);
        builder.block = BB;

        for (auto& statement : this->m_statements) {
            statement->codegen(builder, inner_scope);
        }

        llvm::verifyFunction(*function);

        builder.block = nullptr;
    }

    llvm::Type* FundamentalType::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        switch (this->m_ty) {
        case FundamentalTypeKind::Int:
            return builder.builder->getInt32Ty();
        default:
            return builder.builder->getVoidTy();
        }
    }
}