#include "statement_builders.hpp"

#include "basic_builders.hpp"
#include "basic_nodes.hpp"
#include "doc_builders.hpp"
#include "doc_nodes.hpp"
#include "error_handling.hpp"
#include "expression_builders.hpp"
#include "statement_nodes.hpp"
#include "tokens.hpp"
#include "tree_sitter_wrapper.hpp"
#include "type_builders.hpp"
#include "type_nodes.hpp"
#include <optional>
#include <type_traits>
#include <vector>

namespace builders {

// TODO: return some info ??
// statement+
void build_source_file(parser::ParseTree::Node parser_node,
                       nodes::ExpressionStorage &expression_storage,
                       nodes::TypeStorage &type_storage) {
  auto current_node = parser_node.nth_named_child(0);
  while (!current_node.is_null()) {
    build_statement(current_node, expression_storage, type_storage);
    // TODO: do something with statement
    current_node = current_node.next_named_sibling();
  }
}

// TODO: return some info ??
// import | type_definition | function_definition | typeclass_definition
void build_statement(parser::ParseTree::Node parser_node,
                     nodes::ExpressionStorage &expression_storage,
                     nodes::TypeStorage &type_storage) {
  tokens::Type type = tokens::string_to_type(parser_node.get_type());

  switch (type) {
  case tokens::Type::IMPORT:
    build_import(parser_node);
    return;
  case tokens::Type::TYPE_DEFINITION:
    build_type_definition(parser_node, expression_storage, type_storage);
    return;
  case tokens::Type::FUNCTION_DEFINITION:
    build_function_definition(parser_node, expression_storage, type_storage);
    return;
  case tokens::Type::TYPECLASS_DEFINITION:
    build_typeclass_definition(parser_node, expression_storage, type_storage);
    return;
  default:
    error_handling::handle_parsing_error("Unexprected statement node type",
                                         parser_node);
  }
  error_handling::handle_general_error("Unreachable");
  exit(1); // unreachable} // IN PROCESS
}

// ('::' | 'import') simple_name ('=' simple_name)? (':' identifier*)?
nodes::Import build_import(parser::ParseTree::Node parser_node) {

  auto name_node = parser_node.child_by_field_name("name");
  auto module_node = parser_node.child_by_field_name("module");
  if (module_node.is_null()) {
    module_node = name_node;
  }

  std::vector<nodes::Identifier> symbols;

  auto current_node = module_node.next_named_sibling();
  while (!current_node.is_null()) {
    symbols.push_back(build_identifier(current_node));
    current_node = current_node.next_named_sibling();
  }

  return nodes::Import(build_node(parser_node), build_identifier(name_node),
                       build_identifier(module_node), std::move(symbols));
}

// '?' expression
nodes::Constraint build_constraint(parser::ParseTree::Node parser_node,
                                   nodes::ExpressionStorage &expression_storage,
                                   nodes::TypeStorage &type_storage) {
  return nodes::Constraint(build_node(parser_node),
                           build_expression(parser_node.nth_named_child(0),
                                            expression_storage, type_storage));
}

parser::ParseTree::Node collect_symbol_doc_nodes(
    parser::ParseTree::Node start_parser_node,
    std::optional<parser::ParseTree::Node> &description_node,
    std::vector<parser::ParseTree::Node> &annotation_nodes) {
  auto current_node = start_parser_node;

  description_node = std::nullopt;
  annotation_nodes.clear();

  if (current_node.is_null()) {
    return current_node;
  }

  if (tokens::string_to_type(current_node.get_type()) ==
      tokens::Type::DEFINITION_INFO) {
    description_node = current_node;
    current_node = current_node.next_named_sibling();
  }

  while (!current_node.is_null() &&
         tokens::string_to_type(current_node.get_type()) ==
             tokens::Type::ANNOTATION_INFO) {
    annotation_nodes.push_back(current_node);
    current_node = current_node.next_named_sibling();
  }

  return current_node;
}

// definition_info? annotation_info* '^'? simple_type (argument_type* '='
// variant_type)? ('{' function_definition* '}' | ';')
nodes::TypeDefinition
build_type_definition(parser::ParseTree::Node parser_node,
                      nodes::ExpressionStorage &expression_storage,
                      nodes::TypeStorage &type_storage) {
  bool is_on_heap = parser_node.nth_child(0).get_value() == "^";

  std::optional<parser::ParseTree::Node> description_node;
  std::vector<parser::ParseTree::Node> annotation_nodes;

  auto name_node = collect_symbol_doc_nodes(parser_node.nth_named_child(0),
                                            description_node, annotation_nodes);

  std::vector<nodes::Identifier> arguments;
  std::vector<nodes::FunctionDefinition> methods;

  std::optional<parser::ParseTree::Node> type_node;

  auto current_node = name_node.next_named_sibling();
  while (!current_node.is_null()) {
    switch (tokens::string_to_type(current_node.get_type())) {
    case tokens::Type::ARGUMENT_TYPE_IDENTIFIER:
      arguments.push_back(build_identifier(current_node));
      break;
    case tokens::Type::VARIANT_TYPE:
      if (type_node.has_value()) {
        error_handling::handle_parsing_error(
            "More then one type node in type definition", parser_node);
      }
      type_node = current_node;
      break;
    case tokens::Type::FUNCTION_DEFINITION:
      methods.push_back(build_function_definition(
          current_node, expression_storage, type_storage));
      break;
    default:
      error_handling::handle_parsing_error(
          "Unexprected node type in type definition", parser_node);
    }
    current_node = current_node.next_named_sibling();
  }

  if (!type_node.has_value() && !annotation_nodes.empty()) {
    error_handling::handle_parsing_error("", parser_node);
  }

  std::optional<nodes::VariantType> type =
      type_node.has_value()
          ? build_variant_type(type_node.value(), type_storage)
          : std::optional<nodes::VariantType>();

  std::unordered_set<std::string> annotations;

  // collect annotations from type
  if (type.has_value()) {
    for (size_t i = 0; i < type.value().size(); ++i) {
      auto constructor_annotations = type.value().get(i)->get_all_annotations();
      for (auto &annotation : constructor_annotations) {
        annotations.insert(annotation);
      }
    }
  }

  return nodes::TypeDefinition(
      build_node(parser_node),
      build_symbol_docs(description_node, annotation_nodes, annotations),
      is_on_heap, build_identifier(name_node), std::move(arguments),
      std::move(type), std::move(methods));
}

// definition_info? annotation_info*  (constraint ';')*  _var_let_? (simple_name
// | '(' operator ')') (annotation? _reference_? argument_name '?'?)* (:
// (annotation? _reference_ type)+)?
// ('=' (block | expression ';') | ';')
nodes::FunctionDefinition
build_function_definition(parser::ParseTree::Node parser_node,
                          nodes::ExpressionStorage &expression_storage,
                          nodes::TypeStorage &type_storage) {

  std::optional<parser::ParseTree::Node> description_node;
  std::vector<parser::ParseTree::Node> annotation_nodes;

  std::vector<nodes::Constraint> constraints;

  auto current_node = collect_symbol_doc_nodes(
      parser_node.nth_named_child(0), description_node, annotation_nodes);

  while (!current_node.is_null() &&
         tokens::string_to_type(current_node.get_type()) ==
             tokens::Type::CONSTRAINT) {
    constraints.push_back(
        build_constraint(current_node, expression_storage, type_storage));
    current_node = current_node.next_named_sibling();
  }

  auto name_node = current_node;

  nodes::FunctionDefinition::ModifierType modifier =
      nodes::FunctionDefinition::STATIC;

  current_node = name_node.previous_sibling();
  if (!current_node.is_null() && current_node.is_named()) {
    std::string modifier_str = current_node.get_value();
    if (modifier_str == "%" || modifier_str == "let") {
      modifier = nodes::FunctionDefinition::LET;
    } else if (modifier_str == "$" || modifier_str == "var") {
      modifier = nodes::FunctionDefinition::VAR;
    }
  }

  std::vector<std::optional<std::string>> argument_annotations;
  std::vector<nodes::Modifier> argument_reference_types;
  std::vector<nodes::Identifier> arguments;
  std::vector<bool> optional_arguments;
  std::vector<bool> result_arguments;

  std::vector<std::optional<std::string>> type_annotations;
  std::vector<nodes::Modifier> type_reference_types;
  std::vector<nodes::TypeProxy> types;

  std::optional<parser::ParseTree::Node> expression_node;

  current_node = name_node.next_named_sibling();

  bool at_least_one_argument_annotation_found = false;
  bool at_least_one_argument_modifier_found = false;

  std::optional<std::string> last_annotation;
  nodes::Modifier last_reference_type = nodes::Modifier::NONE;
  while (!current_node.is_null()) {
    auto maybe_reference_node = current_node.previous_sibling();
    if (!maybe_reference_node.is_null() && !maybe_reference_node.is_named()) {
      last_reference_type = build_modifier(maybe_reference_node);
    }

    switch (tokens::string_to_type(current_node.get_type())) {
    case tokens::Type::ANNOTATION_IDENTIFIER:
      last_annotation = build_annotation(current_node);
      break;
    case tokens::Type::ARGUMENT_NAME_IDENTIFIER:
      if (last_annotation.has_value()) {
        at_least_one_argument_annotation_found = true;
        at_least_one_argument_modifier_found = true;
      }
      if (last_reference_type != nodes::Modifier::NONE) {
        at_least_one_argument_modifier_found = true;
      }

      argument_annotations.push_back(last_annotation);
      argument_reference_types.push_back(last_reference_type);
      arguments.push_back(build_identifier(current_node));
      optional_arguments.push_back(!current_node.next_sibling().is_null() &&
                                   current_node.next_sibling().is_named() &&
                                   current_node.next_sibling().get_value() ==
                                       "?");
      result_arguments.push_back(!current_node.next_sibling().is_null() &&
                                 current_node.next_sibling().is_named() &&
                                 current_node.next_sibling().get_value() ==
                                     "!");
      if (optional_arguments.back() || result_arguments.back()) {
        at_least_one_argument_modifier_found = true;
      }

      last_reference_type = nodes::Modifier::NONE;
      last_annotation = std::nullopt;
      break;
    case tokens::Type::TYPE:
      type_annotations.push_back(last_annotation);
      type_reference_types.push_back(last_reference_type);
      types.push_back(build_type(current_node, type_storage));
      last_reference_type = nodes::Modifier::NONE;
      last_annotation = std::nullopt;
      break;
    default:
      if (expression_node.has_value()) {
        error_handling::handle_parsing_error(
            "More then one expression found in function definition",
            parser_node);
      }
      expression_node = current_node;
      break;
    }
    current_node = current_node.next_named_sibling();
  }

  std::vector<std::optional<std::string>> *annotations = nullptr;
  std::vector<nodes::Modifier> *reference_types = nullptr;

  if (!at_least_one_argument_annotation_found) {
    for (size_t i = 0; i < argument_annotations.size(); ++i) {
      std::string argument_annotation = *arguments[i].get();
      argument_annotations[i] =
          argument_annotation.substr(1, argument_annotation.size() - 1);
    }
  }

  if (types.empty()) {
    annotations = &argument_annotations;
    reference_types = &argument_reference_types;
  } else if (at_least_one_argument_modifier_found) {
    error_handling::handle_parsing_error(
        "It is impossible to use argument modifiers (annotations, references, "
        "optional markers, result markers) when types explicitely "
        "defined. Use type annotations instead.",
        parser_node);
  } else {
    annotations = &type_annotations;
    reference_types = &type_reference_types;
  }

  std::unordered_set<std::string> annotations_set;
  for (auto &annotation : *annotations) {
    if (annotation.has_value()) {
      if (!annotations_set.insert(annotation.value()).second) {
        error_handling::handle_parsing_error(
            "Two or more same annotations found in function definition",
            parser_node);
      }
    }
  }

  return nodes::FunctionDefinition(
      build_node(parser_node),
      build_symbol_docs(description_node, annotation_nodes, annotations_set),
      std::move(constraints), modifier, build_identifier(name_node),
      std::move(*annotations), std::move(arguments),
      std::move(*reference_types), std::move(types),
      std::move(optional_arguments), std::move(result_arguments),
      expression_node.has_value()
          ? build_expression(expression_node.value(), expression_storage,
                             type_storage)
          : std::optional<nodes::ExpressionProxy>());
} // TODO: refactor ??

// definition_info? annotation_info* typeclass_identifier (':'
// typeclass_identifier+)? ('{' function_definition* '}' | ';')
nodes::TypeclassDefinition
build_typeclass_definition(parser::ParseTree::Node parser_node,
                           nodes::ExpressionStorage &expression_storage,
                           nodes::TypeStorage &type_storage) {

  std::optional<parser::ParseTree::Node> description_node;
  std::vector<parser::ParseTree::Node> annotation_nodes;

  auto name_node = collect_symbol_doc_nodes(parser_node.nth_named_child(0),
                                            description_node, annotation_nodes);

  if (!annotation_nodes.empty()) {
    error_handling::handle_parsing_error(
        "Typeclass can't have annotation info nodes", parser_node);
  }

  std::vector<nodes::Identifier> base_typeclasses;
  std::vector<nodes::FunctionDefinition> methods;

  auto current_node = name_node.next_named_sibling();
  while (!current_node.is_null()) {
    switch (tokens::string_to_type(current_node.get_type())) {
    case tokens::Type::TYPECLASS_IDENTIFIER:
      base_typeclasses.push_back(build_identifier(current_node));
      break;
    case tokens::Type::FUNCTION_DEFINITION:
      methods.push_back(build_function_definition(
          current_node, expression_storage, type_storage));
      break;
    default:
      error_handling::handle_parsing_error(
          "Unexprected node type in type definition", parser_node);
    }
    current_node = current_node.next_named_sibling();
  }

  return nodes::TypeclassDefinition(
      build_node(parser_node), build_symbol_docs(description_node),
      build_identifier(name_node), std::move(base_typeclasses),
      std::move(methods));
}

} // namespace builders