#include <language/modules/ModuleRepository.hpp>

#include <language/ast/ASTNode.hpp>
#include <language/modules/CoreModule.hpp>
#include <language/modules/LinearSolverModule.hpp>
#include <language/modules/MathModule.hpp>
#include <language/modules/MeshModule.hpp>
#include <language/modules/SchemeModule.hpp>
#include <language/modules/SocketModule.hpp>
#include <language/modules/UtilsModule.hpp>
#include <language/modules/WriterModule.hpp>
#include <language/utils/BasicAffectationRegistrerFor.hpp>
#include <language/utils/BuiltinFunctionEmbedder.hpp>
#include <language/utils/ParseError.hpp>
#include <language/utils/SymbolTable.hpp>
#include <language/utils/TypeDescriptor.hpp>
#include <language/utils/ValueDescriptor.hpp>

#include <utils/PugsAssert.hpp>

#include <algorithm>

void
ModuleRepository::_subscribe(std::unique_ptr<IModule> m)
{
  auto is_keyword = [](const std::string& s) -> bool {
    if (s.size() == 0) {
      return false;
    } else {
      if (not(std::isalpha(s[0]) or s[0] == '_')) {
        return false;
      }
      for (size_t i = 1; i < s.size(); ++i) {
        if (not(std::isalnum(s[0]) or s[0] == '_')) {
          return false;
        }
      }
    }

    return true;
  };

  if (not is_keyword(std::string{m->name()})) {
    std::ostringstream os;
    os << "cannot subscribe module with invalid name: '" << m->name() << "'\n";
    throw UnexpectedError(os.str());
  }

  auto [i_module, success] = m_module_set.emplace(m->name(), std::move(m));
  Assert(success, "module has already been subscribed");
}

ModuleRepository::ModuleRepository()
{
  this->_subscribe(std::make_unique<CoreModule>());
  this->_subscribe(std::make_unique<LinearSolverModule>());
  this->_subscribe(std::make_unique<MathModule>());
  this->_subscribe(std::make_unique<MeshModule>());
  this->_subscribe(std::make_unique<SchemeModule>());
  this->_subscribe(std::make_unique<SocketModule>());
  this->_subscribe(std::make_unique<UtilsModule>());
  this->_subscribe(std::make_unique<WriterModule>());
}

template <typename NameEmbedderMapT, typename EmbedderTableT>
void
ModuleRepository::_populateEmbedderTableT(const ASTNode& module_node,
                                          const std::string& module_name,
                                          const NameEmbedderMapT& name_embedder_map,
                                          const ASTNodeDataType& data_type,
                                          SymbolTable& symbol_table,
                                          EmbedderTableT& embedder_table)
{
  for (auto [symbol_name, embedded] : name_embedder_map) {
    auto [i_symbol, success] = symbol_table.add(symbol_name, module_node.begin());

    if (not success) {
      std::ostringstream error_message;
      error_message << "importing module '" << module_name << "', cannot add symbol '" << symbol_name
                    << "', it is already defined!";
      throw ParseError(error_message.str(), module_node.begin());
    }

    i_symbol->attributes().setDataType(data_type);
    i_symbol->attributes().setIsInitialized();
    i_symbol->attributes().value() = embedder_table.size();

    embedder_table.add(embedded);
  }
}

void
ModuleRepository::_populateSymbolTable(const ASTNode& module_node,
                                       const std::string& module_name,
                                       const IModule::NameValueMap& name_value_descriptor_map,
                                       SymbolTable& symbol_table)
{
  for (auto [symbol_name, value_descriptor] : name_value_descriptor_map) {
    auto [i_symbol, success] = symbol_table.add(symbol_name, module_node.begin());

    if (not success) {
      std::ostringstream error_message;
      error_message << "importing module '" << module_name << "', cannot add symbol '" << symbol_name
                    << "', it is already defined!";
      throw ParseError(error_message.str(), module_node.begin());
    }

    i_symbol->attributes().setDataType(value_descriptor->type());
    i_symbol->attributes().setIsInitialized();
    i_symbol->attributes().value() = value_descriptor->value();
  }
}

void
ModuleRepository::populateSymbolTable(const ASTNode& module_name_node, SymbolTable& symbol_table)
{
  const std::string& module_name = module_name_node.string();

  auto i_module = m_module_set.find(module_name);
  if (i_module != m_module_set.end()) {
    const IModule& populating_module = *i_module->second;

    if (populating_module.isMandatory()) {
      std::ostringstream error_message;
      error_message << "module '" << rang::fgB::blue << module_name << rang::style::reset << rang::style::bold
                    << "' is an autoload " << rang::fgB::yellow << "mandatory" << rang::style::reset
                    << rang::style::bold << " module. It cannot be imported explicitly!";
      throw ParseError(error_message.str(), module_name_node.begin());
    }

    this->_populateEmbedderTableT(module_name_node, module_name, populating_module.getNameBuiltinFunctionMap(),
                                  ASTNodeDataType::build<ASTNodeDataType::builtin_function_t>(), symbol_table,
                                  symbol_table.builtinFunctionEmbedderTable());

    this->_populateEmbedderTableT(module_name_node, module_name, populating_module.getNameTypeMap(),
                                  ASTNodeDataType::build<ASTNodeDataType::type_name_id_t>(), symbol_table,
                                  symbol_table.typeEmbedderTable());

    this->_populateSymbolTable(module_name_node, module_name, populating_module.getNameValueMap(), symbol_table);

    for (auto [symbol_name, embedded] : populating_module.getNameTypeMap()) {
      BasicAffectationRegisterFor<EmbeddedData>(ASTNodeDataType::build<ASTNodeDataType::type_id_t>(symbol_name));
    }

  } else {
    throw ParseError(std::string{"could not find module "} + module_name, std::vector{module_name_node.begin()});
  }
}

void
ModuleRepository::populateMandatorySymbolTable(const ASTNode& root_node, SymbolTable& symbol_table)
{
  for (auto&& [module_name, i_module] : m_module_set) {
    if (i_module->isMandatory()) {
      this->_populateEmbedderTableT(root_node, module_name, i_module->getNameBuiltinFunctionMap(),
                                    ASTNodeDataType::build<ASTNodeDataType::builtin_function_t>(), symbol_table,
                                    symbol_table.builtinFunctionEmbedderTable());

      this->_populateEmbedderTableT(root_node, module_name, i_module->getNameTypeMap(),
                                    ASTNodeDataType::build<ASTNodeDataType::type_name_id_t>(), symbol_table,
                                    symbol_table.typeEmbedderTable());

      this->_populateSymbolTable(root_node, module_name, i_module->getNameValueMap(), symbol_table);

      for (auto [symbol_name, embedded] : i_module->getNameTypeMap()) {
        BasicAffectationRegisterFor<EmbeddedData>(ASTNodeDataType::build<ASTNodeDataType::type_id_t>(symbol_name));
      }

      i_module->registerOperators();
    }
  }
}

std::string
ModuleRepository::getAvailableModules() const
{
  std::stringstream os;
  os << rang::fgB::yellow << "Available modules" << rang::fg::blue << " [modules tagged with a " << rang::style::reset
     << rang::style::bold << '*' << rang::style::reset << rang::fg::blue << " are automatically imported]"
     << rang::style::reset << '\n';
  for (auto& [name, i_module] : m_module_set) {
    if (i_module->isMandatory()) {
      os << rang::style::bold << " *" << rang::style::reset;
    } else {
      os << "  ";
    }
    os << rang::fgB::green << name << rang::style::reset << '\n';
  }

  return os.str();
}

void
ModuleRepository::registerOperators(const std::string& module_name)
{
  auto i_module = m_module_set.find(module_name);
  if (i_module != m_module_set.end()) {
    i_module->second->registerOperators();
  } else {
    throw NormalError(std::string{"could not find module "} + module_name);
  }
}

std::string
ModuleRepository::getModuleInfo(const std::string& module_name) const
{
  auto demangleBuiltinFunction = [](const std::string& mangled_name) -> std::string {
    size_t i = 0;
    for (; i < mangled_name.size(); ++i) {
      if (mangled_name[i] == ':')
        break;
    }
    return mangled_name.substr(0, i);
  };

  std::stringstream os;
  auto i_module = m_module_set.find(module_name);
  if (i_module != m_module_set.end()) {
    os << rang::fgB::yellow << "Module '" << rang::fgB::blue << module_name << rang::fgB::yellow << "' provides"
       << rang::style::reset << '\n';
    const auto& builtin_function_map = i_module->second->getNameBuiltinFunctionMap();
    if (builtin_function_map.size() > 0) {
      os << "  functions\n";
      for (auto& [mangled_name, function] : builtin_function_map) {
        os << "    " << rang::fgB::green << demangleBuiltinFunction(mangled_name) << rang::style::reset << ": ";
        os << dataTypeName(function->getParameterDataTypes());
        os << rang::fgB::yellow << " -> " << rang::style::reset;
        os << dataTypeName(function->getReturnDataType()) << '\n';
      }
    }

    const auto& builtin_type_map = i_module->second->getNameTypeMap();
    if (builtin_type_map.size() > 0) {
      os << "  types\n";
      for (auto& [name, descriptor] : builtin_type_map) {
        os << "    " << rang::fgB::green << name << rang::style::reset << '\n';
      }
    }

  } else {
    throw NormalError(std::string{"could not find module "} + module_name);
  }

  return os.str();
}