#include <atomic>
#include <functional>
#include <memory_resource>
#include <mutex>
#include <type_traits>

namespace modern {

// Functors namespace (similar to ObjectPoolFunctors)
struct ObjectPoolFunctors {
  using DefaultCreator =
      std::function<void(void *, std::pmr::memory_resource *)>;

  template <class TYPE> class Nil {
  public:
    void operator()(TYPE *) const {}
  };

  template <class TYPE> class Reset {
  public:
    void operator()(TYPE *object) const { object->reset(); }
  };

  template <class TYPE> class Clear {
  public:
    void operator()(TYPE *object) const { object->clear(); }
  };
};

// Default creator implementation
template <class TYPE> class DefaultCreator {
private:
  std::pmr::polymorphic_allocator<TYPE> allocator_;

public:
  explicit DefaultCreator(
      std::pmr::memory_resource *mr = std::pmr::get_default_resource())
      : allocator_(mr) {}

  void operator()(void *arena, std::pmr::memory_resource *mr) {
    if constexpr (std::uses_allocator_v<
                      TYPE, std::pmr::polymorphic_allocator<TYPE>>) {
      // TYPE uses allocator - pass it to constructor
      new (arena) TYPE(std::pmr::polymorphic_allocator<TYPE>(mr));
    } else {
      // TYPE doesn't use allocator - default construct
      new (arena) TYPE{};
    }
  }
};

// Main ObjectPool class
template <class TYPE, class CREATOR = DefaultCreator<TYPE>,
          class RESETTER = ObjectPoolFunctors::Nil<TYPE>>
class ObjectPool {
private:
  // Node structure (similar to Bloomberg's design)
  union ObjectNode {
    struct {
      std::atomic<ObjectNode *> next_p;
      std::atomic<int> refCount;
    } inUse;
    alignas(TYPE) char dummy[sizeof(TYPE)]; // Ensure proper alignment
  };

  // Block structure for batch allocation
  struct BlockNode {
    BlockNode *next_p;
    int numObjects;
  };

  // Memory layout constants
  static constexpr int k_ROUNDED_NUM_OBJECTS =
      (sizeof(TYPE) + sizeof(ObjectNode) - 1) / sizeof(ObjectNode);
  static constexpr int k_NUM_OBJECTS_PER_FRAME = 1 + k_ROUNDED_NUM_OBJECTS;
  static constexpr int k_GROW_FACTOR = 2;
  static constexpr int k_MAX_NUM_OBJECTS = -32;

  // Data members
  std::atomic<ObjectNode *> freeObjectsList_{nullptr};
  CREATOR objectCreator_;
  RESETTER objectResetter_;
  int numReplenishObjects_;
  std::atomic<int> numAvailableObjects_{0};
  std::atomic<int> numObjects_{0};
  BlockNode *blockList_{nullptr};
  std::pmr::memory_resource *memoryResource_;
  std::mutex mutex_;

  // Exception-safe cleanup helper
  class AutoCleanup {
    BlockNode *block_;
    ObjectNode *head_;
    std::pmr::memory_resource *mr_;
    int numNodes_;

  public:
    AutoCleanup(BlockNode *block, ObjectNode *head,
                std::pmr::memory_resource *mr, int numNodes = 0)
        : block_(block), head_(head), mr_(mr), numNodes_(numNodes) {}

    ~AutoCleanup() {
      if (head_) {
        ObjectNode *current = head_ + 1; // Skip to first object
        for (int i = 0; i < numNodes_; ++i) {
          reinterpret_cast<TYPE *>(current + 1)->~TYPE();
          current += k_NUM_OBJECTS_PER_FRAME;
        }
        mr_->deallocate(block_,
                        sizeof(BlockNode) + sizeof(ObjectNode) * numNodes_ *
                                                k_NUM_OBJECTS_PER_FRAME,
                        alignof(BlockNode));
      }
    }

    AutoCleanup &operator++() {
      ++numNodes_;
      return *this;
    }
    void release() {
      head_ = nullptr;
      block_ = nullptr;
    }
  };

  void replenish() {
    int numObjects = numReplenishObjects_ >= 0 ? numReplenishObjects_
                                               : -numReplenishObjects_;
    addObjects(numObjects);

    if (numReplenishObjects_ < 0) {
      if (numReplenishObjects_ > k_MAX_NUM_OBJECTS) {
        numReplenishObjects_ *= k_GROW_FACTOR;
      } else {
        numReplenishObjects_ = -numReplenishObjects_;
      }
    }
  }

  void addObjects(int numObjects) {
    const size_t bytesPerBlock =
        sizeof(BlockNode) +
        sizeof(ObjectNode) * numObjects * k_NUM_OBJECTS_PER_FRAME;

    // Allocate block using PMR
    BlockNode *start = static_cast<BlockNode *>(
        memoryResource_->allocate(bytesPerBlock, alignof(BlockNode)));

    start->next_p = blockList_;
    start->numObjects = numObjects;

    ObjectNode *last = reinterpret_cast<ObjectNode *>(start + 1);
    AutoCleanup guard(start, last, memoryResource_, 0);

    // Create objects
    for (int i = 0; i < numObjects; ++i, ++guard) {
      last->inUse.next_p.store(last + k_NUM_OBJECTS_PER_FRAME,
                               std::memory_order_relaxed);
      last->inUse.refCount.store(0, std::memory_order_relaxed);

      // Use creator to construct object
      objectCreator_(last + 1, memoryResource_);

      last += k_NUM_OBJECTS_PER_FRAME;
    }
    last -= k_NUM_OBJECTS_PER_FRAME;
    last->inUse.refCount.store(0, std::memory_order_relaxed);

    guard.release();
    blockList_ = start;

    // Link to free list (lock-free)
    ObjectNode *firstNode = reinterpret_cast<ObjectNode *>(start + 1);
    ObjectNode *expected = freeObjectsList_.load(std::memory_order_acquire);
    do {
      last->inUse.next_p.store(expected, std::memory_order_relaxed);
    } while (!freeObjectsList_.compare_exchange_weak(
        expected, firstNode, std::memory_order_release,
        std::memory_order_acquire));

    numObjects_.fetch_add(numObjects, std::memory_order_relaxed);
    numAvailableObjects_.fetch_add(numObjects, std::memory_order_relaxed);
  }

public:
  // Constructors
  explicit ObjectPool(int growBy = -1, std::pmr::memory_resource *mr =
                                           std::pmr::get_default_resource())
      : objectCreator_(mr), numReplenishObjects_(growBy), memoryResource_(mr) {
    assert(growBy != 0);
  }

  explicit ObjectPool(
      const CREATOR &creator, int growBy = -1,
      std::pmr::memory_resource *mr = std::pmr::get_default_resource())
      : objectCreator_(creator), numReplenishObjects_(growBy),
        memoryResource_(mr) {
    assert(growBy != 0);
  }

  ObjectPool(const CREATOR &creator, const RESETTER &resetter, int growBy = -1,
             std::pmr::memory_resource *mr = std::pmr::get_default_resource())
      : objectCreator_(creator), objectResetter_(resetter),
        numReplenishObjects_(growBy), memoryResource_(mr) {
    assert(growBy != 0);
  }

  ~ObjectPool() {
    // Destroy all objects
    for (BlockNode *block = blockList_; block; block = block->next_p) {
      int numObjects = block->numObjects;
      ObjectNode *node = reinterpret_cast<ObjectNode *>(block + 1);

      for (int i = 0; i < numObjects; ++i) {
        reinterpret_cast<TYPE *>(node + 1)->~TYPE();
        node += k_NUM_OBJECTS_PER_FRAME;
      }
    }

    // Deallocate blocks
    while (blockList_) {
      BlockNode *next = blockList_->next_p;
      const size_t bytesPerBlock =
          sizeof(BlockNode) +
          sizeof(ObjectNode) * blockList_->numObjects * k_NUM_OBJECTS_PER_FRAME;
      memoryResource_->deallocate(blockList_, bytesPerBlock,
                                  alignof(BlockNode));
      blockList_ = next;
    }
  }

  // Get object (lock-free fast path)
  TYPE *getObject() {
    ObjectNode *node;

    do {
      node = freeObjectsList_.load(std::memory_order_acquire);
      if (!node) {
        std::lock_guard<std::mutex> lock(mutex_);
        node = freeObjectsList_.load(std::memory_order_relaxed);
        if (!node) {
          replenish();
          continue;
        }
      }

      // Try to acquire node with atomic reference counting
      int expected = 0;
      if (node->inUse.refCount.compare_exchange_weak(
              expected, 2, std::memory_order_acq_rel,
              std::memory_order_relaxed)) {

        // Try to remove from free list
        ObjectNode *next = node->inUse.next_p.load(std::memory_order_relaxed);
        if (freeObjectsList_.compare_exchange_weak(node, next,
                                                   std::memory_order_release,
                                                   std::memory_order_acquire)) {
          break;
        }

        // Failed to remove, restore ref count
        node->inUse.refCount.store(0, std::memory_order_release);
      }
    } while (true);

    node->inUse.next_p.store(nullptr, std::memory_order_relaxed);
    numAvailableObjects_.fetch_sub(1, std::memory_order_relaxed);
    return reinterpret_cast<TYPE *>(node + 1);
  }

  // Release object back to pool
  void releaseObject(TYPE *object) {
    ObjectNode *node = reinterpret_cast<ObjectNode *>(object) - 1;

    // Reset object state
    objectResetter_(object);

    // Return to free list
    node->inUse.refCount.store(0, std::memory_order_relaxed);

    ObjectNode *expected = freeObjectsList_.load(std::memory_order_relaxed);
    do {
      node->inUse.next_p.store(expected, std::memory_order_relaxed);
    } while (!freeObjectsList_.compare_exchange_weak(
        expected, node, std::memory_order_release, std::memory_order_relaxed));

    numAvailableObjects_.fetch_add(1, std::memory_order_relaxed);
  }

  // Capacity management
  void increaseCapacity(int numObjects) {
    if (numObjects > 0) {
      std::lock_guard<std::mutex> lock(mutex_);
      addObjects(numObjects);
    }
  }

  void reserveCapacity(int numObjects) {
    std::lock_guard<std::mutex> lock(mutex_);
    int needed = numObjects - numObjects_.load(std::memory_order_relaxed);
    if (needed > 0) {
      addObjects(needed);
    }
  }

  // Accessors
  int numAvailableObjects() const {
    return numAvailableObjects_.load(std::memory_order_relaxed);
  }

  int numObjects() const { return numObjects_.load(std::memory_order_relaxed); }

  std::pmr::memory_resource *memory_resource() const { return memoryResource_; }
};

} // namespace modern