#pragma once

#include <algorithm>
#include <atomic>
#include <cstdint>
#include <cstdio>
#include <cstdlib>

/**
 * @brief High-performance thread-safe reference counting system
 *
 * This library provides custom smart pointers with shared/weak semantics,
 * designed for better performance than std::shared_ptr/weak_ptr.
 *
 * Key features:
 * - Thread-safe reference counting using atomic operations
 * - Weak references to break circular dependencies
 * - Single allocation for better cache locality
 * - Optimized for copy/move operations
 * - Compatible with std::shared_ptr semantics
 *
 * Basic usage:
 * @code
 * auto obj = make_ref<MyClass>(args...);     // Create managed object
 * auto copy = obj;                           // Copy (thread-safe)
 * WeakRef<MyClass> weak = obj;               // Create weak reference
 * auto locked = weak.lock();                 // Try to promote to strong
 * @endcode
 *
 * Thread safety: All operations are thread-safe. Multiple threads can
 * safely copy, move, and destroy references to the same object.
 */

namespace detail {
struct ControlBlock {
  std::atomic<uint32_t> strong_count;
  std::atomic<uint32_t> weak_count;

  ControlBlock()
      : strong_count(1), weak_count(1) {
  } // Start with 1 strong, 1 weak (biased)

  /**
   * @brief Increment strong reference count
   * @return Previous strong count
   */
  uint32_t increment_strong() noexcept {
    return strong_count.fetch_add(1, std::memory_order_relaxed);
  }

  /**
   * @brief Decrement strong reference count
   * @return Previous strong count
   */
  uint32_t decrement_strong() noexcept {
    return strong_count.fetch_sub(1, std::memory_order_acq_rel);
  }

  /**
   * @brief Increment weak reference count
   * @return Previous weak count
   */
  uint32_t increment_weak() noexcept {
    return weak_count.fetch_add(1, std::memory_order_relaxed);
  }

  /**
   * @brief Decrement weak reference count
   * @return Previous weak count
   */
  uint32_t decrement_weak() noexcept {
    return weak_count.fetch_sub(1, std::memory_order_acq_rel);
  }
};
} // namespace detail

/**
 * @brief Strong reference to a shared object (similar to std::shared_ptr)
 *
 * Ref<T> manages shared ownership of an object. The object is automatically
 * destroyed when the last Ref pointing to it is destroyed.
 *
 * Usage:
 * - Use make_ref<T>() to create new objects
 * - Copy/assign to share ownership
 * - Use get(), operator*, operator-> to access the object
 * - Use operator bool() to check if valid
 * - Use reset() to release ownership
 *
 * Limitations compared to std::shared_ptr:
 * - Cannot take ownership of raw pointers
 * - Objects can only be created via make_ref<T>() for proper memory layout
 * - No custom deleter support
 * - No enable_shared_from_this / shared_from_this() integration
 * - No aliasing constructor (sharing ownership with different pointer)
 * - No array support (Ref<T[]>)
 * - No atomic operations (atomic_load, atomic_store, etc.)
 *
 * Thread safety: All operations are thread-safe. The managed object
 * itself is NOT automatically thread-safe.
 */
template <typename T> struct Ref {
  /**
   * @brief Get raw pointer to managed object
   */
  T *get() const noexcept { return ptr; }

  /**
   * @brief Dereference operator
   */
  T &operator*() const { return *ptr; }

  /**
   * @brief Arrow operator
   */
  T *operator->() const { return ptr; }

  /**
   * @brief Check if Ref is valid (not empty)
   */
  explicit operator bool() const noexcept { return ptr != nullptr; }

  /**
   * @brief Destructor - decrements strong reference count
   */
  ~Ref() { release(); }

  /**
   * @brief Copy constructor - increments strong reference count
   */
  Ref(const Ref &other) noexcept
      : ptr(other.ptr), control_block(other.control_block) {
    if (control_block) {
      control_block->increment_strong();
    }
  }

  /**
   * @brief Converting copy constructor for polymorphism (Derived -> Base)
   */
  template <typename U>
  Ref(const Ref<U> &other) noexcept
    requires std::is_convertible_v<U *, T *>
      : ptr(other.ptr), control_block(other.control_block) {
    if (control_block) {
      control_block->increment_strong();
    }
  }

  /**
   * @brief Copy assignment operator
   */
  Ref &operator=(const Ref &other) noexcept {
    if (this != &other) {
      release();
      ptr = other.ptr;
      control_block = other.control_block;
      if (control_block) {
        control_block->increment_strong();
      }
    }
    return *this;
  }

  /**
   * @brief Converting assignment operator for polymorphism (Derived -> Base)
   */
  template <typename U>
  Ref &operator=(const Ref<U> &other) noexcept
    requires std::is_convertible_v<U *, T *>
  {
    release();
    ptr = other.ptr;
    control_block = other.control_block;
    if (control_block) {
      control_block->increment_strong();
    }
    return *this;
  }

  /**
   * @brief Move constructor - transfers ownership
   */
  Ref(Ref &&other) noexcept
      : ptr(other.ptr), control_block(other.control_block) {
    other.ptr = nullptr;
    other.control_block = nullptr;
  }

  /**
   * @brief Converting move constructor for polymorphism (Derived -> Base)
   */
  template <typename U>
  Ref(Ref<U> &&other) noexcept
    requires std::is_convertible_v<U *, T *>
      : ptr(other.ptr), control_block(other.control_block) {
    other.ptr = nullptr;
    other.control_block = nullptr;
  }

  /**
   * @brief Move assignment operator
   */
  Ref &operator=(Ref &&other) noexcept {
    if (this != &other) {
      release();
      ptr = other.ptr;
      control_block = other.control_block;
      other.ptr = nullptr;
      other.control_block = nullptr;
    }
    return *this;
  }

  /**
   * @brief Converting move assignment operator for polymorphism (Derived ->
   * Base)
   */
  template <typename U>
  Ref &operator=(Ref<U> &&other) noexcept
    requires std::is_convertible_v<U *, T *>
  {
    release();
    ptr = other.ptr;
    control_block = other.control_block;
    other.ptr = nullptr;
    other.control_block = nullptr;
    return *this;
  }

  /**
   * @brief Reset to empty state
   */
  void reset() noexcept {
    release();
    ptr = nullptr;
    control_block = nullptr;
  }

  /**
   * @brief Equality comparison
   */
  bool operator==(const Ref &other) const noexcept {
    return control_block == other.control_block;
  }

  /**
   * @brief Inequality comparison
   */
  bool operator!=(const Ref &other) const noexcept { return !(*this == other); }

  /**
   * @brief Default constructor - creates empty Ref
   */
  Ref() : ptr(nullptr), control_block(nullptr) {}

private:
  explicit Ref(T *object_ptr, detail::ControlBlock *cb)
      : ptr(object_ptr), control_block(cb) {}

  T *ptr;
  detail::ControlBlock *control_block;

  /**
   * @brief Release current reference and handle cleanup
   */
  void release() noexcept {
    if (control_block) {
      uint32_t prev_strong = control_block->decrement_strong();

      // If this was the last strong reference, destroy the object
      if (prev_strong == 1) {
        // We need to call the destructor before we decrement the weak count, to
        // account for the possibility that T has a WeakRef to itself.
        ptr->~T();

        // Release the bias - decrement weak count for strong references
        uint32_t prev_weak = control_block->decrement_weak();

        // If weak count hits 0, destroy and free control block
        if (prev_weak == 1) {
          control_block->~ControlBlock();
          std::free(control_block);
        }
      }
    }
  }

  template <typename U, typename... Args>
  friend Ref<U> make_ref(Args &&...args);

  template <typename U> friend struct WeakRef;
  template <typename U> friend struct Ref;
};

/**
 * @brief Weak reference to a shared object (similar to std::weak_ptr)
 *
 * WeakRef<T> holds a non-owning reference to an object managed by Ref<T>.
 * It can be used to break circular dependencies and safely observe objects
 * that might be destroyed by other threads.
 *
 * Usage:
 * - Create from Ref<T> to observe without owning
 * - Use lock() to attempt promotion to Ref<T>
 * - Returns empty Ref<T> if object was already destroyed
 * - Use reset() to stop observing
 *
 * Self-referencing pattern: Objects can safely contain WeakRef members
 * pointing to themselves. The implementation ensures proper destruction
 * order to prevent use-after-free when the object destructor runs.
 *
 * Thread safety: All operations are thread-safe. The observed object
 * may be destroyed by other threads at any time.
 */
template <typename T> struct WeakRef {
  /**
   * @brief Attempt to promote WeakRef to Ref
   * @return Valid Ref if object still alive, empty Ref otherwise
   */
  Ref<T> lock() const {
    if (!control_block) {
      return Ref<T>();
    }

    // Try to increment strong count if it's not zero
    uint32_t expected_strong =
        control_block->strong_count.load(std::memory_order_relaxed);
    while (expected_strong > 0) {
      // Try to increment the strong count
      if (control_block->strong_count.compare_exchange_weak(
              expected_strong, expected_strong + 1, std::memory_order_acquire,
              std::memory_order_relaxed)) {
        // Success - we incremented the strong count
        return Ref<T>(get_object_ptr(), control_block);
      }
      // CAS failed, expected_strong now contains the current value, retry
    }

    // Strong count was 0, object is being destroyed
    return Ref<T>();
  }

  /**
   * @brief Destructor - decrements weak reference count
   */
  ~WeakRef() { release(); }

  /**
   * @brief Copy constructor from WeakRef
   */
  WeakRef(const WeakRef &other) noexcept : control_block(other.control_block) {
    if (control_block) {
      control_block->increment_weak();
    }
  }

  /**
   * @brief Copy constructor from Ref
   */
  WeakRef(const Ref<T> &ref) noexcept : control_block(ref.control_block) {
    if (control_block) {
      control_block->increment_weak();
    }
  }

  /**
   * @brief Converting copy constructor from WeakRef for polymorphism
   */
  template <typename U>
  WeakRef(const WeakRef<U> &other) noexcept
    requires std::is_convertible_v<U *, T *>
      : control_block(other.control_block) {
    if (control_block) {
      control_block->increment_weak();
    }
  }

  /**
   * @brief Converting copy constructor from Ref for polymorphism
   */
  template <typename U>
  WeakRef(const Ref<U> &ref) noexcept
    requires std::is_convertible_v<U *, T *>
      : control_block(ref.control_block) {
    if (control_block) {
      control_block->increment_weak();
    }
  }

  /**
   * @brief Converting copy assignment from WeakRef for polymorphism
   */
  template <typename U>
  WeakRef &operator=(const WeakRef<U> &other) noexcept
    requires std::is_convertible_v<U *, T *>
  {
    release();
    control_block = other.control_block;
    if (control_block) {
      control_block->increment_weak();
    }
    return *this;
  }

  /**
   * @brief Converting copy assignment from Ref for polymorphism
   */
  template <typename U>
  WeakRef &operator=(const Ref<U> &ref) noexcept
    requires std::is_convertible_v<U *, T *>
  {
    release();
    control_block = ref.control_block;
    if (control_block) {
      control_block->increment_weak();
    }
    return *this;
  }

  /**
   * @brief Converting move constructor from WeakRef for polymorphism
   */
  template <typename U>
  WeakRef(WeakRef<U> &&other) noexcept
    requires std::is_convertible_v<U *, T *>
      : control_block(other.control_block) {
    other.control_block = nullptr;
  }

  /**
   * @brief Converting move assignment from WeakRef for polymorphism
   */
  template <typename U>
  WeakRef &operator=(WeakRef<U> &&other) noexcept
    requires std::is_convertible_v<U *, T *>
  {
    release();
    control_block = other.control_block;
    other.control_block = nullptr;
    return *this;
  }

  /**
   * @brief Copy assignment from WeakRef
   */
  WeakRef &operator=(const WeakRef &other) noexcept {
    if (this != &other) {
      release();
      control_block = other.control_block;
      if (control_block) {
        control_block->increment_weak();
      }
    }
    return *this;
  }

  /**
   * @brief Copy assignment from Ref
   */
  WeakRef &operator=(const Ref<T> &ref) noexcept {
    release();
    control_block = ref.control_block;
    if (control_block) {
      control_block->increment_weak();
    }
    return *this;
  }

  /**
   * @brief Move constructor
   */
  WeakRef(WeakRef &&other) noexcept : control_block(other.control_block) {
    other.control_block = nullptr;
  }

  /**
   * @brief Move assignment
   */
  WeakRef &operator=(WeakRef &&other) noexcept {
    if (this != &other) {
      release();
      control_block = other.control_block;
      other.control_block = nullptr;
    }
    return *this;
  }

  /**
   * @brief Reset to empty state
   */
  void reset() noexcept {
    release();
    control_block = nullptr;
  }

  /**
   * @brief Default constructor - creates empty WeakRef
   */
  WeakRef() : control_block(nullptr) {}

private:
  explicit WeakRef(detail::ControlBlock *cb) : control_block(cb) {}

  detail::ControlBlock *control_block;

  // Helper to calculate object pointer from control block
  T *get_object_ptr() const {
    if (!control_block)
      return nullptr;
    constexpr size_t cb_size = sizeof(detail::ControlBlock);
    constexpr size_t alignment = alignof(T);
    constexpr size_t padded_cb_size =
        (cb_size + alignment - 1) & ~(alignment - 1);
    return reinterpret_cast<T *>(reinterpret_cast<char *>(control_block) +
                                 padded_cb_size);
  }

  /**
   * @brief Release current weak reference and handle cleanup
   */
  void release() noexcept {
    if (control_block) {
      uint32_t prev_weak = control_block->decrement_weak();

      // If weak count hits 0, destroy and free control block
      if (prev_weak == 1) {
        control_block->~ControlBlock();
        std::free(control_block);
      }
    }
  }

  template <typename U> friend struct Ref;
  template <typename U> friend struct WeakRef;
};

/**
 * @brief Create a new managed object wrapped in Ref<T>
 *
 * This is the only way to create Ref<T> objects. It performs a single
 * allocation for both the control block and object, improving cache locality.
 *
 * @tparam T Type of object to create
 * @tparam Args Types of constructor arguments
 * @param args Arguments forwarded to T's constructor
 * @return Ref<T> managing the newly created object
 *
 * Example:
 * @code
 * auto obj = make_ref<MyClass>(arg1, arg2);
 * auto empty_vec = make_ref<std::vector<int>>();
 * @endcode
 *
 * Thread safety: Safe to call from multiple threads simultaneously.
 */
template <typename T, typename... Args> Ref<T> make_ref(Args &&...args) {
  constexpr size_t cb_size = sizeof(detail::ControlBlock);
  constexpr size_t alignment = alignof(T);
  constexpr size_t padded_cb_size =
      (cb_size + alignment - 1) & ~(alignment - 1);

  constexpr size_t total_alignment =
      std::max(alignof(detail::ControlBlock), alignment);
  constexpr size_t total_size = padded_cb_size + sizeof(T);
  constexpr size_t aligned_total_size =
      (total_size + total_alignment - 1) & ~(total_alignment - 1);

  char *buf = reinterpret_cast<char *>(
      std::aligned_alloc(total_alignment, aligned_total_size));
  if (!buf) {
    std::fprintf(stderr, "Out of memory\n");
    std::abort();
  }

  auto *cb = new (buf) detail::ControlBlock();
  T *obj = new (buf + padded_cb_size) T{std::forward<Args>(args)...};
  return Ref<T>(obj, cb);
}