#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->structs, 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& scope) {
        auto ty = this->m_ty->codegen(builder, scope.structs);

        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_shared<types::FundamentalType>(true, types::FundamentalTypeKind::Char),
            static_cast<uint32_t>(this->m_value.size()) + 1,
            true
        };
        return std::shared_ptr<types::Type> {stack_type};
    }

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

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

        if (scope.is_lvalue) {
            return codegen::StackValue{
                global_str,
                std::unique_ptr<types::Type>{
                    new types::PointerType { false, std::shared_ptr<types::Type> {stack_type } }
                },
            };
        }
        else {
            return codegen::StackValue{
                builder.builder->CreateLoad(stack_type->codegen(builder, scope.structs), 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()) {
            if (value->second.ty->m_kind == types::TypeKind::Pointer) {
                auto ptr_ty = dynamic_cast<types::PointerType*>(value->second.ty.get());
                return ptr_ty->m_inner;
            }
            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, scope.structs);
                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, scope.structs)), 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::Array
            && this->m_ty->m_kind == types::TypeKind::Pointer) {
            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) {
                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{ false, 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, scope.structs);
            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());

            auto gep_value = builder.builder->CreateGEP(ptr_ty->m_inner->codegen(builder, scope.structs), value.value, idx_list, "GEP");

            auto loaded = value.ty->load(builder, gep_value, scope.structs);
            return codegen::StackValue{
                loaded.first,
                loaded.second
            };
        }
        else if (value.ty->m_kind == types::TypeKind::Array) {
            auto array_ty = dynamic_cast<types::ArrayType*>(value.ty.get());
            if (array_ty->m_raw) {
                throw CompileError("Tried indexing a raw array", this->m_meta);
            }

            auto gep_value = builder.builder->CreateGEP(array_ty->m_inner->codegen(builder, scope.structs), value.value, idx_list, "GEP");

            auto ptr_ty = std::shared_ptr<types::Type>{
                new types::PointerType { array_ty->m_inner->m_const, array_ty->m_inner}
            };

            if (scope.is_lvalue) {
                return codegen::StackValue{
                    gep_value,
                    ptr_ty,
                };
            }
            else {
                auto loaded = ptr_ty->load(builder, gep_value, scope.structs);
                return codegen::StackValue{
                    loaded.first,
                    loaded.second
                };
            }
        }

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


    std::shared_ptr<types::Type> FieldAccessExpression::get_codegen_type(codegen::Scope& scope) {
        auto ty = this->m_expr->get_codegen_type(scope);
        if (ty->m_kind == types::TypeKind::Struct) {
            auto struct_ty = dynamic_cast<types::StructType*>(ty.get());
            if (struct_ty->m_fields) {
                for (auto& field : *struct_ty->m_fields) {
                    if (field.first == this->m_field)
                        return field.second;
                }
                throw CompileError("Unknown field", this->m_meta);
            }
            throw CompileError("Cannot access a field of opaque struct", this->m_meta);
        }
        else {
            throw CompileError("Tried accessing a non-struct", this->m_meta);
        }
    }

    codegen::StackValue FieldAccessExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto lvalued = scope.with_lvalue();
        auto struct_ptr = this->m_expr->codegen(builder, lvalued);
        if (struct_ptr.ty->m_kind == types::TypeKind::Pointer) {
            auto ptr_ty = dynamic_cast<types::PointerType*>(struct_ptr.ty.get());
            if (ptr_ty->m_inner->m_kind == types::TypeKind::Struct) {
                auto struct_ty = dynamic_cast<types::StructType*>(ptr_ty->m_inner.get());

                int idx = -1;
                auto field_ty = std::shared_ptr<types::Type>{};

                for (int i = 0; i < static_cast<int>(struct_ty->m_fields->size()); i++) {
                    auto field = (*struct_ty->m_fields)[i];
                    if (field.first == this->m_field) {
                        idx = i;
                        field_ty = field.second;
                        break;
                    }
                }

                auto gep = builder.builder->CreateStructGEP(
                    ptr_ty->m_inner->codegen(builder, scope.structs), struct_ptr.value, idx);

                auto ptr_ty = std::shared_ptr<types::Type>{
                    new types::PointerType { field_ty->m_const, field_ty }
                };

                if (scope.is_lvalue) {
                    return codegen::StackValue{
                        gep,
                        ptr_ty
                    };
                }
                else {
                    auto loaded = ptr_ty->load(builder, gep, scope.structs);
                    return codegen::StackValue{ loaded.first, loaded.second };
                }
            }
            else {
                throw CompileError("Tried field-accessing a non-struct-pointer", this->m_meta);
            }
        }
        else {
            throw CompileError("Tried field-accessing a non-pointer", this->m_meta);
        }
    }



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

    codegen::StackValue ListInitializerExpression::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto value_ptr = builder.builder->CreateAlloca(this->m_ty->codegen(builder, scope.structs));

        if (this->m_ty->m_kind == types::TypeKind::Array) {
            auto array_ty = dynamic_cast<types::ArrayType*>(this->m_ty.get());

            int counter = 0;
            for (auto& expr : this->m_expressions) {
                std::vector<llvm::Value*> indices{};
                indices.push_back(llvm::ConstantInt::get(builder.builder->getInt32Ty(), counter++));
                auto gep = builder.builder->CreateGEP(
                    array_ty->m_inner->codegen(builder, scope.structs), value_ptr, indices, "GEP");
                builder.builder->CreateStore(expr->codegen(builder, scope).value, gep);
            }

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

            if (scope.is_lvalue) {
                return codegen::StackValue{ value_ptr, ptr_ty };
            }
            else {
                auto loaded = ptr_ty->load(builder, value_ptr, scope.structs);
                return codegen::StackValue{ loaded.first, loaded.second };
            }
        }
        else if (this->m_ty->m_kind == types::TypeKind::Struct) {
            auto struct_ty = dynamic_cast<types::StructType*>(this->m_ty.get());

            for (int i = 0; i < static_cast<int>(this->m_expressions.size()); i++) {
                std::vector<llvm::Value*> indices{};
                auto gep = builder.builder->CreateStructGEP(
                    struct_ty->codegen(builder, scope.structs), value_ptr, i, "struct_gep"
                );
                builder.builder->CreateStore(
                    this->m_expressions[i]->codegen(builder, scope).value, gep);
            }

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

            if (scope.is_lvalue) {
                return codegen::StackValue{ value_ptr, ptr_ty };
            }
            else {
                auto loaded = ptr_ty->load(builder, value_ptr, scope.structs);
                return codegen::StackValue{ loaded.first, loaded.second };
            }

        }
        else {
            throw CompileError("Tried to list-initialize a non-array!", 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);

        if (this->m_type->m_kind == types::TypeKind::Array) {
            auto array_ty = dynamic_cast<types::ArrayType*>(this->m_type.get());
            if (!array_ty->m_raw) {
                auto raw_ty = types::ArrayType{ array_ty->m_inner, array_ty->m_size, true };
                auto raw_llvm_ty = raw_ty.codegen(builder, scope.structs);
                auto ptr = builder.builder->CreateAlloca(raw_llvm_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, this->m_type };
                return;
            }
        }

        auto ty = this->m_type->codegen(builder, scope.structs);
        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->m_const, 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{ true, ret_ty_ptr, param_ty_ptrs, this->m_is_vararg }
        };

        auto fn_ty = fn_ty_ptr->codegen(builder, scope.structs);
        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::Array) {
                    auto array_ty = dynamic_cast<types::ArrayType*>(param.second.get());
                    if (!array_ty->m_raw) {
                        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,
                            };
                        }
                        continue;
                    }
                }
                builder.builder->SetInsertPoint(BB);
                auto arg_ptr = builder.builder->CreateAlloca(param_ty_ptrs[counter]->codegen(builder, scope.structs));
                auto param_ty_ptr = std::shared_ptr<types::Type>{
                    new types::PointerType { true, 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;
    }

    void TopLevelTypedef::codegen(codegen::Builder& builder, codegen::Scope& scope) {
        auto ty = this->m_ty->codegen(builder, scope.structs);
        if (this->m_ty->m_kind == types::TypeKind::Struct) {
            auto struct_ty = dynamic_cast<types::StructType*>(this->m_ty.get());
            if (struct_ty->m_name) {
                scope.structs[*struct_ty->m_name] = ty;
            }
        }
    }
}

namespace types {
    llvm::Type* FundamentalType::codegen(codegen::Builder& builder, codegen::TypeMap&) {
        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, codegen::TypeMap& structs) {
        std::vector<llvm::Type*> params{};

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

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

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

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

    llvm::Type* ArrayType::codegen(codegen::Builder& builder, codegen::TypeMap& structs) {
        if (this->m_raw)
            return llvm::ArrayType::get(this->m_inner->codegen(builder, structs), this->m_size);
        return llvm::PointerType::get(*builder.context, 0);
    }

    llvm::Type* StructType::codegen(codegen::Builder& builder, codegen::TypeMap& structs) {
        if (this->m_is_ref) {
            if (this->m_name) {
                return structs[*this->m_name];
            }
            else {
                throw CompileError("reference to nonexistant struct", {});
            }
        }
        if (this->m_fields) {
            std::vector<llvm::Type*> fields{};
            for (auto& field : *this->m_fields) {
                fields.push_back(field.second->codegen(builder, structs));
            }

            auto ty = llvm::StructType::create(*builder.context, fields);
            if (this->m_name)
                ty->setName(*this->m_name);
            return ty;
        }
        else {
            auto ty = llvm::StructType::create(*builder.context);
            if (this->m_name)
                ty->setName(*this->m_name);
            return ty;
        }
    }

    // llvm::Type* StructRef::codegen(codegen::Builder&, codegen::TypeMap& structs) {
    //     return structs[this->m_name];
    // }
}