#include "compiler.hpp"
#include "binop.hpp"
#include "error.hpp"
#include "global_scope.hpp"
#include "parse_driver.hpp"
#include "type.hpp"
#include "type_env.hpp"

#include "llvm/IR/LegacyPassManager.h"
#include "llvm/IR/Verifier.h"
#include "llvm/Support/TargetSelect.h"
#include "llvm/Support/TargetRegistry.h"
#include "llvm/Support/raw_ostream.h"
#include "llvm/Support/FileSystem.h"
#include "llvm/Target/TargetOptions.h"
#include "llvm/Target/TargetMachine.h"

void compiler::add_default_types() {
    global_env->bind_type("Int", type_ptr(new type_base("Int")));
    global_env->bind_type("Bool", type_ptr(new type_base("Bool")));
}

void compiler::add_binop_type(binop op, type_ptr type) {
    auto name = mng.new_mangled_name(op_action(op));
    global_env->bind(op_name(op), std::move(type), visibility::global);
    global_env->set_mangled_name(op_name(op), name);
}

void compiler::add_default_function_types() {
    type_ptr int_type = global_env->lookup_type("Int");
    assert(int_type != nullptr);
    type_ptr int_type_app = type_ptr(new type_app(int_type));

    type_ptr bool_type = global_env->lookup_type("Bool");
    assert(bool_type != nullptr);
    type_ptr bool_type_app = type_ptr(new type_app(bool_type));

    type_ptr closed_int_op_type(
            new type_arr(int_type_app, type_ptr(new type_arr(int_type_app, int_type_app))));
    type_ptr compare_int_op_type(
            new type_arr(int_type_app, type_ptr(new type_arr(int_type_app, bool_type_app))));

    constexpr binop closed_ops[] = { PLUS, MINUS, TIMES, DIVIDE, MODULO };
    constexpr binop compare_ops[] = { EQUALS, LESS_EQUALS };
    for(auto& op : closed_ops) add_binop_type(op, closed_int_op_type);
    for(auto& op : compare_ops) add_binop_type(op, compare_int_op_type);

    for(auto name : { "True", "False" }) {
        global_env->bind(name, bool_type_app, visibility::global);
        global_env->set_mangled_name(name, mng.new_mangled_name(name));
    }
}

void compiler::parse() {
    if(!driver())
        throw compiler_error("failed to open file");
}

void compiler::typecheck() {
    std::set<std::string> free_variables;
    global_defs.find_free(free_variables);
    global_defs.typecheck(type_m, global_env);
}

void compiler::translate() {
    for(auto& data : global_defs.defs_data) {
        data.second->into_globals(global_scp);
    }
    for(auto& defn : global_defs.defs_defn) {
        auto& function = defn.second->into_global(global_scp);
        function.body->env->get_parent()->set_mangled_name(defn.first, function.name);
    }
}

void compiler::compile() {
    global_scp.compile();
}

void compiler::create_llvm_binop(binop op) {
    auto new_function =
        ctx.create_custom_function(global_env->get_mangled_name(op_name(op)), 2);
    std::vector<instruction_ptr> instructions;
    instructions.push_back(instruction_ptr(new instruction_push(1)));
    instructions.push_back(instruction_ptr(new instruction_eval()));
    instructions.push_back(instruction_ptr(new instruction_push(1)));
    instructions.push_back(instruction_ptr(new instruction_eval()));
    instructions.push_back(instruction_ptr(new instruction_binop(op)));
    instructions.push_back(instruction_ptr(new instruction_update(2)));
    instructions.push_back(instruction_ptr(new instruction_pop(2)));
    ctx.get_builder().SetInsertPoint(&new_function->getEntryBlock());
    for(auto& instruction : instructions) {
        instruction->gen_llvm(ctx, new_function);
    }
    ctx.get_builder().CreateRetVoid();
}

void compiler::create_llvm_bool(bool b) {
    auto new_function = ctx.create_custom_function(
            global_env->get_mangled_name(b ? "True" : "False"), 0);
    std::vector<instruction_ptr> instructions;
    instructions.push_back(instruction_ptr(new instruction_pushint(b)));
    instructions.push_back(instruction_ptr(new instruction_update(0)));
    ctx.get_builder().SetInsertPoint(&new_function->getEntryBlock());
    for(auto& instruction : instructions) {
        instruction->gen_llvm(ctx, new_function);
    }
    ctx.get_builder().CreateRetVoid();
}

void compiler::generate_llvm() {
    for(auto op : all_binops) {
        create_llvm_binop(op);
    }
    create_llvm_bool(true);
    create_llvm_bool(false);

    global_scp.generate_llvm(ctx);
}

void compiler::output_llvm(const std::string& into) {
    std::string targetTriple = llvm::sys::getDefaultTargetTriple();

    llvm::InitializeNativeTarget();
    llvm::InitializeNativeTargetAsmParser();
    llvm::InitializeNativeTargetAsmPrinter();

    std::string error;
    const llvm::Target* target =
        llvm::TargetRegistry::lookupTarget(targetTriple, error);
    if (!target) {
        std::cerr << error << std::endl;
    } else {
        std::string cpu = "generic";
        std::string features = "";
        llvm::TargetOptions options;
        std::unique_ptr<llvm::TargetMachine> targetMachine(
                target->createTargetMachine(targetTriple, cpu, features,
                    options, llvm::Optional<llvm::Reloc::Model>()));

        ctx.get_module().setDataLayout(targetMachine->createDataLayout());
        ctx.get_module().setTargetTriple(targetTriple);

        std::error_code ec;
        llvm::raw_fd_ostream file(into, ec, llvm::sys::fs::F_None);
        if (ec) {
            throw compiler_error("failed to open object file for writing");
        } else {
            llvm::CodeGenFileType type = llvm::CGFT_ObjectFile;
            llvm::legacy::PassManager pm;
            if (targetMachine->addPassesToEmitFile(pm, file, NULL, type)) {
                throw compiler_error("failed to add passes to pass manager");
            } else {
                pm.run(ctx.get_module());
                file.close();
            }
        }
    }
}

compiler::compiler(const std::string& filename)
    : file_m(), global_defs(), driver(file_m, global_defs, filename),
      global_env(new type_env), type_m(), mng(), global_scp(mng), ctx() {
    add_default_types();
    add_default_function_types();
}

void compiler::operator()(const std::string& into) {
    parse();
    typecheck();
    translate();
    compile();
    generate_llvm();
    output_llvm(into);
}

file_mgr& compiler::get_file_manager() {
    return file_m;
}

type_mgr& compiler::get_type_manager() {
    return type_m;
}