import Metal

/// Buffer pool for reusing MTLBuffers to reduce allocation overhead
/// 
/// Metal buffer allocation can be expensive, especially when done frequently
/// during inference. This pool caches and reuses buffers of common sizes.
public final class BufferPool: @unchecked Sendable {
    private let device: MTLDevice
    private var availableBuffers: [Int: [MTLBuffer]] = [:]  // size -> [buffers]
    private let lock = NSLock()
    
    /// Statistics
    public private(set) var totalAllocations: Int = 0
    public private(set) var totalReuses: Int = 0
    public private(set) var peakBufferCount: Int = 0
    
    public init(device: MTLDevice) {
        self.device = device
    }
    
    /// Acquire a buffer of the specified size
    /// Returns a reusable buffer if available, otherwise creates a new one
    public func acquire(length: Int) -> MTLBuffer {
        lock.lock()
        defer { lock.unlock() }
        
        // Round up to nearest 256 bytes for alignment
        let alignedLength = (length + 255) & ~255
        
        // Check for available buffer
        if var buffers = availableBuffers[alignedLength], !buffers.isEmpty {
            let buffer = buffers.removeLast()
            availableBuffers[alignedLength] = buffers
            totalReuses += 1
            return buffer
        }
        
        // Create new buffer
        totalAllocations += 1
        guard let buffer = device.makeBuffer(
            length: alignedLength,
            options: .storageModeShared
        ) else {
            fatalError("Failed to allocate Metal buffer of size \(alignedLength)")
        }
        
        // Track peak
        let currentCount = totalAllocations - totalReuses
        if currentCount > peakBufferCount {
            peakBufferCount = currentCount
        }
        
        return buffer
    }
    
    /// Release a buffer back to the pool for reuse
    public func release(_ buffer: MTLBuffer) {
        lock.lock()
        defer { lock.unlock() }
        
        let length = buffer.length
        if availableBuffers[length] == nil {
            availableBuffers[length] = []
        }
        availableBuffers[length]?.append(buffer)
    }
    
    /// Clear all cached buffers (useful for memory pressure situations)
    public func clear() {
        lock.lock()
        defer { lock.unlock() }
        availableBuffers.removeAll()
    }
    
    /// Get pool statistics
    public var stats: String {
        lock.lock()
        defer { lock.unlock() }
        let totalBuffers = availableBuffers.values.reduce(0) { $0 + $1.count }
        return """
        BufferPool Stats:
          Allocations: \(totalAllocations)
          Reuses: \(totalReuses)
          Available: \(totalBuffers)
          Peak: \(peakBufferCount)
          Hit Rate: \(totalAllocations > 0 ? String(format: "%.1f%%", Float(totalReuses) / Float(totalAllocations) * 100) : "0%")
        """
    }
}