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

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

namespace codegen {
    Scope Scope::with_lvalue() {
        return Scope{ this->values, true };
    }
}

namespace AST {
    codegen::StackValue IntLiteralExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto ty = builder.builder->getInt32Ty();

        auto stack_type = new types::FundamentalType{ types::FundamentalTypeKind::Int };

        return codegen::StackValue{
            llvm::ConstantInt::get(ty, this->m_value),
             std::unique_ptr<types::Type>{stack_type}
        };
    }

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

    codegen::StackValue BinaryOperationExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto lvalued = scope.with_lvalue();
        auto lhs = this->m_lhs->codegen(builder, this->m_binop == types::BinOp::Assignment ? lvalued : scope);
        auto rhs = this->m_rhs->codegen(builder, scope);
        switch (this->m_binop) {
        case types::BinOp::Assignment:
            builder.builder->CreateStore(rhs.value, lhs.value, false);
            return rhs;
        case types::BinOp::Add:
            return codegen::StackValue{
                lhs.ty->add(builder, lhs.value, rhs.value),
                lhs.ty
            };
        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.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);
        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.value, ptr, false);
        }

        scope.values[this->m_name] = codegen::StackValue{ ptr, std::move(this->m_type) };
    }

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

        auto ret_ty = this->m_return_ty->codegen(builder);
        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;
    }
}

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