#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() {
    body->find_free(free_variables);
    for(auto& param : params) {
        free_variables.erase(param);
    }
}

void definition_defn::insert_types(type_mgr& mgr, type_env_ptr& env, visibility v) {
    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);
    }
    env->bind(name, full_type, v);
}

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


global_function& definition_defn::into_global(global_scope& scope) {
    std::vector<std::string> all_params;
    for(auto& free : free_variables) {
        if(env->is_global(free)) continue;
        all_params.push_back(free);
    }
    all_params.insert(all_params.end(), params.begin(), params.end());
    body->translate(scope);
    return scope.add_function(name, std::move(all_params), std::move(body));
}

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::into_globals(global_scope& scope) {
    for(auto& constructor : constructors) {
        global_constructor& c = scope.add_constructor(
                constructor->name, constructor->tag, constructor->types.size());
        env->set_mangled_name(constructor->name, c.name);
    }
}

void definition_group::find_free(std::set<std::string>& into) {
    for(auto& def_pair : defs_defn) {
        def_pair.second->find_free();
        for(auto& free_var : def_pair.second->free_variables) {
            if(defs_defn.find(free_var) == defs_defn.end()) {
                into.insert(free_var);
            } else {
                def_pair.second->nearby_variables.insert(free_var);
            }
        }
    }
}

void definition_group::typecheck(type_mgr& mgr, type_env_ptr& env) {
    this->env = type_scope(env);

    for(auto& def_data : defs_data) {
        def_data.second->insert_types(this->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();
        dependency_graph.add_function(def_defn.second->name);

        for(auto& dependency : def_defn.second->nearby_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, this->env, vis);
        }
        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) {
            this->env->generalize(def_defnn_name, *group, mgr);
        }
    }
}