// Author: Simon-Pierre Boucher — contact@spboucher.ai
#include "core/allocator.h"

#include <Metal/Metal.hpp>

#include <cstdio>
#include <cstdlib>

namespace forge {

Allocator::Allocator(MTL::Device* device) : device_(device) {}

Allocator::~Allocator() {
    trim();
}

size_t Allocator::bucket_size(size_t nbytes) {
    size_t b = 256;
    while (b < nbytes) b <<= 1;
    return b;
}

MTL::Buffer* Allocator::acquire(size_t nbytes) {
    const size_t bucket = bucket_size(nbytes);
    {
        std::lock_guard<std::mutex> lock(mutex_);
        auto it = pool_.find(bucket);
        if (it != pool_.end() && !it->second.empty()) {
            MTL::Buffer* buf = it->second.back();
            it->second.pop_back();
            bytes_pooled_ -= bucket;
            bytes_live_ += bucket;
            return buf;
        }
    }
    MTL::Buffer* buf = device_->newBuffer(bucket, MTL::ResourceStorageModeShared);
    if (!buf) {
        std::fprintf(stderr, "forge: Metal buffer allocation failed (%zu bytes)\n", bucket);
        std::abort();
    }
    std::lock_guard<std::mutex> lock(mutex_);
    bytes_live_ += bucket;
    return buf;
}

void Allocator::release(MTL::Buffer* buffer) {
    if (!buffer) return;
    const size_t bucket = buffer->length();
    std::lock_guard<std::mutex> lock(mutex_);
    if (defer_) {
        retired_.push_back(buffer);
    } else {
        pool_[bucket].push_back(buffer);
        bytes_pooled_ += bucket;
    }
    bytes_live_ -= bucket;
}

void Allocator::set_defer(bool defer) {
    std::lock_guard<std::mutex> lock(mutex_);
    defer_ = defer;
}

void Allocator::flush_retired() {
    std::lock_guard<std::mutex> lock(mutex_);
    for (MTL::Buffer* buf : retired_) {
        pool_[buf->length()].push_back(buf);
        bytes_pooled_ += buf->length();
    }
    retired_.clear();
}

void Allocator::trim() {
    std::lock_guard<std::mutex> lock(mutex_);
    for (auto& [bucket, buffers] : pool_) {
        for (MTL::Buffer* buf : buffers) buf->release();
    }
    pool_.clear();
    bytes_pooled_ = 0;
}

} // namespace forge