#pragma once

#include "basic_nodes.hpp"
#include "name_tree.hpp"
#include "statement_nodes.hpp"

#include <optional>
#include <vector>

namespace types {

class Type;
class TypeStorage;

class TypeProxy {
  friend TypeStorage;

public:
  Type *get();

  const Type *get() const;

private:
  TypeProxy(TypeStorage &type_storage, size_t id)
      : type_storage_(&type_storage), id_(id) {}

private:
  TypeStorage *type_storage_;
  size_t id_;
};

class Defined {
public:
  Defined(names::StatementProxy<nodes::TypeDefinition> definition)
      : definition_(definition) {}

  nodes::TypeDefinition *get_definition() { return definition_.get(); }

  const nodes::TypeDefinition *get_definition() const {
    return definition_.get();
  }

private:
  names::StatementProxy<nodes::TypeDefinition> definition_;
};

class Abstract {
public:
  Abstract(std::optional<std::string> name,
           std::vector<std::string> typeclasses)
      : name_(name), typeclasses_(typeclasses) {}

  //

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

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

  //

  size_t typeclasses_size() const { return typeclasses_.size(); }

  std::string *get_typeclass(size_t id) { return &typeclasses_.at(id); }

  const std::string *get_typeclass(size_t id) const {
    return &typeclasses_.at(id);
  }

private:
  std::optional<std::string> name_;
  std::vector<std::string> typeclasses_;
};

class Modified {
public:
  Modified(nodes::Modifier modifier, TypeProxy type)
      : modifier_(modifier), type_(type) {}

  nodes::Modifier get_modifier() const { return modifier_; }

  Type *get_type() { return type_.get(); }

  const Type *get_type() const { return type_.get(); }

  const TypeProxy &get_type_proxy() const { return type_; }

private:
  nodes::Modifier modifier_;
  TypeProxy type_;
};

class Container {
public:
  enum ContainerType {
    OR,
    AND,
  };

  Container(ContainerType type, std::vector<TypeProxy> &&fields)
      : type_(type), fields_(std::move(fields)) {}

  Container(ContainerType type, const std::vector<TypeProxy> &fields)
      : type_(type), fields_(std::move(fields)) {}

  //

  ContainerType get_type() const { return type_; }

  //

  size_t fields_size() const { return fields_.size(); }

  Type *get_field(size_t id) { return fields_.at(id).get(); }

  const Type *get_field(size_t id) const { return fields_.at(id).get(); }

  const TypeProxy &get_field_proxy(size_t id) const { return fields_.at(id); }

private:
  ContainerType type_;
  std::vector<TypeProxy> fields_; // or constructors
};

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

  template <typename T>
  explicit Type(T &&type) : type_(std::forward<T>(type)) {}

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

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

  auto get_any() { return &type_; }

  auto get_any() const { return &type_; }

private:
  std::variant<Defined, Abstract, Modified, Container> type_;
};

class TypeStorage {
  friend TypeProxy;

public:
  TypeProxy add_type(const Type &type) {
    storage_.push_back(type);
    return TypeProxy(*this, storage_.size() - 1);
  }

  TypeProxy add_type(Type &&type) {
    storage_.push_back(std::move(type));
    return TypeProxy(*this, storage_.size() - 1);
  }

private:
  Type *get_type(size_t id) { return &storage_.at(id); }

  const Type *get_type(size_t id) const { return &storage_.at(id); }

private:
  std::vector<Type> storage_;
};

}; // namespace types