module;
#include <atomic>
#include <cstdint>

export module VUI:Ref;

namespace VUI {

export struct RefCounted {
  virtual ~RefCounted() = default;

  void incrementRefCount() const { ++m_RefCount; }
  void decrementRefCount() const { --m_RefCount; }

  [[nodiscard]] std::uint32_t getRefCount() const { return m_RefCount.load(); }

private:
  mutable std::atomic<uint32_t> m_RefCount = 0;
};

template <class T>
concept IsRefCounted = std::is_base_of_v<RefCounted, T>;

export template <typename T> class Ref {
public:
  static_assert(IsRefCounted<T>, "Should be RefCounted");
  Ref() : m_Instance(nullptr) {}

  explicit Ref(std::nullptr_t) : m_Instance(nullptr) {}

  explicit Ref(T *instance) : m_Instance(instance) { incRef(); }

  template <typename T2> explicit Ref(const Ref<T2> &other) {
    m_Instance = static_cast<T *>(other.m_Instance);
    incRef();
  }

  template <typename T2> explicit Ref(Ref<T2> &&other) {
    m_Instance = static_cast<T *>(other.m_Instance);
    other.m_Instance = nullptr;
  }

  ~Ref() { decRef(); }

  Ref(const Ref &other) : m_Instance(other.m_Instance) { incRef(); }

  Ref &operator=(std::nullptr_t) {
    decRef();
    m_Instance = nullptr;
    return *this;
  }

  Ref &operator=(const Ref &other) {
    if (this == &other)
      return *this;

    other.incRef();
    decRef();

    m_Instance = other.m_Instance;
    return *this;
  }

  template <typename T2> Ref &operator=(const Ref<T2> &other) {
    other.incRef();
    decRef();

    m_Instance = other.m_Instance;
    return *this;
  }

  template <typename T2> Ref &operator=(Ref<T2> &&other) {
    decRef();

    m_Instance = other.m_Instance;
    other.m_Instance = nullptr;
    return *this;
  }

  operator bool() { return m_Instance != nullptr; }
  operator bool() const { return m_Instance != nullptr; }

  T *operator->() { return m_Instance; }
  const T *operator->() const { return m_Instance; }

  T &operator*() { return *m_Instance; }
  const T &operator*() const { return *m_Instance; }

  T *value() { return m_Instance; }
  const T *value() const { return m_Instance; }

  void reset(T *instance = nullptr) {
    decRef();
    m_Instance = instance;
  }

  template <typename T2> Ref<T2> as() const { return Ref<T2>(*this); }

  template <typename... Args> static Ref Create(Args &&...args) {
    return Ref(new T(std::forward<Args>(args)...));
  }

  bool operator==(const Ref &other) const {
    return m_Instance == other.m_Instance;
  }

  bool operator!=(const Ref &other) const { return !(*this == other); }

  bool equals(const Ref &other) {
    if (!m_Instance || !other.m_Instance)
      return false;

    return m_Instance == other.m_Instance;
  }

private:
  void incRef() const {
    if (m_Instance) {
      m_Instance->incrementRefCount();
    }
  }

  void decRef() const {
    if (m_Instance) {
      m_Instance->decrementRefCount();

      if (m_Instance->getRefCount() == 0) {
        delete m_Instance;
        m_Instance = nullptr;
      }
    }
  }

  template <typename T2> friend class Ref;
  mutable T *m_Instance;
};
} // namespace VUI