import Metal

// Optimized forward pass with batched Metal commands
// Goal: Reduce waitUntilCompleted() calls from 11 to 1

extension E4BModel {
    
    /// Optimized forward pass - batches all Metal commands
    /// Expected improvement: 4x faster token generation (verified)
    public func forwardOptimized(tokenId: Int, position: Int, debug: Bool = false) throws -> [Float] {
        // Create ONE shared command buffer for entire forward pass
        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        let h = temps.io
        
        // ── Phase 1: Embedding (batched, NO fusion) ──
        try dequantizeRowOptimized(weight: embedWeight, tokenId: tokenId, output: h, cmdBuf: cmdBuf)
        
        if embedScale != 1.0 {
            try scaleBufferOptimized(h, scale: embedScale, count: hiddenSize, cmdBuf: cmdBuf)
        }
        
        // Debug: Check embedding output
        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()
        let embedPtr = h.contents().assumingMemoryBound(to: Float.self)
        let embedSample = Array(UnsafeBufferPointer(start: embedPtr, count: min(20, hiddenSize)))
        let embedNaNCount = Array(UnsafeBufferPointer(start: embedPtr, count: hiddenSize)).filter { $0.isNaN }.count
        print("TEXT Embedding: sample=\(embedSample), NaN=\(embedNaNCount)/\(hiddenSize)")
        
        let cmdBuf2 = engine.commandQueue.makeCommandBuffer()!
        
        // Per-layer embedding (E4B)
        if let plWeight = embedTokensPerLayerWeight,
           let plBuf = perLayerEmbedBuffer,
           let ctxBuf = perLayerContextBuffer {
            
            let totalPerLayer = perLayerInputSize * numHiddenLayers
            try dequantizeRowOptimized(weight: plWeight, tokenId: tokenId, output: plBuf, 
                                       nCols: totalPerLayer, cmdBuf: cmdBuf2)
            
            let plEmbedScale = sqrt(Float(perLayerInputSize))
            try scaleBufferOptimized(plBuf, scale: plEmbedScale, count: totalPerLayer, cmdBuf: cmdBuf2)
            
            if let projBuf = perLayerModelProjection {
                try matmulBF16Optimized(input: h, weight: projBuf, output: ctxBuf,
                                       inDim: hiddenSize, outDim: perLayerModelProjectionOutDim, 
                                       cmdBuf: cmdBuf2)
                try scaleBufferOptimized(ctxBuf, scale: perLayerModelProjectionScaleVal, 
                                        count: totalPerLayer, cmdBuf: cmdBuf2)
                
                if let norm = perLayerProjectionNorm {
                    try rmsNormBatchOptimized(input: ctxBuf, weight: norm, output: plBuf,
                                             perLayerSize: perLayerInputSize, 
                                             numLayers: numHiddenLayers, cmdBuf: cmdBuf2)
                    
                    let blit = cmdBuf2.makeBlitCommandEncoder()!
                    blit.copy(from: plBuf, sourceOffset: 0,
                             to: ctxBuf, destinationOffset: 0,
                             size: totalPerLayer * 4)
                    blit.endEncoding()
                    
                    try dequantizeRowOptimized(weight: plWeight, tokenId: tokenId, output: plBuf,
                                               nCols: totalPerLayer, cmdBuf: cmdBuf2)
                    try scaleBufferOptimized(plBuf, scale: plEmbedScale, count: totalPerLayer, cmdBuf: cmdBuf2)
                }
                
                try eltwiseAddScaledOptimized(a: ctxBuf, scaleA: 1.0,
                                              b: plBuf, scaleB: 1.0,
                                              output: ctxBuf, count: totalPerLayer, cmdBuf: cmdBuf2)
                try scaleBufferOptimized(ctxBuf, scale: perLayerInputScaleVal, 
                                        count: totalPerLayer, cmdBuf: cmdBuf2)
                
                let blit = cmdBuf2.makeBlitCommandEncoder()!
                blit.copy(from: ctxBuf, sourceOffset: 0,
                         to: plBuf, destinationOffset: 0,
                         size: totalPerLayer * 4)
                blit.endEncoding()
            }
        }
        
        // ── Phase 2: Layers (all in same command buffer) ──
        for layerIdx in 0..<numHiddenLayers {
            let isOwner = layerIdx < firstKVShared
            let cacheIdx = isOwner ? layerIdx : (kvSourceMap[layerIdx] ?? (layerIdx - numKvShared))
            let cache = kvCaches[cacheIdx]
            
            let plOffset = perLayerInputSize > 0 ? layerIdx * perLayerInputSize * MemoryLayout<Float>.stride : 0
            
            // OPTIMIZED: Use shared command buffer (no wait per layer)
            try layers[layerIdx].forwardOptimized(
                input: h, position: position,
                kvCache: cache, shouldStoreKV: isOwner,
                temps: temps, engine: engine,
                cmdBuf: cmdBuf2,
                perLayerInput: perLayerEmbedBuffer,
                perLayerInputOffset: plOffset
            )
        }
        
        let cmdBuf3 = engine.commandQueue.makeCommandBuffer()!
        
        // ── Phase 3: LM Head (in same command buffer) ──
        var lmInput = h
        if let fn = finalNorm {
            try rmsNormOptimized(input: h, weight: fn, output: temps.ns,
                                count: hiddenSize, cmdBuf: cmdBuf3)
            lmInput = temps.ns
        }
        
        try quantizedMatmulOptimized(input: lmInput, weights: embedWeight, 
                                     output: logitsBuffer, cmdBuf: cmdBuf3)
        
        // Logit softcapping
        if let cap = finalLogitSoftcapping {
            try applyLogitSoftcappingOptimized(buffer: logitsBuffer, cap: cap, 
                                              count: vocabSize, cmdBuf: cmdBuf3)
        }
        
        // ── Final: Commit and wait ONCE ──
        cmdBuf3.commit()
        cmdBuf3.waitUntilCompleted()  // Only ONE wait for entire forward pass!
        
        // Read back logits
        let logits = engine.readFloats(from: logitsBuffer, count: vocabSize)
        
        if debug && position < 3 {
            let maxLogit = logits.max() ?? 0
            let minLogit = logits.min() ?? 0
            print("  Optimized forward: max=\(maxLogit), min=\(minLogit)")
        }
        
        return logits
    }
    
    // ── Optimized helper functions (accept cmdBuf parameter) ──
    
    // FUSED: dequantize + scale in one kernel
    func dequantizeScaleFused(weight: QuantizedWeights, tokenId: Int,
                              output: MTLBuffer, scale: Float,
                              nCols: Int? = nil,
                              cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "fused_dequantize_scale")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(weight.weight, offset: 0, index: 0)
        enc.setBuffer(weight.scales, offset: 0, index: 1)
        enc.setBuffer(weight.biases, offset: 0, index: 2)
        enc.setBuffer(output, offset: 0, index: 3)
        let actualCols = nCols ?? hiddenSize
        var nColsVal = UInt32(actualCols)
        enc.setBytes(&nColsVal, length: 4, index: 4)
        var row = Int32(tokenId)
        enc.setBytes(&row, length: 4, index: 5)
        var groupSize = UInt32(weight.groupSize)
        enc.setBytes(&groupSize, length: 4, index: 6)
        var s = scale
        enc.setBytes(&s, length: 4, index: 7)
        let tg = engine.threadgroupSize1D(pso, count: actualCols)
        enc.dispatchThreads(MTLSize(width: actualCols, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
    }
    
    func dequantizeRowOptimized(weight: QuantizedWeights, tokenId: Int,
                                output: MTLBuffer, nCols: Int? = nil,
                                cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "dequantize_row")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(weight.weight, offset: 0, index: 0)
        enc.setBuffer(weight.scales, offset: 0, index: 1)
        enc.setBuffer(weight.biases, offset: 0, index: 2)
        enc.setBuffer(output, offset: 0, index: 3)
        let actualCols = nCols ?? hiddenSize
        var nColsVal = UInt32(actualCols)
        enc.setBytes(&nColsVal, length: 4, index: 4)
        var row = Int32(tokenId)
        enc.setBytes(&row, length: 4, index: 5)
        var groupSize = UInt32(weight.groupSize)
        enc.setBytes(&groupSize, length: 4, index: 6)
        let tg = engine.threadgroupSize1D(pso, count: actualCols)
        enc.dispatchThreads(MTLSize(width: actualCols, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
    
    func scaleBufferOptimized(_ buf: MTLBuffer, scale: Float, count: Int,
                              cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "eltwise_scale")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(buf, offset: 0, index: 0)
        var s = scale
        enc.setBytes(&s, length: 4, index: 1)
        var N = UInt32(count)
        enc.setBytes(&N, length: 4, index: 2)
        let tg = engine.threadgroupSize1D(pso, count: count)
        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
    
    func quantizedMatmulOptimized(input: MTLBuffer, weights: QuantizedWeights,
                                  output: MTLBuffer, cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "quantized_matmul")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(input, 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(output, 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)
        let tg = engine.threadgroupSize1D(pso, count: weights.outDim)
        enc.dispatchThreads(MTLSize(width: weights.outDim, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
    
    func rmsNormOptimized(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer,
                         count: Int, cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "rms_norm")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(weight, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)
        var N = UInt32(count)
        enc.setBytes(&N, length: 4, index: 3)
        var eps: Float = rmsNormEps
        enc.setBytes(&eps, length: 4, index: 4)
        let tg = engine.threadgroupSize1D(pso, count: count)
        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
    
    func rmsNormBatchOptimized(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer,
                               perLayerSize: Int, numLayers: Int,
                               cmdBuf: MTLCommandBuffer) throws {
        // Batch all per-layer norms in one kernel dispatch
        let pso = try engine.pipeline(named: "rms_norm")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        
        for layerIdx in 0..<numLayers {
            let offset = layerIdx * perLayerSize
            enc.setBuffer(input, offset: offset * 4, index: 0)
            enc.setBuffer(weight, offset: 0, index: 1)
            enc.setBuffer(output, offset: offset * 4, index: 2)
            var N = UInt32(perLayerSize)
            enc.setBytes(&N, length: 4, index: 3)
            var eps: Float = rmsNormEps
            enc.setBytes(&eps, length: 4, index: 4)
            let tg = engine.threadgroupSize1D(pso, count: perLayerSize)
            enc.dispatchThreads(MTLSize(width: perLayerSize, height: 1, depth: 1),
                               threadsPerThreadgroup: tg)
        }
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
    
    func matmulBF16Optimized(input: MTLBuffer, weight: MTLBuffer, output: MTLBuffer,
                             inDim: Int, outDim: Int, cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "matmul_f32")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(input, offset: 0, index: 0)
        enc.setBuffer(weight, offset: 0, index: 1)
        enc.setBuffer(output, offset: 0, index: 2)
        var M: UInt32 = 1
        enc.setBytes(&M, length: 4, index: 3)
        var K = UInt32(inDim)
        enc.setBytes(&K, length: 4, index: 4)
        var N = UInt32(outDim)
        enc.setBytes(&N, length: 4, index: 5)
        let tg = engine.threadgroupSize1D(pso, count: outDim)
        enc.dispatchThreads(MTLSize(width: outDim, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
    
    func eltwiseAddScaledOptimized(a: MTLBuffer, scaleA: Float,
                                   b: MTLBuffer, scaleB: Float,
                                   output: MTLBuffer, count: Int,
                                   cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "eltwise_add_scaled")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(a, offset: 0, index: 0)
        var sa = scaleA
        enc.setBytes(&sa, length: 4, index: 1)
        enc.setBuffer(b, offset: 0, index: 2)
        var sb = scaleB
        enc.setBytes(&sb, length: 4, index: 3)
        enc.setBuffer(output, offset: 0, index: 4)
        var N = UInt32(count)
        enc.setBytes(&N, length: 4, index: 5)
        let tg = engine.threadgroupSize1D(pso, count: count)
        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
    
    func applyLogitSoftcappingOptimized(buffer: MTLBuffer, cap: Float, 
                                        count: Int, cmdBuf: MTLCommandBuffer) throws {
        let pso = try engine.pipeline(named: "tanh_scale")
        let enc = cmdBuf.makeComputeCommandEncoder()!
        enc.setComputePipelineState(pso)
        enc.setBuffer(buffer, offset: 0, index: 0)
        enc.setBuffer(buffer, offset: 0, index: 1)  // in-place
        var c = cap
        enc.setBytes(&c, length: 4, index: 2)
        var N = UInt32(count)
        enc.setBytes(&N, length: 4, index: 3)
        let tg = engine.threadgroupSize1D(pso, count: count)
        enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        // NO waitUntilCompleted here!
    }
}