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

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

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

namespace AST {
    std::shared_ptr<types::Type> IntLiteralExpression::get_codegen_type(codegen::Scope&) {
        return this->m_ty;
    }

    codegen::StackValue IntLiteralExpression::codegen(codegen::Builder& builder, codegen::Scope&) {
        auto ty = this->m_ty->codegen(builder);

        return codegen::StackValue{
            llvm::ConstantInt::get(ty, this->m_value),
            this->m_ty,
        };
    }

    std::shared_ptr<types::Type> StringLiteralExpression::get_codegen_type(codegen::Scope&) {
        auto stack_type = new types::ArrayType{
            std::make_unique<types::FundamentalType>(types::FundamentalTypeKind::Char),
            static_cast<uint32_t>(this->m_value.size()) + 1
        };
        return std::shared_ptr<types::Type> {stack_type};
    }

    codegen::StackValue StringLiteralExpression::codegen(codegen::Builder& builder, codegen::Scope&) {
        auto stack_type = new types::ArrayType{
            std::make_unique<types::FundamentalType>(types::FundamentalTypeKind::Char),
            static_cast<uint32_t>(this->m_value.size()) + 1
        };

        auto str = llvm::StringRef{ this->m_value.c_str() };
        auto global_str = builder.builder->CreateGlobalString(str);

        return codegen::StackValue{
            builder.builder->CreateLoad(stack_type->codegen(builder), global_str, "literal"),
            std::unique_ptr<types::Type>{stack_type},
        };
    }

    std::shared_ptr<types::Type> ValueReferenceExpression::get_codegen_type(codegen::Scope& scope) {
        auto value = scope.values.find(this->m_name);
        if (value != scope.values.end()) {
            return value->second.ty;
        }
        else {
            throw CompileError("Value " + this->m_name + " not found", this->m_meta);
        }
    }

    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 {
                auto loaded = value->second.ty->load(builder, value->second.value);
                return codegen::StackValue{
                    loaded.first,
                    loaded.second
                };
            }
        }
        else {
            throw CompileError("Value " + this->m_name + " not found", this->m_meta);
        }
    }

    std::shared_ptr<types::Type> BinaryOperationExpression::get_codegen_type(codegen::Scope& scope) {
        auto lhs = this->m_lhs->get_codegen_type(scope);
        auto rhs = this->m_rhs->get_codegen_type(scope);
        try {
            switch (this->m_binop) {
            case types::BinOp::Assignment:
                return rhs;
            default:
                auto binop = types::find_binop(
                    scope.binops,
                    lhs,
                    this->m_binop,
                    rhs);
                if (binop) {
                    return binop->result(*binop, lhs, rhs);
                }
                throw CompileError("invalid binop", this->m_meta);
            }
        }
        catch (std::runtime_error& error) {
            throw CompileError(error.what(), this->m_meta);
        }
    }

    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);
        try {
            switch (this->m_binop) {
            case types::BinOp::Assignment:
                builder.builder->CreateStore(rhs.value, lhs.value, false);
                return rhs;
            default:
                auto binop = types::find_binop(
                    scope.binops,
                    lhs.ty,
                    this->m_binop,
                    rhs.ty);
                if (binop) {
                    return codegen::StackValue{
                        binop->codegen(builder, lhs.value, rhs.value),
                        binop->result(*binop, lhs.ty, rhs.ty)
                    };
                }
                throw CompileError("invalid binop", this->m_meta);
            }
        }
        catch (std::runtime_error& error) {
            throw CompileError(error.what(), this->m_meta);
        }
    }

    std::shared_ptr<types::Type> FunctionCallExpression::get_codegen_type(codegen::Scope& scope) {
        auto fn_ty = this->m_fn_expr->get_codegen_type(scope);
        return *fn_ty->return_type();
    }

    codegen::StackValue FunctionCallExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        std::vector<llvm::Value*> args{};
        for (auto& arg : this->m_args) {
            args.push_back(arg->codegen(builder, scope).value);
        }
        auto function = this->m_fn_expr->codegen(builder, scope);

        auto value = builder.builder->CreateCall(llvm::dyn_cast<llvm::FunctionType>(function.ty->codegen(builder)), function.value, args, "call");
        return codegen::StackValue{
            value,
            *function.ty->return_type(),
        };
    }

    std::shared_ptr<types::Type> CastExpression::get_codegen_type(codegen::Scope&) {
        return this->m_ty;
    }

    codegen::StackValue CastExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto expr_ty = this->m_expr->get_codegen_type(scope);
        if (expr_ty->m_kind == types::TypeKind::Pointer && this->m_ty->m_kind == types::TypeKind::Pointer) {
            std::cout << expr_ty->formatted() << std::endl;
            auto lvalued = scope.with_lvalue();
            auto expr = this->m_expr->codegen(builder, lvalued);
            auto cast = types::find_cast(scope.casts, expr.ty, this->m_ty);
            if (cast) {
                return codegen::StackValue{
                    cast->codegen(builder, cast->target_ty, expr.value),
                    cast->target_ty
                };
            }
            else {
                return expr;
            }
        }
        else {
            auto expr = this->m_expr->codegen(builder, scope);
            auto cast = types::find_cast(scope.casts, expr.ty, this->m_ty);
            if (cast) {
                // if (cast->target_ty->m_kind == types::TypeKind::Pointer && cast->casted_ty->m_kind == types::TypeKind::Array) {
                //     auto allocated = builder.builder->CreateAlloca(cast->casted_ty->codegen(builder));
                //     builder.builder->CreateStore(expr.value, allocated);
                //     return codegen::StackValue{
                //         cast->codegen(builder, cast->target_ty, allocated),
                //         cast->target_ty
                //     };
                // }
                return codegen::StackValue{
                    cast->codegen(builder, cast->target_ty, expr.value),
                    cast->target_ty
                };
            }
            return expr;
        }
    }

    std::shared_ptr<types::Type> RefExpression::get_codegen_type(codegen::Scope& scope) {
        return std::shared_ptr<types::Type> {
            new types::PointerType{ this->m_expr->get_codegen_type(scope) }
        };
    }

    codegen::StackValue RefExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto with_lvalue = scope.with_lvalue();
        return this->m_expr->codegen(builder, with_lvalue);
    }

    std::shared_ptr<types::Type> DerefExpression::get_codegen_type(codegen::Scope& scope) {
        auto ty = this->m_expr->get_codegen_type(scope);
        if (ty->m_kind == types::TypeKind::Pointer) {
            auto ptr_ty = dynamic_cast<types::PointerType*>(ty.get());
            return ptr_ty->m_inner;
        }
        else {
            throw new CompileError("Tried to deref a non-pointer!", this->m_meta);
        }
    }

    codegen::StackValue DerefExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto value = this->m_expr->codegen(builder, scope);
        if (value.ty->m_kind == types::TypeKind::Pointer) {
            auto loaded = value.ty->load(builder, value.value);
            return codegen::StackValue{
                loaded.first,
                loaded.second
            };
        }
        else {
            throw new CompileError("Tried to deref a non-pointer!", this->m_meta);
        }
    }

    std::shared_ptr<types::Type> IndexAccessExpression::get_codegen_type(codegen::Scope& scope) {
        auto ty = this->m_expr->get_codegen_type(scope);
        if (ty->m_kind == types::TypeKind::Array) {
            auto ptr_ty = dynamic_cast<types::ArrayType*>(ty.get());
            return ptr_ty->m_inner;
        }
        else if (ty->m_kind == types::TypeKind::Pointer) {
            auto ptr_ty = dynamic_cast<types::PointerType*>(ty.get());
            return ptr_ty->m_inner;
        }
        else {
            throw CompileError("Tried indexing a non-pointer", this->m_meta);
        }
    }


    codegen::StackValue IndexAccessExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto lvalue_scope = scope.with_lvalue();
        auto value = this->m_expr->codegen(builder, lvalue_scope);
        std::shared_ptr<types::Type> inner_ty{};

        std::vector<llvm::Value*> idx_list{ };
        idx_list.push_back(llvm::ConstantInt::get(builder.builder->getInt32Ty(), this->m_num));

        if (value.ty->m_kind == types::TypeKind::Pointer) {
            auto ptr_ty = dynamic_cast<types::PointerType*>(value.ty.get());

            if (ptr_ty->m_inner->m_kind == types::TypeKind::Pointer) {
                auto inner_ptr_ty = dynamic_cast<types::PointerType*>(ptr_ty->m_inner.get());

                auto gep_value = builder.builder->CreateGEP(inner_ptr_ty->m_inner->codegen(builder), value.value, idx_list, "GEP");
                if (scope.is_lvalue) {
                    return codegen::StackValue{
                        gep_value,
                        value.ty,
                    };
                }
                else {
                    auto loaded = value.ty->load(builder, gep_value);
                    return codegen::StackValue{
                        loaded.first,
                        loaded.second
                    };
                }
            }
            else if (ptr_ty->m_inner->m_kind == types::TypeKind::Array) {
                std::cout << ptr_ty->m_inner->formatted() << std::endl;

                // Must be an array otherwise
                auto arr_ty = dynamic_cast<types::ArrayType*>(ptr_ty->m_inner.get());
                auto gep_value = builder.builder->CreateGEP(arr_ty->m_inner->codegen(builder), value.value, idx_list, "GEP");
                if (scope.is_lvalue) {
                    return codegen::StackValue{
                        gep_value,
                        value.ty,
                    };
                }
                else {
                    auto inner_ptr_ty = types::PointerType{ arr_ty->m_inner };
                    auto loaded = inner_ptr_ty.load(builder, gep_value);
                    return codegen::StackValue{
                        loaded.first,
                        loaded.second
                    };
                }
            }
            else {
                auto gep_value = builder.builder->CreateGEP(ptr_ty->m_inner->codegen(builder), value.value, idx_list, "GEP");
                if (scope.is_lvalue) {
                    return codegen::StackValue{
                        gep_value,
                        value.ty,
                    };
                }
                else {
                    auto loaded = value.ty->load(builder, gep_value);
                    return codegen::StackValue{
                        loaded.first,
                        loaded.second
                    };
                }
            }
        }

        throw CompileError("Tried indexing a non-pointer", this->m_meta);
    }


    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);
        }

        auto ptr_ty = std::shared_ptr<types::Type>{ new types::PointerType{ this->m_type } };

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

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

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

        auto condition = this->m_condition->codegen(builder, scope);

        auto function = builder.block->getParent();
        auto then_block = llvm::BasicBlock::Create(*builder.context, "then", function);

        std::optional<llvm::BasicBlock*> else_block{};
        if (this->m_else.has_value())
            else_block = llvm::BasicBlock::Create(*builder.context, "else", function);

        auto after_block = llvm::BasicBlock::Create(*builder.context, "after", function);

        builder.builder->CreateCondBr(condition.value, then_block, else_block.value_or(after_block));

        builder.block = then_block;
        builder.builder->SetInsertPoint(then_block);
        this->m_then->codegen(builder, scope);
        builder.builder->CreateBr(after_block);

        if (else_block.has_value()) {
            builder.block = *else_block;
            builder.builder->SetInsertPoint(*else_block);
            this->m_else->get()->codegen(builder, scope);
            builder.builder->CreateBr(after_block);
        }

        builder.block = after_block;
        builder.builder->SetInsertPoint(after_block);
    }


    void Function::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        std::shared_ptr<types::Type> ret_ty_ptr{ this->m_return_ty };
        std::vector<std::shared_ptr<types::Type>> param_ty_ptrs{};

        for (auto& param : this->m_params) {
            param_ty_ptrs.push_back(param.second);
        }

        auto fn_ty_ptr = std::shared_ptr<types::Type>{ new types::FunctionType{ ret_ty_ptr, param_ty_ptrs, this->m_is_vararg } };

        auto fn_ty = fn_ty_ptr->codegen(builder);
        auto function = llvm::Function::Create(
            llvm::dyn_cast<llvm::FunctionType>(fn_ty),
            llvm::GlobalValue::LinkageTypes::ExternalLinkage,
            this->m_name,
            builder.mod.get()
        );

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

        if (this->m_statements) {
            auto BB = llvm::BasicBlock::Create(*builder.context, "entry", function, nullptr);
            builder.block = BB;

            codegen::Scope inner_scope{ scope };

            int counter = 0;
            for (auto& param : this->m_params) {
                if (param.second->m_kind == types::TypeKind::Pointer) {
                    auto ty = param_ty_ptrs[counter];
                    auto arg = function->getArg(counter++);
                    if (param.first) {
                        arg->setName(*param.first);
                        inner_scope.values[*param.first] = codegen::StackValue{
                            arg,
                            ty,
                        };
                    }
                }
                else {
                    builder.builder->SetInsertPoint(BB);
                    auto arg_ptr = builder.builder->CreateAlloca(param_ty_ptrs[counter]->codegen(builder));
                    auto param_ty_ptr = std::shared_ptr<types::Type>{
                        new types::PointerType { param_ty_ptrs[counter]}
                    };
                    auto arg = function->getArg(counter++);
                    builder.builder->CreateStore(arg, arg_ptr);
                    if (param.first) {
                        arg->setName(*param.first);
                        inner_scope.values[*param.first] = codegen::StackValue{
                            arg_ptr,
                            param_ty_ptr,
                        };
                    }
                }
            }

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

        builder.builder->CreateRetVoid();

        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();
        case FundamentalTypeKind::Bool:
            return builder.builder->getInt1Ty();
        case FundamentalTypeKind::Char:
            return builder.builder->getInt8Ty();
        case FundamentalTypeKind::Void:
            return builder.builder->getVoidTy();
        default:
            return builder.builder->getVoidTy();
        }
    }

    llvm::Type* FunctionType::codegen(codegen::Builder& builder) {
        std::vector<llvm::Type*> params{};

        for (auto& param : this->m_param_tys) {
            params.push_back(param->codegen(builder));
        }

        auto ret_ty = this->m_ret_ty->codegen(builder);

        return llvm::FunctionType::get(ret_ty, params, this->m_vararg);
    }

    llvm::Type* PointerType::codegen(codegen::Builder& builder) {
        return llvm::PointerType::get(*builder.context, 0);
    }

    llvm::Type* ArrayType::codegen(codegen::Builder& builder) {
        return llvm::ArrayType::get(this->m_inner->codegen(builder), this->m_size);
    }
}