#pragma once

#include <cstddef>
#include <cstdint>
#include <memory>
#include <mutex>
#include <queue>
#include <vector>

namespace softbus::core::memory
{

class MemoryPool
{
public:
    struct PoolBlock
    {
        std::uint8_t* data{nullptr};
        std::size_t capacity{0};
        std::size_t size{0};
        std::size_t slotIndex{0};
    };

    using BlockPtr = std::shared_ptr<PoolBlock>;

    explicit MemoryPool(std::size_t blockSize = 4096, std::size_t blockCount = 1024)
        : m_blockSize(blockSize), m_storage(blockSize * blockCount), m_inUse(blockCount, false)
    {
        for (std::size_t i = 0; i < blockCount; ++i) {
            m_freeSlots.push(i);
        }
    }

    BlockPtr allocate(std::size_t requiredSize)
    {
        if (requiredSize > m_blockSize) {
            return {};
        }

        std::size_t slot = 0;
        {
            std::lock_guard<std::mutex> lk(m_mutex);
            if (m_freeSlots.empty()) {
                return {};
            }
            slot = m_freeSlots.front();
            m_freeSlots.pop();
            m_inUse[slot] = true;
        }

        auto* raw = new PoolBlock{
            m_storage.data() + slot * m_blockSize,
            m_blockSize,
            requiredSize,
            slot,
        };
        return BlockPtr(raw, [this](PoolBlock* b) {
            if (!b) return;
            this->release(*b);
            delete b;
        });
    }

    void release(const PoolBlock& block)
    {
        std::lock_guard<std::mutex> lk(m_mutex);
        if (block.slotIndex >= m_inUse.size()) {
            return;
        }
        if (!m_inUse[block.slotIndex]) {
            return;
        }
        m_inUse[block.slotIndex] = false;
        m_freeSlots.push(block.slotIndex);
    }

    std::size_t blockSize() const { return m_blockSize; }
    std::size_t totalBlocks() const { return m_inUse.size(); }

    std::size_t inUseBlocks() const
    {
        std::lock_guard<std::mutex> lk(m_mutex);
        std::size_t used = 0;
        for (const bool v : m_inUse) {
            if (v) ++used;
        }
        return used;
    }

private:
    std::size_t m_blockSize;
    std::vector<std::uint8_t> m_storage;
    mutable std::mutex m_mutex;
    std::queue<std::size_t> m_freeSlots;
    std::vector<bool> m_inUse;
};

} // namespace softbus::core::memory
