#include "definition.hpp"
#include "error.hpp"
#include "ast.hpp"
#include "instruction.hpp"
#include "llvm_context.hpp"
#include "type.hpp"
#include "type_env.hpp"
#include "graph.hpp"
#include <llvm/IR/DerivedTypes.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Type.h>

void definition_defn::find_free(type_mgr& mgr, type_env_ptr& env) {
    this->env = env;

    var_env = type_scope(env);
    return_type = mgr.new_type();
    full_type = return_type;

    for(auto it = params.rbegin(); it != params.rend(); it++) {
        type_ptr param_type = mgr.new_type();
        full_type = type_ptr(new type_arr(param_type, full_type));
        var_env->bind(*it, param_type);
    }

    body->find_free(mgr, var_env, free_variables);
}

void definition_defn::insert_types(type_mgr& mgr) {
    env->bind(name, full_type);
}

void definition_defn::typecheck(type_mgr& mgr) {
    type_ptr body_type = body->typecheck(mgr);
    mgr.unify(return_type, body_type);
}

void definition_defn::compile() {
    env_ptr new_env = env_ptr(new env_offset(0, nullptr));
    for(auto it = params.rbegin(); it != params.rend(); it++) {
        new_env = env_ptr(new env_var(*it, new_env));
    }
    body->compile(new_env, instructions);
    instructions.push_back(instruction_ptr(new instruction_update(params.size())));
    instructions.push_back(instruction_ptr(new instruction_pop(params.size())));
}

void definition_defn::declare_llvm(llvm_context& ctx) {
    generated_function = ctx.create_custom_function(name, params.size());
}

void definition_defn::generate_llvm(llvm_context& ctx) {
    ctx.builder.SetInsertPoint(&generated_function->getEntryBlock());
    for(auto& instruction : instructions) {
        instruction->gen_llvm(ctx, generated_function);
    }
    ctx.builder.CreateRetVoid();
}

void definition_data::insert_types(type_env_ptr& env) {
    this->env = env;
    env->bind_type(name, type_ptr(new type_data(name, vars.size())));
}

void definition_data::insert_constructors() const {
    type_ptr this_type_ptr = env->lookup_type(name);
    type_data* this_type = static_cast<type_data*>(this_type_ptr.get());
    int next_tag = 0;

    std::set<std::string> var_set;
    type_app* return_app = new type_app(std::move(this_type_ptr));
    type_ptr return_type(return_app);
    for(auto& var : vars) {
        if(var_set.find(var) != var_set.end()) throw 0;
        var_set.insert(var);
        return_app->arguments.push_back(type_ptr(new type_var(var)));
    }

    for(auto& constructor : constructors) {
        constructor->tag = next_tag;
        this_type->constructors[constructor->name] = { next_tag++ };

        type_ptr full_type = return_type;
        for(auto it = constructor->types.rbegin(); it != constructor->types.rend(); it++) {
            type_ptr type = (*it)->to_type(var_set, env);
            full_type = type_ptr(new type_arr(type, full_type));
        }

        type_scheme_ptr full_scheme(new type_scheme(std::move(full_type)));
        full_scheme->forall.insert(full_scheme->forall.begin(), vars.begin(), vars.end());
        env->bind(constructor->name, full_scheme);
    }
}

void definition_data::generate_llvm(llvm_context& ctx) {
    for(auto& constructor : constructors) {
        auto new_function =
            ctx.create_custom_function(constructor->name, constructor->types.size());
        std::vector<instruction_ptr> instructions;
        instructions.push_back(instruction_ptr(
                new instruction_pack(constructor->tag, constructor->types.size())
        ));
        instructions.push_back(instruction_ptr(new instruction_update(0)));
        ctx.builder.SetInsertPoint(&new_function->getEntryBlock());
        for (auto& instruction : instructions) {
            instruction->gen_llvm(ctx, new_function);
        }
        ctx.builder.CreateRetVoid();
    }
}

void definition_group::find_free(type_mgr& mgr, type_env_ptr& env, std::set<std::string>& into) {
    this->env = type_scope(env);

    for(auto& def_pair : defs_defn) {
        def_pair.second->find_free(mgr, env);
        std::set<std::string> local_dependencies;
        for(auto& free_var : def_pair.second->free_variables) {
            if(defs_defn.find(free_var) == defs_defn.end()) {
                into.insert(free_var);
            } else {
                local_dependencies.insert(free_var);
            }
        }
        std::swap(def_pair.second->free_variables, local_dependencies);
    }
}

void definition_group::typecheck(type_mgr& mgr) {
    for(auto& def_data : defs_data) {
        def_data.second->insert_types(env);
    }
    for(auto& def_data : defs_data) {
        def_data.second->insert_constructors();
    }

    function_graph dependency_graph;

    for(auto& def_defn : defs_defn) {
        def_defn.second->find_free(mgr, env);
        dependency_graph.add_function(def_defn.second->name);

        for(auto& dependency : def_defn.second->free_variables) {
            if(defs_defn.find(dependency) == defs_defn.end())
                throw 0;
            dependency_graph.add_edge(def_defn.second->name, dependency);
        }
    }

    std::vector<group_ptr> groups = dependency_graph.compute_order();
    for(auto it = groups.rbegin(); it != groups.rend(); it++) {
        auto& group = *it;
        for(auto& def_defnn_name : group->members) {
            auto& def_defn = defs_defn.find(def_defnn_name)->second;
            def_defn->insert_types(mgr);
        }
        for(auto& def_defnn_name : group->members) {
            auto& def_defn = defs_defn.find(def_defnn_name)->second;
            def_defn->typecheck(mgr);
        }
        for(auto& def_defnn_name : group->members) {
            env->generalize(def_defnn_name, mgr);
        }
    }
}