#include "llvm_context.hpp"
#include <llvm/IR/DerivedTypes.h>

using namespace llvm;

void llvm_context::create_types() {
    stack_type = StructType::create(ctx, "stack");
    stack_ptr_type = PointerType::getUnqual(stack_type);
    tag_type = IntegerType::getInt8Ty(ctx);
    struct_types["node_base"] = StructType::create(ctx, "node_base");
    struct_types["node_app"] = StructType::create(ctx, "node_app");
    struct_types["node_num"] = StructType::create(ctx, "node_num");
    struct_types["node_global"] = StructType::create(ctx, "node_global");
    struct_types["node_ind"] = StructType::create(ctx, "node_ind");
    struct_types["node_data"] = StructType::create(ctx, "node_data");
    node_ptr_type = PointerType::getUnqual(struct_types.at("node_base"));
    function_type = FunctionType::get(Type::getVoidTy(ctx), { stack_ptr_type }, false);

    struct_types.at("node_base")->setBody(
            IntegerType::getInt32Ty(ctx)
    );
    struct_types.at("node_app")->setBody(
            struct_types.at("node_base"),
            node_ptr_type,
            node_ptr_type
    );
    struct_types.at("node_num")->setBody(
            struct_types.at("node_base"),
            IntegerType::getInt32Ty(ctx)
    );
    struct_types.at("node_global")->setBody(
            struct_types.at("node_base"),
            FunctionType::get(Type::getVoidTy(ctx), { stack_ptr_type }, false)
    );
    struct_types.at("node_ind")->setBody(
            struct_types.at("node_base"),
            node_ptr_type
    );
    struct_types.at("node_data")->setBody(
            struct_types.at("node_base"),
            IntegerType::getInt8Ty(ctx),
            PointerType::getUnqual(node_ptr_type)
    );
}

void llvm_context::create_functions() {
    auto void_type = Type::getVoidTy(ctx);
    auto sizet_type = IntegerType::get(ctx, sizeof(size_t) * 8);
    functions["stack_init"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_init",
            &module
    );
    functions["stack_free"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_free",
            &module
    );
    functions["stack_push"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type, node_ptr_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_push",
            &module
    );
    functions["stack_pop"] = Function::Create(
            FunctionType::get(node_ptr_type, { stack_ptr_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_pop",
            &module
    );
    functions["stack_peek"] = Function::Create(
            FunctionType::get(node_ptr_type, { stack_ptr_type, sizet_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_peek",
            &module
    );
    functions["stack_popn"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type, sizet_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_popn",
            &module
    );
    functions["stack_slide"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type, sizet_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_slide",
            &module
    );
    functions["stack_update"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type, sizet_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_update",
            &module
    );
    functions["stack_alloc"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type, sizet_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_alloc",
            &module
    );
    functions["stack_pack"] = Function::Create(
            FunctionType::get(void_type, { stack_ptr_type, sizet_type, tag_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_pack",
            &module
    );
    functions["stack_split"] = Function::Create(
            FunctionType::get(node_ptr_type, { stack_ptr_type, sizet_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "stack_split",
            &module
    );

    auto int32_type = IntegerType::getInt32Ty(ctx);
    functions["alloc_app"] = Function::Create(
            FunctionType::get(node_ptr_type, { node_ptr_type, node_ptr_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "alloc_app",
            &module
    );
    functions["alloc_num"] = Function::Create(
            FunctionType::get(node_ptr_type, { int32_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "alloc_num",
            &module
    );
    functions["alloc_global"] = Function::Create(
            FunctionType::get(node_ptr_type, { function_type, int32_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "alloc_global",
            &module
    );
    functions["alloc_ind"] = Function::Create(
            FunctionType::get(node_ptr_type, { node_ptr_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "alloc_ind",
            &module
    );

    functions["eval"] = Function::Create(
            FunctionType::get(node_ptr_type, { node_ptr_type }, false),
            Function::LinkageTypes::ExternalLinkage,
            "eval",
            &module
    );
}

ConstantInt* llvm_context::create_i8(int8_t i) {
    return ConstantInt::get(ctx, APInt(8, i));
}
ConstantInt* llvm_context::create_i32(int32_t i) {
    return ConstantInt::get(ctx, APInt(32, i));
}
ConstantInt* llvm_context::create_size(size_t i) {
    return ConstantInt::get(ctx, APInt(sizeof(size_t) * 8, i));
}

Value* llvm_context::create_pop(Function* f) {
    auto pop_f = functions.at("stack_pop");
    return builder.CreateCall(pop_f, { f->arg_begin() });
}
Value* llvm_context::create_peek(Function* f, Value* off) {
    auto peek_f = functions.at("stack_peek");
    return builder.CreateCall(peek_f, { f->arg_begin(), off });
}
void llvm_context::create_push(Function* f, Value* v) {
    auto push_f = functions.at("stack_push");
    builder.CreateCall(push_f, { f->arg_begin(), v });
}
void llvm_context::create_popn(Function* f, Value* off) {
    auto popn_f = functions.at("stack_popn");
    builder.CreateCall(popn_f, { f->arg_begin(), off });
}
void llvm_context::create_update(Function* f, Value* off) {
    auto update_f = functions.at("stack_update");
    builder.CreateCall(update_f, { f->arg_begin(), off });
}
void llvm_context::create_pack(Function* f, Value* c, Value* t) {
    auto pack_f = functions.at("stack_pack");
    builder.CreateCall(pack_f, { f->arg_begin(), c, t });
}
void llvm_context::create_split(Function* f, Value* c) {
    auto split_f = functions.at("stack_split");
    builder.CreateCall(split_f, { f->arg_begin(), c });
}
void llvm_context::create_slide(Function* f, Value* off) {
    auto slide_f = functions.at("stack_slide");
    builder.CreateCall(slide_f, { f->arg_begin(), off });
}
void llvm_context::create_alloc(Function* f, Value* n) {
    auto alloc_f = functions.at("stack_alloc");
    builder.CreateCall(alloc_f, { f->arg_begin(), n });
}

Value* llvm_context::create_eval(Value* e) {
    auto eval_f = functions.at("eval");
    return builder.CreateCall(eval_f, { e });
}

Value* llvm_context::unwrap_num(Value* v) {
    auto num_ptr_type = PointerType::getUnqual(struct_types.at("node_num"));
    auto cast = builder.CreatePointerCast(v, num_ptr_type);
    auto offset_0 = create_i32(0);
    auto offset_1 = create_i32(1);
    auto int_ptr = builder.CreateGEP(cast, { offset_0, offset_1 });
    return builder.CreateLoad(int_ptr);
}
Value* llvm_context::create_num(Value* v) {
    auto alloc_num_f = functions.at("alloc_num");
    return builder.CreateCall(alloc_num_f, { v });
}

Value* llvm_context::unwrap_data_tag(Value* v) {
    auto data_ptr_type = PointerType::getUnqual(struct_types.at("node_data"));
    auto cast = builder.CreatePointerCast(v, data_ptr_type);
    auto offset_0 = create_i32(0);
    auto offset_1 = create_i32(1);
    auto tag_ptr = builder.CreateGEP(cast, { offset_0, offset_1 });
    return builder.CreateLoad(tag_ptr);
}

Value* llvm_context::create_global(Value* f, Value* a) {
    auto alloc_global_f = functions.at("alloc_global");
    return builder.CreateCall(alloc_global_f, { f, a });
}

Value* llvm_context::create_app(Value* l, Value* r) {
    auto alloc_app_f = functions.at("alloc_app");
    return builder.CreateCall(alloc_app_f, { l, r });
}

llvm::Function* llvm_context::create_custom_function(std::string name, int32_t arity) {
    auto void_type = llvm::Type::getVoidTy(ctx);
    auto function_type =
        llvm::FunctionType::get(void_type, { stack_ptr_type }, false);
    auto new_function = llvm::Function::Create(
            function_type,
            llvm::Function::LinkageTypes::ExternalLinkage,
            "f_" + name,
            &module
    );
    auto start_block = llvm::BasicBlock::Create(ctx, "entry", new_function);

    auto new_custom_f = custom_function_ptr(new custom_function());
    new_custom_f->arity = arity;
    new_custom_f->function = new_function;
    custom_functions["f_" + name] = std::move(new_custom_f);

    return new_function;
}