#include "Compiler.hpp"
#include <llvm/IR/BasicBlock.h>
#include <llvm/IR/Verifier.h>
#include <llvm/IR/Type.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()
  {

  }

  /*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_module->setDataLayout(target_machine->createDataLayout());
    m_module->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_module);
    pass.run(*m_module);

    dest.flush();
  }

  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_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 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);
        auto body = node->child(3);

        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(ident, fun_type, node->loc());

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

        auto* bb = llvm::BasicBlock::Create(*m_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_module->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_context,
                                      llvm::APInt(32,
                                                  std::stoi(node->repr()),
                                                  true));
      } break;

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