#pragma once

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

#include <vector>

// IN PROGRESS

namespace nodes {

enum class CombineResult {
  DIFFERENT_NAME_ERROR,
  DIFFERNENT_MODIFIER_ERROR,
  MORE_THEN_ONE_DOCS_ERROR,
  MORE_THEN_ONE_CONSTRAINTS_ERROR,
  MORE_THEN_ONE_DEFINITION_BODY_ERROR,
  ARGUMENTS_ERROR,
  DIFFERENT_STATEMENT_TYPES,
  STATEMENTS_CANT_BE_COMBINED_ERROR,
  GENERIC_ERROR,
  OK,
};

// 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:
  class Argument {
  public:
    Argument(const std::optional<std::string> &annotation, Identifier &&name,
             Modifier before_modifier = Modifier::NONE,
             Modifier after_modifier = Modifier::NONE)
        : annotation_(annotation), name_(std::move(name)),
          before_modifier_(before_modifier), after_modifier_(after_modifier) {}

    Argument(const std::optional<std::string> &annotation,
             const Identifier &name, Modifier before_modifier = Modifier::NONE,
             Modifier after_modifier = Modifier::NONE)
        : annotation_(annotation), name_(name),
          before_modifier_(before_modifier), after_modifier_(after_modifier) {}

    Argument(const std::optional<std::string> &annotation, TypeProxy type,
             Modifier before_modifier = Modifier::NONE)
        : annotation_(annotation), type_(type),
          before_modifier_(before_modifier),
          after_modifier_(type.get()->get_modifier()) {}

    //

    // add type with same annotation and same before_modifier
    bool update_type_from(const Argument &other_argument) {
      if (type_.has_value() || !other_argument.type_.has_value()) {
        return false;
      }

      if (annotation_.has_value() &&
          (!other_argument.annotation_.has_value() ||
           annotation_.value() != other_argument.annotation_.value())) {
        return false;
      }

      if (before_modifier_ != other_argument.before_modifier_) {
        return false;
      }

      if (after_modifier_ != other_argument.after_modifier_) {
        return false;
      }

      annotation_ = other_argument.annotation_;
      type_ = other_argument.type_;
      before_modifier_ = other_argument.before_modifier_;
      after_modifier_ = other_argument.after_modifier_;

      return true;
    }

    // check, that argument modifiers are none or same to different from type
    // modifiers
    bool add_type(const std::optional<std::string> &annotation, TypeProxy type,
                  nodes::Modifier before_modifier = Modifier::NONE) {
      if (type_.has_value()) {
        return false;
      }

      if (annotation_.has_value() &&
          (!annotation.has_value() ||
           annotation_.value() != annotation.value())) {
        return false;
      }

      if (before_modifier_ != Modifier::NONE &&
          before_modifier_ != before_modifier) {
        return false;
      }

      if (after_modifier_ != Modifier::NONE &&
          after_modifier_ != type.get()->get_modifier()) {
        return false;
      }

      annotation_ = annotation;
      type_ = type;
      before_modifier_ = before_modifier;
      after_modifier_ = type.get()->get_modifier();

      return true;
    }

    bool add_annotation(const std::string &annotation) {
      if (annotation_.has_value()) {
        return false;
      }

      annotation_ = annotation;

      return true;
    }

    //

    std::optional<std::string *> get_annotation() {
      if (annotation_.has_value()) {
        return &annotation_.value();
      }
      return std::nullopt;
    }

    std::optional<const std::string *> get_annotation() const {
      if (annotation_.has_value()) {
        return &annotation_.value();
      }
      return std::nullopt;
    }

    //

    std::optional<Identifier *> get_name() {
      if (name_.has_value()) {
        return &name_.value();
      }
      return std::nullopt;
    }

    std::optional<const Identifier *> get_name() const {
      if (name_.has_value()) {
        return &name_.value();
      }
      return std::nullopt;
    }

    //

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

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

    //

    Modifier get_before_modifier() const { return before_modifier_; }

    Modifier get_after_modifier() const { return after_modifier_; }

  private:
    std::optional<std::string> annotation_;
    std::optional<Identifier> name_; // no name for output arguments
    std::optional<TypeProxy> type_;  // no type if it is deduced
    Modifier before_modifier_ =
        Modifier::NONE; // in, out, ref, none // sync with type
    Modifier after_modifier_ =
        Modifier::NONE; // optional, result, none // sync with type
  };

  FunctionDefinition(Node node, SymbolDocs &&docs,
                     std::vector<Constraint> &&constraints,
                     Modifier return_modifier, bool is_method,
                     const std::optional<Identifier> &name_prefix,
                     const Identifier &name, std::vector<Argument> &&arguments,
                     bool are_annotations_same_to_names,
                     std::optional<ExpressionProxy> expression)
      : Node(node), docs_(std::move(docs)),
        constraints_(std::move(constraints)), return_modifier_(return_modifier),
        is_method_(is_method), name_(name), full_name_(name),
        arguments_(std::move(arguments)),
        are_annotations_same_to_names_(are_annotations_same_to_names),
        expression_(expression) {
    if (name_prefix.has_value()) {
      full_name_.append_before(*name_prefix.value().get());
    }
  }

  //

  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);
  }

  //

  Modifier get_return_modifier() const { return return_modifier_; }

  bool is_method() const { return is_method_; }

  //

  Identifier *get_name() { return &name_; }

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

  //

  Identifier *get_full_name() { return &full_name_; }

  const Identifier *get_full_name() const { return &full_name_; }

  //

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

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

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

  //

  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;
  }

  //

  bool are_annotations_same_to_names() const {
    return are_annotations_same_to_names_;
  }

  //

  bool is_same_to(const FunctionDefinition &other_function_definition) const;

  CombineResult combine(FunctionDefinition &&other_function_definition);

private:
  SymbolDocs docs_;
  std::vector<Constraint> constraints_;
  Modifier return_modifier_;
  bool is_method_;
  Identifier name_;
  Identifier full_name_;
  std::vector<Argument> arguments_;
  bool are_annotations_same_to_names_; // needed for easier prinitng process
  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> &&typeclasses,
                 std::vector<Identifier> &&arguments,
                 std::optional<VariantType> &&type)
      : Node(node), docs_(std::move(docs)), is_on_heap_(is_on_heap),
        name_(name), typeclasses_(typeclasses),
        arguments_(std::move(arguments)), type_(std::move(type)) {}

  //

  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;
  }

  //

  bool is_typeclass() { return name_.get_type() == Identifier::TYPECLASS; }

  //

  bool is_same_to(const TypeDefinition &other_type_definition) const;

  CombineResult combine(TypeDefinition &&other_type_definition);

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

class Statement {
public:
  Statement(const Statement &) = default;
  Statement(Statement &&) = default;
  Statement &operator=(const Statement &) = default;

  template <typename T>
  explicit Statement(T &&statement) : expression_(std::forward<T>(statement)) {}

  template <typename T> std::optional<T *> get() {
    if (std::holds_alternative<T>(expression_)) {
      return &std::get<T>(expression_);
    }
    return std::nullopt;
  }

  template <typename T> std::optional<const T *> get() const {
    if (std::holds_alternative<T>(expression_)) {
      return &std::get<T>(expression_);
    }
    return std::nullopt;
  }

  auto get_any() { return &expression_; }

  auto get_any() const { return &expression_; }

  bool is_same_to(const Statement &other_statement) const;

  CombineResult combine(Statement &&other_statement);

private:
  std::variant<Import, TypeDefinition, FunctionDefinition, Extra, EmptyLines>
      expression_;
};

} // namespace nodes