#include "ast.hpp"
#include <iostream>
#include "binop.hpp"
#include "definition.hpp"
#include "graph.hpp"
#include "instruction.hpp"
#include "llvm_context.hpp"
#include "parser.hpp"
#include "error.hpp"
#include "type.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 yy::parser::error(const yy::location& loc, const std::string& msg) {
    std::cout << "An error occured: " << msg << std::endl;
}

extern definition_group global_defs;

void typecheck_program(
        definition_group& defs,
        type_mgr& mgr, type_env_ptr& env) {
    type_ptr int_type = type_ptr(new type_base("Int")); 
    env->bind_type("Int", int_type);
    type_ptr int_type_app = type_ptr(new type_app(int_type));

    type_ptr binop_type = type_ptr(new type_arr(
                int_type_app,
                type_ptr(new type_arr(int_type_app, int_type_app))));
    env->bind("+", binop_type, visibility::global);
    env->bind("-", binop_type, visibility::global);
    env->bind("*", binop_type, visibility::global);
    env->bind("/", binop_type, visibility::global);

    std::set<std::string> free;
    defs.find_free(free);
    defs.typecheck(mgr, env);

    for(auto& pair : defs.env->names) {
        std::cout << pair.first << ": ";
        pair.second.type->print(mgr, std::cout);
        std::cout << std::endl;
    }
}

global_scope translate_program(definition_group& group) {
    global_scope scope;
    for(auto& data : group.defs_data) {
        data.second->into_globals(scope);
    }
    for(auto& defn : group.defs_defn) {
        auto& function = defn.second->into_global(scope);
        function.body->env->parent->set_mangled_name(defn.first, function.name);
    }
    return scope;
}

void gen_llvm_internal_op(llvm_context& ctx, binop op) {
    auto new_function = ctx.create_custom_function(op_action(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.builder.SetInsertPoint(&new_function->getEntryBlock());
    for(auto& instruction : instructions) {
        instruction->gen_llvm(ctx, new_function);
    }
    ctx.builder.CreateRetVoid();
}

void output_llvm(llvm_context& ctx, const std::string& filename) {
    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;
        llvm::TargetMachine* targetMachine =
            target->createTargetMachine(targetTriple, cpu, features,
                    options, llvm::Optional<llvm::Reloc::Model>());

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

        std::error_code ec;
        llvm::raw_fd_ostream file(filename, ec, llvm::sys::fs::F_None);
        if (ec) {
            throw 0;
        } else {
            llvm::CodeGenFileType type = llvm::CGFT_ObjectFile;
            llvm::legacy::PassManager pm;
            if (targetMachine->addPassesToEmitFile(pm, file, NULL, type)) {
                throw 0;
            } else {
                pm.run(ctx.module);
                file.close();
            }
        }
    }
}

void gen_llvm(global_scope& scope) {
    llvm_context ctx;
    gen_llvm_internal_op(ctx, PLUS);
    gen_llvm_internal_op(ctx, MINUS);
    gen_llvm_internal_op(ctx, TIMES);
    gen_llvm_internal_op(ctx, DIVIDE);

    scope.generate_llvm(ctx);

    ctx.module.print(llvm::outs(), nullptr);
    output_llvm(ctx, "program.o");
}

int main() {
    yy::parser parser;
    type_mgr mgr;
    type_env_ptr env(new type_env);

    parser.parse();
    for(auto& def_defn : global_defs.defs_defn) {
        std::cout << def_defn.second->name;
        for(auto& param : def_defn.second->params) std::cout << " " << param;
        std::cout << ":" << std::endl;

        def_defn.second->body->print(1, std::cout);
    }
    try {
        typecheck_program(global_defs, mgr, env);
        global_scope scope = translate_program(global_defs);
        scope.compile();
        gen_llvm(scope);
    } catch(unification_error& err) {
        std::cout << "failed to unify types: " << std::endl;
        std::cout << "  (1) \033[34m";
        err.left->print(mgr, std::cout);
        std::cout << "\033[0m" << std::endl;
        std::cout << "  (2) \033[32m";
        err.right->print(mgr, std::cout);
        std::cout << "\033[0m" << std::endl;
    } catch(type_error& err) {
        std::cout << "failed to type check program: " << err.description << std::endl;
    }
}