#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 "utils.hpp"

#include <optional>
#include <type_traits>
#include <vector>

namespace builders {

// statement+
std::vector<nodes::Statement>
build_source_file(parser::ParseTree::Node parser_node,
                  nodes::ExpressionStorage &expression_storage,
                  nodes::TypeStorage &type_storage,
                  names::NameTree &name_tree) {
  std::vector<nodes::Statement> statements;

  auto current_node = parser_node.nth_named_child(0);
  while (!current_node.is_null()) {
    statements.push_back(build_statement(current_node, expression_storage,
                                         type_storage, name_tree));
    current_node = current_node.next_named_sibling();
  }

  return statements;
}

// import | type_definition | function_definition | typeclass_definition
nodes::Statement build_statement(parser::ParseTree::Node parser_node,
                                 nodes::ExpressionStorage &expression_storage,
                                 nodes::TypeStorage &type_storage,
                                 names::NameTree &name_tree) {
  tokens::Type type = tokens::string_to_type(parser_node.get_type());

  std::optional<std::string> statement_name;
  std::optional<nodes::Statement> statement;

  switch (type) {
  case tokens::Type::IMPORT:
    statement = nodes::Statement(build_import(parser_node));
    statement_name = *statement.value()
                          .get<nodes::Import>()
                          .value()
                          ->get_module_name()
                          ->get();
    break;
  case tokens::Type::TYPE_DEFINITION:
    statement = nodes::Statement(
        build_type_definition(parser_node, expression_storage, type_storage));
    statement_name = *statement.value()
                          .get<nodes::TypeDefinition>()
                          .value()
                          ->get_name()
                          ->get();
    break;
  case tokens::Type::FUNCTION_DEFINITION:
    statement = nodes::Statement(build_function_definition(
        parser_node, expression_storage, type_storage));
    statement_name = *statement.value()
                          .get<nodes::FunctionDefinition>()
                          .value()
                          ->get_name()
                          ->get();
    break;
  case tokens::Type::TYPECLASS_DEFINITION:
    statement = nodes::Statement(build_typeclass_definition(
        parser_node, expression_storage, type_storage));
    statement_name = *statement.value()
                          .get<nodes::TypeclassDefinition>()
                          .value()
                          ->get_name()
                          ->get();
    break;
  case tokens::Type::EMPTY_LINES:
    statement = build_empty_lines(parser_node);
    break;
  default:
    error_handling::handle_parsing_error("Unexprected statement node type",
                                         parser_node);
  }

  if (!statement.has_value()) {
    error_handling::handle_general_error("Unreachable");
  }

  if (statement_name.has_value()) {
    // TODO: combine statements with same name
    if (!name_tree.insert(statement_name.value(), statement.value())) {
      error_handling::handle_parsing_error(
          "Two or more statements in file have the same name", parser_node);
    }
  }

  return std::move(statement.value());
}

// ('::' | '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<nodes::FunctionDefinition::Argument> arguments;
  std::vector<nodes::FunctionDefinition::Argument> argument_types;

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

  current_node = name_node.next_named_sibling();

  bool at_least_one_argument_annotation_found = false;

  size_t current_type_id = 0;

  std::optional<std::string> last_annotation;
  nodes::Modifier last_before_modifier = nodes::Modifier::NONE;
  nodes::Modifier last_after_modifier = nodes::Modifier::NONE;
  while (!current_node.is_null()) {
    // update last before modifier
    auto maybe_reference_node = current_node.previous_sibling();
    if (!maybe_reference_node.is_null() && !maybe_reference_node.is_named()) {
      last_before_modifier = build_modifier(maybe_reference_node);

      // only out, in, ref allowed
      if (last_before_modifier != nodes::Modifier::OUT &&
          last_before_modifier != nodes::Modifier::IN &&
          last_before_modifier != nodes::Modifier::REF) {
        last_before_modifier = nodes::Modifier::NONE;
      }
    }

    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:
      // update last after modifier
      maybe_reference_node = current_node.next_sibling();
      if (!maybe_reference_node.is_null() && !maybe_reference_node.is_named()) {
        last_after_modifier = build_modifier(maybe_reference_node);

        // only optional, result allowed
        if (last_after_modifier != nodes::Modifier::OPTIONAL &&
            last_after_modifier != nodes::Modifier::RESULT) {
          last_after_modifier = nodes::Modifier::NONE;
        }
      }

      // update conditions
      if (last_annotation.has_value()) {
        at_least_one_argument_annotation_found = true;
      }

      arguments.push_back(nodes::FunctionDefinition::Argument(
          last_annotation, build_identifier(current_node), last_before_modifier,
          last_after_modifier));

      last_annotation = std::nullopt;
      break;
    case tokens::Type::TYPE:
      if (current_type_id >= arguments.size()) {
        arguments.push_back(nodes::FunctionDefinition::Argument(
            last_annotation, build_type(current_node, type_storage),
            last_before_modifier));
      } else {
        if (!arguments[current_type_id].add_type(
                last_annotation, build_type(current_node, type_storage),
                last_before_modifier)) {
          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.",
              current_node);
        }
      }

      last_annotation = std::nullopt;

      ++current_type_id;
      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();
  }

  if (current_type_id > 0 && current_type_id < arguments.size()) {
    error_handling::handle_parsing_error(
        "Less types then arguments in function definition", parser_node);
  }

  // automatic annotations
  bool are_annotations_same_to_names =
      (!at_least_one_argument_annotation_found && current_type_id == 0);

  if (are_annotations_same_to_names) {
    for (size_t i = 0; i < arguments.size(); ++i) {
      std::string new_annotation = *arguments[i].get_name().value()->get();
      if (!arguments[i].add_annotation(
              new_annotation.substr(1, new_annotation.size() - 1))) {
        error_handling::handle_parsing_error(
            "no annotations provided ( => all annotations same to names), but "
            "can't add name annotation",
            current_node);
      }
    }
  }

  std::unordered_set<std::string> annotations_set;
  for (auto &argument : arguments) {
    if (argument.get_annotation().has_value()) {
      if (!annotations_set.insert(*argument.get_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(arguments), are_annotations_same_to_names,
      expression_node.has_value()
          ? build_expression(expression_node.value(), expression_storage,
                             type_storage)
          : std::optional<nodes::ExpressionProxy>());
}

// 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