lang/include/statement_nodes.hpp

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

  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()) {}

    //

    bool add_type(const std::optional<std::string> &annotation, TypeProxy type,
                  nodes::Modifier before_modifier = Modifier::NONE) {
      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,
                     ModifierType modifier, 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)), modifier_(modifier), name_(name),
        arguments_(std::move(arguments)),
        are_annotations_same_to_names_(are_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_; }

  //

  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 combine(FunctionDefinition &&other_function_definition);

private:
  SymbolDocs docs_;
  std::vector<Constraint> constraints_;
  ModifierType modifier_;
  Identifier 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> &&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);
  }

  //

  bool combine(TypeDefinition &&other_type_definition);

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

  //

  bool combine(TypeclassDefinition &&other_typeclass_definition);

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

class Statement {
public:
  template <typename T>
  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_; }

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

} // namespace nodes