import XCTest
@testable import MarkBase

final class BatchKernelTest: XCTestCase {
    
    func testBatchKernelCompilation() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Batch Metal Kernel Compilation Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Test batch kernel compilation
        print("Testing batch kernel compilation...")
        
        do {
            let pso1 = try engine.pipeline(named: "quantized_matmul_batch")
            print("  ✓ quantized_matmul_batch: compiled")
            print("    Threadgroup size: \(pso1.maxTotalThreadsPerThreadgroup)")
        } catch {
            print("  ✗ quantized_matmul_batch: NOT FOUND")
            print("    Error: \(error)")
        }
        
        do {
            let pso2 = try engine.pipeline(named: "rms_norm_batch")
            print("  ✓ rms_norm_batch: compiled")
            print("    Threadgroup size: \(pso2.maxTotalThreadsPerThreadgroup)")
        } catch {
            print("  ✗ rms_norm_batch: NOT FOUND")
            print("    Error: \(error)")
        }
        
        do {
            let pso3 = try engine.pipeline(named: "sliding_attention_batch")
            print("  ✓ sliding_attention_batch: compiled")
            print("    Threadgroup size: \(pso3.maxTotalThreadsPerThreadgroup)")
        } catch {
            print("  ✗ sliding_attention_batch: NOT FOUND")
            print("    Error: \(error)")
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("Batch kernel compilation test complete")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    func testBatchMatmulSimple() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Simple Batch Matmul Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let device = engine.device
        
        // Create simple test data
        let batchSize = 2
        let inDim = 256
        let outDim = 512
        
        // Input: [2, 256]
        let inputs = device.makeBuffer(length: batchSize * inDim * 4)!
        let inputPtr = inputs.contents().assumingMemoryBound(to: Float.self)
        for i in 0..<batchSize * inDim {
            inputPtr[i] = Float(i) / 100.0
        }
        
        // Output: [2, 512]
        let outputs = device.makeBuffer(length: batchSize * outDim * 4)!
        
        // Simple identity weights (for testing)
        // Note: This is NOT a real quantized weight test, just kernel validation
        let weights = device.makeBuffer(length: outDim * inDim)!
        let scales = device.makeBuffer(length: outDim * (inDim / 64) * 4)!
        let biases = device.makeBuffer(length: outDim * 4)!
        
        // Initialize weights to identity pattern
        let weightPtr = weights.contents().assumingMemoryBound(to: UInt8.self)
        for i in 0..<outDim * inDim {
            weightPtr[i] = 128  // Zero in quantized space
        }
        
        let scalePtr = scales.contents().assumingMemoryBound(to: Float.self)
        for i in 0..<outDim * (inDim / 64) {
            scalePtr[i] = 1.0
        }
        
        let biasPtr = biases.contents().assumingMemoryBound(to: Float.self)
        for i in 0..<outDim {
            biasPtr[i] = 0.0
        }
        
        // Try to run batch matmul kernel
        do {
            let pso = try engine.pipeline(named: "quantized_matmul_batch")
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            let enc = cmdBuf.makeComputeCommandEncoder()!
            enc.setComputePipelineState(pso)
            
            enc.setBuffer(inputs, offset: 0, index: 0)
            enc.setBuffer(weights, offset: 0, index: 1)
            enc.setBuffer(scales, offset: 0, index: 2)
            enc.setBuffer(biases, offset: 0, index: 3)
            enc.setBuffer(outputs, offset: 0, index: 4)
            
            var inDimVal = UInt32(inDim)
            enc.setBytes(&inDimVal, length: 4, index: 5)
            var outDimVal = UInt32(outDim)
            enc.setBytes(&outDimVal, length: 4, index: 6)
            var groupSize = UInt32(64)
            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: outDim, depth: 1)
            enc.dispatchThreads(grid, threadsPerThreadgroup: tg)
            enc.endEncoding()
            
            cmdBuf.commit()
            cmdBuf.waitUntilCompleted()
            
            // Check outputs
            let outputPtr = outputs.contents().assumingMemoryBound(to: Float.self)
            print("Output values (first 10):")
            for i in 0..<10 {
                print("  outputs[\(i)] = \(outputPtr[i])")
            }
            
            // Check for NaN
            let hasNaN = (0..<batchSize * outDim).contains { i in
                outputPtr[i].isNaN || outputPtr[i].isInfinite
            }
            
            if hasNaN {
                print("  ✗ Output has NaN or Inf!")
            } else {
                print("  ✓ Output is valid (no NaN)")
            }
            
            print("\n═══════════════════════════════════════════════════════════════════")
            print("Batch matmul kernel works!")
            print("═══════════════════════════════════════════════════════════════════\n")
            
        } catch {
            print("  ✗ Kernel execution failed: \(error)")
            print("\n═══════════════════════════════════════════════════════════════════")
            print("Batch kernel not ready - needs Metal compilation")
            print("═══════════════════════════════════════════════════════════════════\n")
        }
    }
}