#include "Compiler.hpp"
#include <llvm/Support/DynamicLibrary.h>
#include <llvm/IR/BasicBlock.h>
#include <llvm/IR/Verifier.h>
#include <llvm/IR/Type.h>
#include <llvm/IR/Function.h>
#include <llvm/IR/Value.h>
#include <llvm/TargetParser/Host.h>
#include <llvm/Target/TargetOptions.h>
#include <llvm/Target/TargetMachine.h>
#include <llvm/Support/TargetSelect.h>
#include <llvm/Support/FileSystem.h>
#include <llvm/MC/TargetRegistry.h>
#include <llvm/IR/LegacyPassManager.h>

namespace wg
{
  /*explicit*/ Compiler::Compiler(Mod& mod)
    : m_mod { mod }
  {

  }

  /*virtual*/ Compiler::~Compiler()
  {
  }

  void Compiler::gen(std::filesystem::path obj)
  {
    auto target_triple = llvm::sys::getDefaultTargetTriple();
    llvm::InitializeAllTargetInfos();
    llvm::InitializeAllTargets();
    llvm::InitializeAllTargetMCs();
    llvm::InitializeAllAsmParsers();
    llvm::InitializeAllAsmPrinters();

    std::string err;
    auto target = llvm::TargetRegistry::lookupTarget(target_triple, err);

    if (!target)
      {
        llvm::errs() << err;
        abort();
      }

    auto cpu = "generic";
    auto features = "";

    llvm::TargetOptions opt;
    auto rm = std::optional<llvm::Reloc::Model>();
    auto target_machine = target->createTargetMachine(target_triple,
                                                      cpu,
                                                      features,
                                                      opt,
                                                      rm);

    m_mod.mod().setDataLayout(target_machine->createDataLayout());
    m_mod.mod().setTargetTriple(target_triple);

    auto filename = obj.string();
    std::error_code ec;

    llvm::raw_fd_ostream dest(filename, ec, llvm::sys::fs::OF_None);
    WG_ASSERT(!ec, "cannot write output file.");

    llvm::legacy::PassManager pass;
    auto file_type = llvm::CodeGenFileType::CGFT_ObjectFile;

    if (target_machine->addPassesToEmitFile(pass, dest, nullptr, file_type))
      {
        llvm::errs() << "Target machine cannot emit a file of this type";
      }

    llvm::verifyModule(m_mod.mod());

    pass.run(m_mod.mod());

    dest.flush();
  }


  void Compiler::execute(std::shared_ptr<Node> node)
  {
    imports(node);
    scan(node);
    compile(node);
  }

  void Compiler::imports(std::shared_ptr<Node> node)
  {
    switch (node->type())
      {
      case NODE_DIR: {
        if (node->child(0)->repr() == "import")
          {
            auto mod_name = node->child(1)->repr();
            auto lib_path = "/usr/lib/libwongola-std.so";
            auto lib = llvm::sys::DynamicLibrary::getLibrary(lib_path);

            if (lib.isValid())
              {
                std::string entry_name = "lib_" + mod_name;
                typedef std::unique_ptr<Mod>(*func)();

                func entry = (func) lib.getAddressOfSymbol
                  (entry_name.c_str());

                if (entry == nullptr)
                  {
                    node->loc().error<compile_error>
                      ("cannot import module '"
                       + mod_name
                       + "'");
                  }

                if (node->size() == 3)
                  {
                    mod_name = node->child(2)->repr();
                  }

                m_imports.insert({mod_name, entry()});
              }
            else
              {
                node->loc().error<compile_error>("cannot stdlib");
              }
          }

      } break;
      default: {
        for (size_t i=0; i<node->size(); i++)
          {
            imports(node->child(i));
          }
      } break;
      }
  }

  void Compiler::scan(std::shared_ptr<Node> node)
  {
    switch (node->type())
      {
      case NODE_FUNDECL: {
        std::string name = node->child(0)->repr();
        std::vector<llvm::Type*> params;
        std::vector<std::shared_ptr<Node>> buffer;

        for (size_t i=0; i<node->child(1)->size(); i++)
          {
            auto param = node->child(1)->child(i);

            if (param->type() == NODE_IDENT)
              {
                buffer.push_back(param);
              }
            else if (param->type() == NODE_TYPE)
              {
                for (auto p: buffer)
                  {
                    params.push_back((llvm::Type*)
                                     llvm::Type::getInt32Ty(m_mod.context()));
                  }

                buffer.clear();
              }
          }

        auto* ret = llvm::Type::getVoidTy(m_mod.context());

        if (node->child(2)->size() > 0)
          {
            ret = (llvm::Type*) llvm::Type::getInt32Ty(m_mod.context());
          }

        auto* fun_ty = llvm::FunctionType::get(ret, params, false);

        m_mod.sym().declare_prototype(name, fun_ty, node->loc());

        llvm::Function::Create(fun_ty,
                               llvm::Function::ExternalLinkage,
                               name,
                               m_mod.mod());

      } break;

      default: {
        for (size_t i=0; i<node->size(); i++)
          {
            scan(node->child(i));
          }
      } break;
      }
  }

  llvm::Value* Compiler::compile(std::shared_ptr<Node> node)
  {
    switch (node->type())
      {
      case NODE_PROG: {
        for (size_t i=0; i<node->size(); i++)
          {
            compile(node->child(i));
          }

        return nullptr;
      } break;

      case NODE_BLOCK: {
        for (size_t i=0; i<node->size(); i++)
          {
            compile(node->child(i));
          }
        return nullptr;
      } break;

      case NODE_DIR: {
        return nullptr;
      } break;

      case NODE_NS: {
        auto mod_name = node->child(0)->repr();
        auto ident = node->child(1)->child(0)->repr();

        if (auto itr=m_imports.find(mod_name);
            itr != std::end(m_imports))
          {
            auto fun = itr->second->mod().getFunction(ident);
            WG_ASSERT(fun, "cannot call unknown function '" + ident + "'");

            std::vector<llvm::Value*> values;

            auto arg_node = node->child(1)->child(1);
            for (size_t i=0; i<arg_node->size(); i++)
              {
                auto arg = arg_node->child(i);
                auto val = compile(arg);
                values.push_back(val);
              }

            return m_builder->CreateCall(fun, values);
          }
        else
          {
            node->loc().error<compile_error>("cannot find module '"
                                             + mod_name
                                             + "'");
            abort();
          }

      } break;

      case NODE_RETURN: {
        return m_builder->CreateRet(compile(node->child(0)));
      } break;

        /*case NODE_EXTERN: {
        auto ident = node->child(0)->repr();
        auto params = node->child(1);
        auto ret = node->child(2);

        std::vector<std::string> names;
        std::vector<llvm::Type*> types;

        for (size_t i=0; i<params->size(); i++)
          {
            auto param = params->child(i);

            if (param->type() == NODE_IDENT)
              {
                names.push_back(param->repr());
              }
            else if (param->type() == NODE_TYPE)
              {
                auto ty = llvm::Type::getInt32Ty(*m_context);

                for (auto name: names)
                  {
                    m_sym->declare(name, ty, node->loc());
                    types.push_back(ty);
                  }

                names.clear();
              }
          }

        llvm::Type* ret_type = llvm::Type::getVoidTy(*m_context);

        if (ret->size() > 0)
          {
            ret_type = llvm::Type::getInt32Ty(*m_context);
          }

        auto fun_type = llvm::FunctionType::get(ret_type, types, false);
        auto fun = llvm::Function::Create(fun_type,
                                          llvm::Function::ExternalLinkage,
                                          ident,
                                          *m_module);

        m_sym->declare_prototype(ident, fun_type, node->loc());
        return fun;
        } break;*/

      case NODE_FUNDECL: {
        auto ident = node->child(0)->repr();
        auto body = node->child(3);
        auto fun = m_mod.mod(). getFunction(ident);


        llvm::BasicBlock* old_bb = m_builder->GetInsertBlock();

        auto* bb = llvm::BasicBlock::Create(m_mod.context(),
                                            "entry",
                                            fun);
        m_builder->SetInsertPoint(bb);

        compile(body);

        m_builder->SetInsertPoint(old_bb);

        llvm::verifyFunction(*fun);

        return fun;
      } break;

      case NODE_CALL: {
        std::string ident = node->child(0)->repr();

        auto fun = m_mod.mod().getFunction(ident);
        WG_ASSERT(fun, "cannot call unknown function '" + ident + "'");

        std::vector<llvm::Value*> values;

        for (size_t i=0; i<node->child(1)->size(); i++)
          {
            auto arg = node->child(1)->child(i);
            auto val = compile(arg);
            values.push_back(val);
          }

        return m_builder->CreateCall(fun, values);
      } break;

      case NODE_ADD: {
        auto* lhs = compile(node->child(0));
        auto* rhs = compile(node->child(1));
        return m_builder->CreateAdd(lhs, rhs);
      } break;

      case NODE_SUB: {
        auto* lhs = compile(node->child(0));
        auto* rhs = compile(node->child(1));
        return m_builder->CreateSub(lhs, rhs);
      } break;

      case NODE_MUL: {
        auto* lhs = compile(node->child(0));
        auto* rhs = compile(node->child(1));
        return m_builder->CreateMul(lhs, rhs);
      } break;

      case NODE_DIV: {
        auto* lhs = compile(node->child(0));
        auto* rhs = compile(node->child(1));
        return m_builder->CreateSDiv(lhs, rhs);
      } break;

      case NODE_MOD: {
        auto* lhs = compile(node->child(0));
        auto* rhs = compile(node->child(1));
        return m_builder->CreateSRem(lhs, rhs);
      } break;

      case NODE_INT: {
        return llvm::ConstantInt::get(m_mod.context(),
                                      llvm::APInt(32,
                                                  std::stoi(node->repr()),
                                                  true));
      } break;

      case NODE_BOOL: {
        return llvm::ConstantInt::getBool(m_mod.context(),
                                          node->repr() == "true");
      } break;

      case NODE_AND: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateAnd(lhs, rhs);
      } break;

      case NODE_OR: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateOr(lhs, rhs);
      } break;

      case NODE_NOT: {
        auto lhs = compile(node->child(0));
        return m_builder->CreateNot(lhs);
      } break;

      case NODE_EQ: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateICmpEQ(lhs, rhs);
      } break;

      case NODE_NE: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateICmpNE(lhs, rhs);
      } break;

      case NODE_LT: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateICmpSLT(lhs, rhs);
      } break;

      case NODE_LE: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateICmpSLE(lhs, rhs);
      } break;

      case NODE_GT: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateICmpSGT(lhs, rhs);
      } break;


      case NODE_GE: {
        auto lhs = compile(node->child(0));
        auto rhs = compile(node->child(1));
        return m_builder->CreateICmpSGE(lhs, rhs);
      } break;

      default:
        WG_ASSERT(false,
                  std::string()
                  + "cannot compile unknown node '"
                  + NodeTypeStr[node->type()]
                  + "'");
      }
  }
}