import Metal

/// Per-layer KV cache supporting sliding window (rotating) and full (growing) attention.
public final class KVCache {
    let isSliding: Bool
    let maxLength: Int
    let nKvHeads: Int
    let headDim: Int
    let buffer: MTLBuffer  // contiguous K then V: [2 * maxLength * nKvHeads * headDim]
    private(set) var currentLength: Int = 0

    init(device: MTLDevice, isSliding: Bool, maxContextLength: Int, nKvHeads: Int, headDim: Int) {
        self.isSliding = isSliding
        self.maxLength = isSliding ? 512 : maxContextLength
        self.nKvHeads = nKvHeads
        self.headDim = headDim

        let perStep = nKvHeads * headDim * MemoryLayout<Float>.stride
        let total = 2 * self.maxLength * perStep
        self.buffer = device.makeBuffer(length: total, options: .storageModeShared)!
    }

    var effectiveLength: Int {
        isSliding ? min(currentLength, 512) : currentLength
    }

    /// Key buffer start offset (in bytes, from buffer start)
    var keyBaseOffset: Int { 0 }

    /// Value buffer start offset (in bytes, from buffer start) — immediately after K
    var valueBaseOffset: Int {
        maxLength * nKvHeads * headDim * MemoryLayout<Float>.stride
    }

    /// Byte offset for a given logical position in the key region.
    func keyOffset(for position: Int) -> Int {
        let p = isSliding ? (position % maxLength) : position
        return p * nKvHeads * headDim * MemoryLayout<Float>.stride
    }

    func valueOffset(for position: Int) -> Int {
        valueBaseOffset + keyOffset(for: position)
    }

    /// Store K,V into cache at the given logical position.
    func store(key: MTLBuffer, keySrcOffset: Int,
               value: MTLBuffer, valueSrcOffset: Int,
               position: Int,
               commandBuffer: MTLCommandBuffer) {
        let stepBytes = nKvHeads * headDim * MemoryLayout<Float>.stride
        currentLength = max(currentLength, position + 1)

        let blit = commandBuffer.makeBlitCommandEncoder()!
        blit.copy(from: key, sourceOffset: keySrcOffset,
                  to: buffer, destinationOffset: keyOffset(for: position),
                  size: stepBytes)
        blit.copy(from: value, sourceOffset: valueSrcOffset,
                  to: buffer, destinationOffset: valueOffset(for: position),
                  size: stepBytes)
        blit.endEncoding()
    }
    
    /// Reset cache for new sequence
    func reset() {
        currentLength = 0
    }
}