#pragma once

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

/**
 * @brief Thread-safe reference counting abstraction with shared/weak pointer
 * semantics
 *
 * TODO: Implement custom reference counting system with:
 * - Thread-safe reference counting using atomic operations
 * - Weak reference support to break circular dependencies
 * - Move semantics for efficient transfers
 * - Custom deleter support
 * - Zero-overhead when not using weak references
 */

namespace detail {
struct ControlBlock {
  // Least significant 32 bits are strong reference count
  // Most significant 32 bits are weak reference count
  std::atomic<uint64_t> ref_counts;

  ControlBlock() : ref_counts(1) {} // Start with 1 strong reference

  /**
   * @brief Increment strong reference count
   * @return Previous ref_counts value (both strong and weak counts)
   */
  uint64_t increment_strong() noexcept {
    uint64_t old_value;
    uint64_t new_value;
    do {
      old_value = ref_counts.load(std::memory_order_relaxed);
      uint32_t strong_count = static_cast<uint32_t>(old_value);
      uint32_t weak_count = static_cast<uint32_t>(old_value >> 32);
      new_value =
          (static_cast<uint64_t>(weak_count) << 32) | (strong_count + 1);
    } while (!ref_counts.compare_exchange_weak(old_value, new_value,
                                               std::memory_order_relaxed));

    return old_value;
  }

  /**
   * @brief Decrement strong reference count
   * @return Previous ref_counts value (both strong and weak counts)
   */
  uint64_t decrement_strong() noexcept {
    uint64_t old_value;
    uint64_t new_value;
    do {
      old_value = ref_counts.load(std::memory_order_relaxed);
      uint32_t strong_count = static_cast<uint32_t>(old_value);
      uint32_t weak_count = static_cast<uint32_t>(old_value >> 32);
      new_value =
          (static_cast<uint64_t>(weak_count) << 32) | (strong_count - 1);
    } while (!ref_counts.compare_exchange_weak(old_value, new_value,
                                               std::memory_order_acq_rel));

    return old_value;
  }

  /**
   * @brief Increment weak reference count
   * @return Previous ref_counts value (both strong and weak counts)
   */
  uint64_t increment_weak() noexcept {
    uint64_t old_value;
    uint64_t new_value;
    do {
      old_value = ref_counts.load(std::memory_order_relaxed);
      uint32_t strong_count = static_cast<uint32_t>(old_value);
      uint32_t weak_count = static_cast<uint32_t>(old_value >> 32);
      new_value = (static_cast<uint64_t>(weak_count + 1) << 32) | strong_count;
    } while (!ref_counts.compare_exchange_weak(old_value, new_value,
                                               std::memory_order_relaxed));

    return old_value;
  }

  /**
   * @brief Decrement weak reference count
   * @return Previous ref_counts value (both strong and weak counts)
   */
  uint64_t decrement_weak() noexcept {
    uint64_t old_value;
    uint64_t new_value;
    do {
      old_value = ref_counts.load(std::memory_order_relaxed);
      uint32_t strong_count = static_cast<uint32_t>(old_value);
      uint32_t weak_count = static_cast<uint32_t>(old_value >> 32);
      new_value = (static_cast<uint64_t>(weak_count - 1) << 32) | strong_count;
    } while (!ref_counts.compare_exchange_weak(old_value, new_value,
                                               std::memory_order_acq_rel));

    return old_value;
  }
};
} // namespace detail

template <typename T> struct Ref {
  /**
   * @brief Get raw pointer to managed object
   */
  T *get() 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 Dereference operator
   */
  T &operator*() const { return *get(); }

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

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

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

  /**
   * @brief Copy constructor - increments strong reference count
   */
  Ref(const Ref &other) noexcept : 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();
      control_block = other.control_block;
      if (control_block) {
        control_block->increment_strong();
      }
    }
    return *this;
  }

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

  /**
   * @brief Move assignment operator
   */
  Ref &operator=(Ref &&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 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() : control_block(nullptr) {}

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

  detail::ControlBlock *control_block;

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

      // If this was the last strong reference, destroy the object
      if (prev_strong == 1) {
        T *obj = get();
        obj->~T();

        // If no weak references either, free the entire allocation
        uint32_t prev_weak = static_cast<uint32_t>(prev >> 32);
        if (prev_weak == 0) {
          std::free(control_block);
        }
      }
    }
  }

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

  template <typename U> friend struct WeakRef;
};

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() {
    if (!control_block) {
      return Ref<T>();
    }
    uint64_t old_value;
    uint64_t new_value;
    do {
      // Use acquire ordering to ensure that any subsequent use of the returned
      // Ref (like dereferencing the object pointer) cannot be reordered before
      // this safety check. This would ideally use memory_order_consume for
      // dependency ordering, but the folk wisdom is "don't use that".
      old_value = control_block->ref_counts.load(std::memory_order_acquire);
      uint32_t strong_count = static_cast<uint32_t>(old_value);

      // If strong count is 0, object is being destroyed
      if (strong_count == 0) {
        return Ref<T>();
      }

      uint32_t weak_count = static_cast<uint32_t>(old_value >> 32);
      new_value =
          (static_cast<uint64_t>(weak_count) << 32) | (strong_count + 1);
    } while (!control_block->ref_counts.compare_exchange_weak(
        old_value, new_value, std::memory_order_relaxed));

    return Ref<T>(control_block);
  }

  /**
   * @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 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;

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

      // If this was the last weak reference and no strong references, free
      // control block
      if (prev_weak == 1 && prev_strong == 0) {
        std::free(control_block);
      }
    }
  }

  template <typename U> friend struct Ref;
};

/**
 * @brief Create a new Ref with object constructed in-place after control block
 */
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();
  new (buf + padded_cb_size) T{std::forward<Args>(args)...};
  return Ref<T>(cb);
}