import XCTest
@testable import MarkBase

final class MoEForwardWithFixedExpertTest: XCTestCase {
    
    func testMoEForwardWithFixedExpert() throws {
        print("\n═══════════════════════════════════════")
        print("      MoE Forward Pass - Fixed Expert")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Step 1: Load model...")
        fflush(stdout)
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        print("  ✓ Model loaded (30 layers, MoE)")
        fflush(stdout)
        
        print("\nStep 2: Get layer 0...")
        let layer0 = model.layers[0]
        print("  ✓ Layer 0 retrieved (useMoE: \(layer0.useMoE))")
        fflush(stdout)
        
        print("\nStep 3: Create buffers...")
        let hs = model.hiddenSize
        
        // Create input buffer (normalized hidden state)
        let inputBuffer = engine.device.makeBuffer(length: hs * MemoryLayout<Float>.size)!
        var inputData = [Float](repeating: 0.1, count: hs)
        memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
        
        print("  ✓ Input buffer created (\(hs) floats)")
        fflush(stdout)
        
        print("\nStep 4: Test MoE forward (layer 0)...")
        fflush(stdout)
        
        let start = Date()
        
        do {
            print("  Creating command buffer...")
            fflush(stdout)
            let cmdBuf = engine.commandQueue.makeCommandBuffer()!
            
            print("  Creating KV cache...")
            fflush(stdout)
            let kvCache = try KVCache(
                device: engine.device,
                isSliding: false,
                maxContextLength: 32,
                nKvHeads: 8,
                headDim: 256
            )
            
            print("  Calling layer0.forward...")
            fflush(stdout)
            
            try layer0.forward(
                input: inputBuffer,
                position: 0,
                kvCache: kvCache,
                shouldStoreKV: false,
                temps: model.temps,
                engine: engine
            )
            
            print("  Committing command buffer...")
            fflush(stdout)
            cmdBuf.commit()
            
            print("  Waiting for completion...")
            fflush(stdout)
            cmdBuf.waitUntilCompleted()
            
            let elapsed = Date().timeIntervalSince(start)
            print("\n  ✓ Forward pass completed in \(String(format: "%.3f", elapsed))s")
            print("    Command buffer status: \(cmdBuf.status.rawValue)")
            fflush(stdout)
            
            if elapsed > 10.0 {
                print("  ⚠️ Forward pass took >10s")
                fflush(stdout)
            }
            
            if cmdBuf.status != .completed {
                print("  ❌ Forward pass did not complete!")
                fflush(stdout)
                XCTFail("Forward pass failed")
                return
            }
            
        } catch {
            print("  ❌ Forward pass error: \(error)")
            fflush(stdout)
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ MoE forward pass works!")
        print("═══════════════════════════════════════\n")
        fflush(stdout)
    }
}