import XCTest
@testable import MarkBase

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