#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