#pragma once

#include "basic_nodes.hpp"
#include "doc_nodes.hpp"
#include "expression_nodes.hpp"
#include "type_nodes.hpp"

#include <vector>

// IN PROGRESS

namespace nodes {

// IN PROGRESS: add another constructors ??
class Import : public Node {
public:
  Import(Node node, const Identifier &import_name,
         const Identifier &module_name, std::vector<Identifier> &&symbols = {})
      : Node(node), import_name_(import_name), module_name_(module_name),
        symbols_(std::move(symbols)) {}

  //

  Identifier *get_import_name() { return &import_name_; }

  const Identifier *get_import_name() const { return &import_name_; }

  //

  Identifier *get_module_name() { return &module_name_; }

  const Identifier *get_module_name() const { return &module_name_; }

  //

  size_t symbols_size() const { return symbols_.size(); }

  Identifier *get_symbol(size_t id) { return &symbols_.at(id); }

  const Identifier *get_symbol(size_t id) const { return &symbols_.at(id); }

private:
  Identifier import_name_;
  Identifier module_name_;
  std::vector<Identifier> symbols_;
};

class Constraint : public Node {
public:
  Constraint(Node node, ExpressionProxy expression)
      : Node(node), expression_(expression) {}

  Expression *get_expression() { return expression_.get(); }

  const Expression *get_expression() const { return expression_.get(); }

private:
  ExpressionProxy expression_;
};

class FunctionDefinition : public Node {
public:
  enum ModifierType {
    STATIC,
    LET,
    VAR,
  };

  FunctionDefinition(Node node, SymbolDocs &&docs,
                     std::vector<Constraint> &&constraints,
                     ModifierType modifier, const Identifier &name,
                     std::vector<std::optional<std::string>> &&annotations,
                     std::vector<Identifier> &&arguments,
                     std::vector<Modifier> &&reference_types,
                     std::vector<TypeProxy> &&types,
                     std::vector<bool> &&optional_arguments,
                     std::vector<bool> &&result_arguments,
                     bool is_annotations_same_to_names,
                     std::optional<ExpressionProxy> expression)
      : Node(node), docs_(std::move(docs)),
        constraints_(std::move(constraints)), modifier_(modifier), name_(name),
        annotations_(std::move(annotations)), arguments_(std::move(arguments)),
        reference_types_(std::move(reference_types)), types_(std::move(types)),
        optional_arguments_(optional_arguments),
        result_arguments_(result_arguments),
        is_annotations_same_to_names_(is_annotations_same_to_names),
        expression_(expression) {}

  //

  SymbolDocs *get_docs() { return &docs_; }

  const SymbolDocs *get_docs() const { return &docs_; }

  //

  size_t get_constraints_size() const { return constraints_.size(); }

  Constraint *get_constraint(size_t id) { return &constraints_.at(id); }

  const Constraint *get_constraint(size_t id) const {
    return &constraints_.at(id);
  }

  //

  ModifierType get_modifier() const { return modifier_; }

  //

  Identifier *get_name() { return &name_; }

  const Identifier *get_name() const { return &name_; }

  //

  std::optional<std::string *> get_argument_annotation(size_t id) {
    if (annotations_.at(id).has_value()) {
      return &annotations_[id].value();
    }
    return std::nullopt;
  }

  std::optional<const std::string *> get_argument_annotation(size_t id) const {
    if (annotations_.at(id).has_value()) {
      return &annotations_[id].value();
    }
    return std::nullopt;
  }

  //

  size_t get_arguments_size() const { return arguments_.size(); }

  Identifier *get_argument(size_t id) { return &arguments_.at(id); }

  const Identifier *get_argument(size_t id) const { return &arguments_.at(id); }

  //

  Modifier get_argument_reference_type(size_t id) const {
    return reference_types_.at(id);
  }

  //

  size_t get_argument_types_size() const { return types_.size(); }

  Type *get_argument_type(size_t id) { return types_.at(id).get(); }

  const Type *get_argument_type(size_t id) const { return types_.at(id).get(); }

  //

  bool is_argument_optional(size_t id) const {
    return optional_arguments_.at(id);
  }

  //

  bool is_argument_result(size_t id) const { return result_arguments_.at(id); }

  //

  bool is_annotations_same_to_names() const {
    return is_annotations_same_to_names_;
  }

  //

  std::optional<Expression *> get_expression() {
    if (expression_.has_value()) {
      return expression_.value().get();
    }
    return std::nullopt;
  }

  std::optional<const Expression *> get_expression() const {
    if (expression_.has_value()) {
      return expression_.value().get();
    }
    return std::nullopt;
  }

private:
  SymbolDocs docs_;
  std::vector<Constraint> constraints_;
  ModifierType modifier_;
  Identifier name_;
  std::vector<std::optional<std::string>> annotations_;
  std::vector<Identifier> arguments_;
  std::vector<Modifier> reference_types_;
  std::vector<TypeProxy> types_;
  std::vector<bool> optional_arguments_;
  std::vector<bool> result_arguments_;
  bool is_annotations_same_to_names_;
  std::optional<ExpressionProxy> expression_;
}; // refactor ??

class TypeDefinition : public Node {
public:
  TypeDefinition(Node node, SymbolDocs &&docs, bool is_on_heap,
                 const Identifier &name, std::vector<Identifier> &&arguments,
                 std::optional<VariantType> &&type,
                 std::vector<FunctionDefinition> &&methods)
      : Node(node), docs_(std::move(docs)), is_on_heap_(is_on_heap),
        name_(name), arguments_(std::move(arguments)), type_(std::move(type)),
        methods_(std::move(methods)) {}

  //

  SymbolDocs *get_docs() { return &docs_; }

  const SymbolDocs *get_docs() const { return &docs_; }

  //

  bool is_on_heap() const { return is_on_heap_; }

  //

  Identifier *get_name() { return &name_; }

  const Identifier *get_name() const { return &name_; }

  //

  size_t get_arguments_size() const { return arguments_.size(); }

  Identifier *get_argument(size_t id) { return &arguments_.at(id); }

  const Identifier *get_argument(size_t id) const { return &arguments_.at(id); }

  //

  std::optional<VariantType *> get_type() {
    if (type_.has_value()) {
      return &type_.value();
    }
    return std::nullopt;
  }

  std::optional<const VariantType *> get_type() const {
    if (type_.has_value()) {
      return &type_.value();
    }
    return std::nullopt;
  }

  //

  size_t get_methods_size() const { return methods_.size(); }

  FunctionDefinition *get_method(size_t id) { return &methods_.at(id); }

  const FunctionDefinition *get_method(size_t id) const {
    return &methods_.at(id);
  }

private:
  SymbolDocs docs_;
  bool is_on_heap_;
  Identifier name_;
  std::vector<Identifier> arguments_;
  std::optional<VariantType> type_; // TupleType is VariantType with one variant
  std::vector<FunctionDefinition> methods_;
};

class TypeclassDefinition : public Node {
public:
  TypeclassDefinition(Node node, SymbolDocs &&docs, const Identifier &name,
                      std::vector<Identifier> &&base_typeclasses,
                      std::vector<FunctionDefinition> &&methods)
      : Node(node), docs_(std::move(docs)), name_(name),
        base_typeclasses_(std::move(base_typeclasses)),
        methods_(std::move(methods)) {}

  //

  SymbolDocs *get_docs() { return &docs_; }

  const SymbolDocs *get_docs() const { return &docs_; }

  //

  Identifier *get_name() { return &name_; }

  const Identifier *get_name() const { return &name_; }

  //

  size_t get_base_typeclasses_size() const { return base_typeclasses_.size(); }

  Identifier *get_base_typeclass(size_t id) {
    return &base_typeclasses_.at(id);
  }

  const Identifier *get_base_typeclass(size_t id) const {
    return &base_typeclasses_.at(id);
  }

  //

  size_t get_methods_size() const { return methods_.size(); }

  FunctionDefinition *get_method(size_t id) { return &methods_.at(id); }

  const FunctionDefinition *get_method(size_t id) const {
    return &methods_.at(id);
  }

private:
  SymbolDocs docs_;
  Identifier name_;
  std::vector<Identifier> base_typeclasses_;
  std::vector<FunctionDefinition> methods_;
};

} // namespace nodes