import Metal

// ══════════════════════════════════════════════════════════════════
// Batch Layer Processing - TRUE parallel layer forward pass
// Process multiple tokens through entire layer simultaneously
// ══════════════════════════════════════════════════════════════════

extension E4BLayer {
    
    /// Batch forward pass - process N tokens through entire layer in parallel
    /// Expected: 8-15x speedup for batch inference
    public func forwardBatchTrue(
        batchInput: MTLBuffer,          // [batchSize, hiddenSize]
        positions: [Int],
        batchSize: Int,
        kvCache: KVCache,
        shouldStoreKV: Bool,
        temps: ForwardTemps,
        batchTemps: BatchTemps,         // Batch-specific buffers
        engine: MarkBaseEngine,
        cmdBuf: MTLCommandBuffer
    ) throws {
        // Note: This is a simplified implementation focusing on FFN batch processing
        // Attention still needs sequential KV cache updates
        
        // ── Phase 1: Batch Input Norm ──
        guard let inputLN = inputLayernorm else {
            throw NSError(domain: "LayerBatch", code: -1,
                         userInfo: [NSLocalizedDescriptionKey: "inputLayernorm required for batch processing"])
        }
        
        try batchLayerRMSNorm(
            batchInput: batchInput,
            weights: inputLN,
            batchOutput: batchTemps.hBatch,
            hiddenSize: config.hiddenSize,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // ── Phase 2: Batch Attention (Sequential for KV cache) ──
        // Note: Attention needs per-token KV cache updates, so we process sequentially
        // But we can batch Q/K/V projections
        
        try batchQuantizedMatmul(
            batchInput: batchTemps.hBatch,
            weights: qProj,
            batchOutput: batchTemps.qBatch,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // Batch grouped norm for Q
        guard let qN = qNorm else {
            throw NSError(domain: "LayerBatch", code: -2,
                         userInfo: [NSLocalizedDescriptionKey: "qNorm required for batch processing"])
        }
        
        try batchGroupedRMSNorm(
            batchInput: batchTemps.qBatch,
            weights: qN,
            batchOutput: batchTemps.nsBatch,
            count: config.nHeads * config.headDim,
            groupSize: config.headDim,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // Sequential RoPE and attention (KV cache dependency)
        for i in 0..<batchSize {
            let pos = positions[i]
            let offset = i * config.nHeads * config.headDim
            
            // Get Q for this token
            let qToken = engine.device.makeBuffer(
                bytes: batchTemps.nsBatch.contents() + offset * 4,
                length: config.nHeads * config.headDim * 4,
                options: .storageModeShared
            )!
            
            // Apply RoPE
            try applyRoPEQ(engine: engine, cmdBuf: cmdBuf, q: qToken, position: pos)
            
            // K/V projections (batched, but we need per-token results)
            let hToken = engine.device.makeBuffer(
                bytes: batchTemps.hBatch.contents() + i * config.hiddenSize * 4,
                length: config.hiddenSize * 4,
                options: .storageModeShared
            )!
            
            try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: hToken, weights: kProj, output: temps.k)
            if let vp = vProj {
                try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: hToken, weights: vp, output: temps.v)
            }
            
            // K/V norms
            try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf, input: temps.k, weight: kNorm, output: temps.up,
                               count: config.nKvHeads * config.headDim, groupSize: config.headDim, eps: rmsNormEps)
            
            // RoPE K
            try applyRoPEK(engine: engine, cmdBuf: cmdBuf, k: temps.up, position: pos)
            
            // Store KV
            if shouldStoreKV {
                let valueBuf = vNorm != nil ? temps.gate : temps.v
                kvCache.store(key: temps.up, keySrcOffset: 0, value: valueBuf, valueSrcOffset: 0,
                              position: pos, commandBuffer: cmdBuf)
            }
            
            // Attention
            let curK = temps.up
            let curV = vNorm != nil ? temps.gate : temps.v
            if config.isSliding {
                if shouldStoreKV {
                    try slidingAttention(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache, position: pos)
                } else {
                    try slidingAttentionWithCurrent(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache,
                                                     curK: curK, curV: curV, position: pos)
                }
            } else {
                if shouldStoreKV {
                    try fullAttention(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache, position: pos)
                } else {
                    try fullAttentionWithCurrent(engine: engine, cmdBuf: cmdBuf, q: qToken, cache: kvCache,
                                                  curK: curK, curV: curV, position: pos)
                }
            }
            
            // O projection (write back to batch buffer)
            try quantizedMatmul(engine: engine, cmdBuf: cmdBuf, input: temps.attn, weights: oProj, output: temps.h)
            
            // Copy to batch position
            let batchOffset = i * config.hiddenSize * 4
            memcpy(batchInput.contents() + batchOffset, temps.h.contents(), config.hiddenSize * 4)
        }
        
        // ── Phase 3: Batch FFN (TRUE batch processing) ──
        // This is where we get the big speedup
        
        // Post-attention norm (batched)
        guard let postAttnLN = postAttentionLayernorm else {
            throw NSError(domain: "LayerBatch", code: -3,
                         userInfo: [NSLocalizedDescriptionKey: "postAttentionLayernorm required"])
        }
        
        try batchLayerRMSNorm(
            batchInput: batchInput,
            weights: postAttnLN,
            batchOutput: batchTemps.hBatch,
            hiddenSize: config.hiddenSize,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // Pre-FFN norm (batched)
        guard let preFFNLN = preFeedforwardLayernorm else {
            throw NSError(domain: "LayerBatch", code: -4,
                         userInfo: [NSLocalizedDescriptionKey: "preFeedforwardLayernorm required"])
        }
        
        try batchLayerRMSNorm(
            batchInput: batchTemps.hBatch,
            weights: preFFNLN,
            batchOutput: batchTemps.nsBatch,
            hiddenSize: config.hiddenSize,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // Batch FFN: Gate + Up (fused)
        try batchFusedGateUp(
            batchInput: batchTemps.nsBatch,
            gateWeights: gateProj,
            upWeights: upProj,
            batchOutput: batchTemps.interBatch,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // Batch Down projection
        try batchDownProjection(
            batchInter: batchTemps.interBatch,
            downWeights: downProj,
            batchOutput: batchTemps.hBatch,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // Batch residual add
        try batchEltwiseAdd(
            batchA: batchInput,
            batchB: batchTemps.hBatch,
            batchOutput: batchInput,
            size: config.hiddenSize,
            batchSize: batchSize,
            cmdBuf: cmdBuf,
            engine: engine
        )
        
        // Layer scalar (if needed)
        if layerScalar != 1.0 {
            try batchScaleBuffer(
                batchBuffer: batchInput,
                scale: layerScalar,
                size: config.hiddenSize,
                batchSize: batchSize,
                cmdBuf: cmdBuf,
                engine: engine
            )
        }
    }
    
    // ── Batch Layer Helper Functions ──
    
    private func batchLayerRMSNorm(
        batchInput: MTLBuffer,
        weights: MTLBuffer,
        batchOutput: MTLBuffer,
        hiddenSize: Int,
        batchSize: Int,
        cmdBuf: MTLCommandBuffer,
        engine: MarkBaseEngine
    ) throws {
        let pso = try engine.pipeline(named: "batch_layer_rms_norm")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        
        enc.setBuffer(batchInput, offset: 0, index: 0)
        enc.setBuffer(weights, offset: 0, index: 1)
        enc.setBuffer(batchOutput, offset: 0, index: 2)
        
        var hs = UInt32(hiddenSize)
        enc.setBytes(&hs, length: 4, index: 3)
        var eps: Float = rmsNormEps
        enc.setBytes(&eps, length: 4, index: 4)
        var batch = UInt32(batchSize)
        enc.setBytes(&batch, length: 4, index: 5)
        
        let tg = MTLSize(width: 256, height: 1, depth: 1)
        let grid = MTLSize(width: batchSize, height: hiddenSize, depth: 1)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
    
    private func batchQuantizedMatmul(
        batchInput: MTLBuffer,
        weights: QuantizedWeights,
        batchOutput: MTLBuffer,
        batchSize: Int,
        cmdBuf: MTLCommandBuffer,
        engine: MarkBaseEngine
    ) throws {
        let pso = try engine.pipeline(named: "batch_layer_quantized_matmul")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        
        enc.setBuffer(batchInput, offset: 0, index: 0)
        enc.setBuffer(weights.weight, offset: 0, index: 1)
        enc.setBuffer(weights.scales, offset: 0, index: 2)
        enc.setBuffer(weights.biases, offset: 0, index: 3)
        enc.setBuffer(batchOutput, offset: 0, index: 4)
        
        var inDim = UInt32(weights.inDim)
        enc.setBytes(&inDim, length: 4, index: 5)
        var outDim = UInt32(weights.outDim)
        enc.setBytes(&outDim, length: 4, index: 6)
        var groupSize = UInt32(weights.groupSize)
        enc.setBytes(&groupSize, length: 4, index: 7)
        var batch = UInt32(batchSize)
        enc.setBytes(&batch, length: 4, index: 8)
        
let tg = MTLSize(width: 256, height: 1, depth: 1)
        let grid = MTLSize(width: batchSize, height: Int(weights.outDim), depth: 1)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
    
    private func batchGroupedRMSNorm(
        batchInput: MTLBuffer,
        weights: MTLBuffer,
        batchOutput: MTLBuffer,
        count: Int,
        groupSize: Int,
        batchSize: Int,
        cmdBuf: MTLCommandBuffer,
        engine: MarkBaseEngine
    ) throws {
        // Use existing grouped_rms_norm kernel with batch iteration
        // For now, process sequentially (can optimize later)
        let inputPtr = batchInput.contents().assumingMemoryBound(to: Float.self)
        let outputPtr = batchOutput.contents().assumingMemoryBound(to: Float.self)
        
        for i in 0..<batchSize {
            let offset = i * count
            let tokenInput = engine.device.makeBuffer(
                bytes: inputPtr + offset,
                length: count * 4,
                options: .storageModeShared
            )!
            let tokenOutput = engine.device.makeBuffer(
                bytes: outputPtr + offset,
                length: count * 4,
                options: .storageModeShared
            )!
            
            try groupedRmsNorm(engine: engine, cmdBuf: cmdBuf, input: tokenInput, weight: weights,
                               output: tokenOutput, count: count, groupSize: groupSize, eps: rmsNormEps)
        }
    }
    
    private func batchFusedGateUp(
        batchInput: MTLBuffer,
        gateWeights: QuantizedWeights,
        upWeights: QuantizedWeights,
        batchOutput: MTLBuffer,
        batchSize: Int,
        cmdBuf: MTLCommandBuffer,
        engine: MarkBaseEngine
    ) throws {
        let pso = try engine.pipeline(named: "batch_fused_gate_up")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        
        enc.setBuffer(batchInput, offset: 0, index: 0)
        enc.setBuffer(gateWeights.weight, offset: 0, index: 1)
        enc.setBuffer(gateWeights.scales, offset: 0, index: 2)
        enc.setBuffer(gateWeights.biases, offset: 0, index: 3)
        enc.setBuffer(upWeights.weight, offset: 0, index: 4)
        enc.setBuffer(upWeights.scales, offset: 0, index: 5)
        enc.setBuffer(upWeights.biases, offset: 0, index: 6)
        enc.setBuffer(batchOutput, offset: 0, index: 7)
        
        var hiddenSize = UInt32(gateWeights.inDim)
        enc.setBytes(&hiddenSize, length: 4, index: 8)
        var intermediateSize = UInt32(gateWeights.outDim)
        enc.setBytes(&intermediateSize, length: 4, index: 9)
        var groupSize = UInt32(gateWeights.groupSize)
        enc.setBytes(&groupSize, length: 4, index: 10)
        var batch = UInt32(batchSize)
        enc.setBytes(&batch, length: 4, index: 11)
        
        let tg = MTLSize(width: 256, height: 1, depth: 1)
        let grid = MTLSize(width: batchSize, height: Int(gateWeights.outDim), depth: 1)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
    
    private func batchDownProjection(
        batchInter: MTLBuffer,
        downWeights: QuantizedWeights,
        batchOutput: MTLBuffer,
        batchSize: Int,
        cmdBuf: MTLCommandBuffer,
        engine: MarkBaseEngine
    ) throws {
        let pso = try engine.pipeline(named: "batch_down_projection")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        
        enc.setBuffer(batchInter, offset: 0, index: 0)
        enc.setBuffer(downWeights.weight, offset: 0, index: 1)
        enc.setBuffer(downWeights.scales, offset: 0, index: 2)
        enc.setBuffer(downWeights.biases, offset: 0, index: 3)
        enc.setBuffer(batchOutput, offset: 0, index: 4)
        
        var hiddenSize = UInt32(downWeights.outDim)
        enc.setBytes(&hiddenSize, length: 4, index: 5)
        var intermediateSize = UInt32(downWeights.inDim)
        enc.setBytes(&intermediateSize, length: 4, index: 6)
        var groupSize = UInt32(downWeights.groupSize)
        enc.setBytes(&groupSize, length: 4, index: 7)
        var batch = UInt32(batchSize)
        enc.setBytes(&batch, length: 4, index: 8)
        
        let tg = MTLSize(width: 256, height: 1, depth: 1)
        let grid = MTLSize(width: batchSize, height: Int(downWeights.outDim), depth: 1)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
    
    private func batchEltwiseAdd(
        batchA: MTLBuffer,
        batchB: MTLBuffer,
        batchOutput: MTLBuffer,
        size: Int,
        batchSize: Int,
        cmdBuf: MTLCommandBuffer,
        engine: MarkBaseEngine
    ) throws {
        let pso = try engine.pipeline(named: "batch_eltwise_add")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        
        enc.setBuffer(batchA, offset: 0, index: 0)
        enc.setBuffer(batchB, offset: 0, index: 1)
        enc.setBuffer(batchOutput, offset: 0, index: 2)
        
        var s = UInt32(size)
        enc.setBytes(&s, length: 4, index: 3)
        var batch = UInt32(batchSize)
        enc.setBytes(&batch, length: 4, index: 4)
        
        let tg = MTLSize(width: 256, height: 1, depth: 1)
        let grid = MTLSize(width: batchSize * size, height: 1, depth: 1)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
    
    private func batchScaleBuffer(
        batchBuffer: MTLBuffer,
        scale: Float,
        size: Int,
        batchSize: Int,
        cmdBuf: MTLCommandBuffer,
        engine: MarkBaseEngine
    ) throws {
        let pso = try engine.pipeline(named: "eltwise_scale")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        
        enc.setBuffer(batchBuffer, offset: 0, index: 0)
        var s = scale
        enc.setBytes(&s, length: 4, index: 1)
        var count = UInt32(batchSize * size)
        enc.setBytes(&count, length: 4, index: 2)
        
        let tg = MTLSize(width: 256, height: 1, depth: 1)
        let grid = MTLSize(width: batchSize * size, height: 1, depth: 1)
        enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
}