#pragma once

#include <cstdint>
#include <functional>
#include <memory>
#include <span>
#include <vector>

class HeapMemoryFixedSizePool {
public:
    // Define our custom allocator and deleter signatures
    using AllocFunc = std::function<void *(size_t size, int32_t alignment)>;
    using FreeFunc  = std::function<void(void *ptr)>;

    class MemorySegmentInfo {
    public:
        MemorySegmentInfo(void *data, size_t size);

        [[nodiscard]] void *data() const;
        [[nodiscard]] size_t size() const;

    private:
        void *mData;
        size_t mSize;
    };

    HeapMemoryFixedSizePool();

    // Updated constructors to take the custom functions
    HeapMemoryFixedSizePool(const AllocFunc &allocFunc, FreeFunc freeFunc, std::initializer_list<size_t> segmentSizes);
    HeapMemoryFixedSizePool(const AllocFunc &allocFunc, FreeFunc freeFunc, std::span<const std::size_t> segmentSizes);

    // Delete the copy constructor and copy assignment operator
    HeapMemoryFixedSizePool(const HeapMemoryFixedSizePool &) = delete;
    HeapMemoryFixedSizePool &operator=(const HeapMemoryFixedSizePool &) = delete;

    HeapMemoryFixedSizePool(HeapMemoryFixedSizePool &&other) noexcept;

    HeapMemoryFixedSizePool &operator=(HeapMemoryFixedSizePool &&other) noexcept;

    [[nodiscard]] uint32_t numberOfSegments() const;

    explicit operator bool() const;

    MemorySegmentInfo operator[](size_t idx) const;

    [[nodiscard]] std::span<const uint8_t> dataView() const;

private:
    struct CustomDeleter {
        FreeFunc freeFunc;

        void operator()(uint8_t *ptr) const {
            if (ptr != nullptr && freeFunc) {
                freeFunc(ptr);
            }
        }
    };

    // Initialize with a null pointer and an empty CustomDeleter
    std::unique_ptr<uint8_t[], CustomDeleter> mData{nullptr, CustomDeleter{nullptr}};
    std::size_t mTotalSize{};
    std::vector<MemorySegmentInfo> mSegmentInfos;
};