#pragma once

#include "basic_nodes.hpp"

#include <memory>
#include <optional>
#include <string>
#include <unordered_map>
#include <variant>
#include <vector>

namespace nodes {

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 Type : public Node {
public:
  Type(Node node, Identifier &&identifier, bool is_on_heap = false,
       bool is_optional = false)
      : Node(node), name_(std::move(identifier)), is_on_heap_(is_on_heap),
        is_optional_(is_optional) {}

  Type(Node node, const Identifier &identifier, bool is_on_heap = false,
       bool is_optional = false)
      : Node(node), name_(identifier), is_on_heap_(is_on_heap),
        is_optional_(is_optional) {}

  Type(Node node, Identifier &&identifier, std::vector<TypeProxy> &&parameters,
       bool is_on_heap = false, bool is_optional = false)
      : Node(node), name_(std::move(identifier)),
        parameters_(std::move(parameters)), is_on_heap_(is_on_heap),
        is_optional_(is_optional) {}

  Type(Node node, const Identifier &identifier,
       std::vector<TypeProxy> &&parameters, bool is_on_heap = false,
       bool is_optional = false)
      : Node(node), name_(identifier), parameters_(std::move(parameters)),
        is_on_heap_(is_on_heap), is_optional_(is_optional) {}

  std::string *get_name() { return name_.get(); }

  const std::string *get_name() const { return name_.get(); }

  size_t get_parametrs_size() { return parameters_.size(); }

  Type *get_parameter(size_t id) { return parameters_.at(id).get(); }

  const Type *get_parameter(size_t id) const {
    return parameters_.at(id).get();
  }

  bool is_on_heap() { return is_on_heap_; }

  bool is_optional() { return is_optional_; }

private:
  Identifier name_;
  std::vector<TypeProxy> parameters_;
  // or use allocator ??
  bool is_on_heap_ = false;
  bool is_optional_ = false;
};

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

class TupleType : public Node {
public:
  TupleType(Node node,
            const std::vector<std::pair<std::optional<std::string>, TypeProxy>>
                &fields)
      : Node(node) {
    fields_.reserve(fields.size());
    for (auto &field : fields) {
      if (field.first.has_value()) {
        annotation_ids_[field.first.value()] = fields_.size();
      }
      fields_.push_back(field.second);
    }
  }

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

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

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

  Type *get(const std::string &annotation) {
    return fields_.at(annotation_ids_.at(annotation)).get();
  }

  const Type *get(const std::string &annotation) const {
    return fields_.at(annotation_ids_.at(annotation)).get();
  }

private:
  std::unordered_map<std::string, size_t> annotation_ids_; // Annotations
  std::vector<TypeProxy> fields_;
};

class VariantType : public Node {
public:
  VariantType(Node node, std::vector<TupleType> &&constructors_)
      : Node(node), constructors_(std::move(constructors_)) {}

  VariantType(Node node, const std::vector<TupleType> &constructors_)
      : Node(node), constructors_(constructors_) {}

  size_t size() { return constructors_.size(); }

  TupleType *get(size_t id) { return &constructors_.at(id); }

  const TupleType *get(size_t id) const { return &constructors_.at(id); }

private:
  // named constructors ??
  std::vector<TupleType> constructors_;
};

} // namespace nodes