// Author: Simon-Pierre Boucher — contact@spboucher.ai
#pragma once

#include <cstddef>
#include <mutex>
#include <unordered_map>
#include <vector>

namespace MTL {
class Buffer;
class Device;
}

namespace forge {

// Size-bucketed MTLBuffer pool. Every tensor allocation goes through here so
// the training loop never touches newBuffer (a kernel-level VM allocation).
// Buffers use MTLStorageModeShared: the CPU pointer is valid at all times.
//
// Sizes are rounded up to power-of-two buckets (min 256 bytes, which also
// satisfies every setBuffer offset-alignment rule). release() returns the
// buffer to its bucket; it is the caller's job (Storage refcounting, and
// later the command-buffer retire list) to only release buffers no in-flight
// command buffer references.
class Allocator {
public:
    explicit Allocator(MTL::Device* device);
    ~Allocator();

    Allocator(const Allocator&) = delete;
    Allocator& operator=(const Allocator&) = delete;

    // Returns a retained buffer of capacity >= nbytes.
    MTL::Buffer* acquire(size_t nbytes);
    void release(MTL::Buffer* buffer);

    // While deferring (a command buffer is being encoded / in flight),
    // released buffers park on a retire list instead of the free pool:
    // recycling them would let CPU-side writes (zeros, fills) race pending
    // GPU reads. The Stream flips this on at first encode and flushes at
    // sync.
    void set_defer(bool defer);
    void flush_retired();

    // Drop all pooled (free) buffers back to the OS.
    void trim();

    size_t bytes_live() const { return bytes_live_; }
    size_t bytes_pooled() const { return bytes_pooled_; }

private:
    static size_t bucket_size(size_t nbytes);

    MTL::Device* device_; // borrowed, owned by forge::Device
    std::mutex mutex_;
    std::unordered_map<size_t, std::vector<MTL::Buffer*>> pool_;
    std::vector<MTL::Buffer*> retired_;
    bool defer_ = false;
    size_t bytes_live_ = 0;
    size_t bytes_pooled_ = 0;
};

} // namespace forge