import XCTest
@testable import MarkBase

final class MetalKernelCompilationTest: XCTestCase {
    
    func testBasicMetalCompilation() throws {
        print("\n═══════════════════════════════════════")
        print("      Metal Kernel Compilation Test")
        print("═══════════════════════════════════════\n")
        
        print("Step 1: Create Metal engine...")
        let engine = try MarkBaseEngine()
        print("  ✓ Engine created")
        
        print("\nStep 2: Compile Metal kernels...")
        try engine.compileSource(MetalKernels.combinedSource)
        print("  ✓ Metal source compiled")
        
        print("\nStep 3: Test standard kernel (quantized_matmul_simd)...")
        do {
            let pso = try engine.pipeline(named: "quantized_matmul_simd")
            print("  ✓ Standard kernel compiled successfully")
            print("    Pipeline state: \(pso)")
        } catch {
            print("  ❌ Standard kernel failed: \(error)")
            throw error
        }
        
        print("\nStep 4: Test MoE kernel (quantized_matmul_gate_up)...")
        do {
            let pso = try engine.pipeline(named: "quantized_matmul_gate_up")
            print("  ✓ MoE 4-bit kernel compiled successfully")
            print("    Pipeline state: \(pso)")
        } catch {
            print("  ❌ MoE 4-bit kernel failed: \(error)")
            print("    This might be why generation hangs!")
            throw error
        }
        
        print("\nStep 5: Test MoE 8-bit kernel (quantized_matmul_gate_up_8bit)...")
        do {
            let pso = try engine.pipeline(named: "quantized_matmul_gate_up_8bit")
            print("  ✓ MoE 8-bit kernel compiled successfully")
            print("    Pipeline state: \(pso)")
        } catch {
            print("  ❌ MoE 8-bit kernel failed: \(error)")
            print("    26B-A4B uses 8-bit router, this is critical!")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ All Metal kernels compile successfully")
        print("═══════════════════════════════════════\n")
    }
    
    func testMetalKernelExecution() throws {
        print("\n═══════════════════════════════════════")
        print("      Metal Kernel Execution Test")
        print("═══════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine()
        try engine.compileSource(MetalKernels.combinedSource)
        
        print("Creating test buffers...")
        let device = engine.device
        
        // Create small test buffers
        let inputSize = 64
        let outputSize = 128
        
        let inputBuffer = device.makeBuffer(length: inputSize * MemoryLayout<Float>.size)!
        let outputBuffer = device.makeBuffer(length: outputSize * MemoryLayout<Float>.size)!
        
        // Initialize input
        var inputData = [Float](repeating: 1.0, count: inputSize)
        memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
        
        print("  ✓ Test buffers created")
        
        print("\nTesting standard kernel execution...")
        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        
        do {
            let pso = try engine.pipeline(named: "quantized_matmul_simd")
            let enc = cmdBuf.makeComputeCommandEncoder()!
            enc.setComputePipelineState(pso)
            enc.setBuffer(inputBuffer, offset: 0, index: 0)
            enc.setBuffer(outputBuffer, offset: 0, index: 4)
            
            let count = outputSize
            let tg = MTLSize(width: 256, height: 1, depth: 1)
            enc.dispatchThreads(MTLSize(width: count, height: 1, depth: 1), threadsPerThreadgroup: tg)
            enc.endEncoding()
            
            cmdBuf.commit()
            cmdBuf.waitUntilCompleted()
            
            print("  ✓ Standard kernel executed successfully")
            print("    Command buffer status: \(cmdBuf.status.rawValue)")
            
            if cmdBuf.status != .completed {
                print("  ❌ Command buffer did not complete!")
                let errorMsg = cmdBuf.error?.localizedDescription ?? "unknown error"
                print("    Error: \(errorMsg)")
                XCTFail("Command buffer execution failed")
            }
            
        } catch {
            print("  ❌ Kernel execution failed: \(error)")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ Metal kernel execution works")
        print("═══════════════════════════════════════\n")
    }
}