import XCTest
@testable import MarkBase

final class MoEForwardTests: XCTestCase {
    
    func test26BA4BModelLoading() throws {
        print("\n═══════════════════════════════════════")
        print("      26B-A4B MoE Model Loading Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Step 1: Initialize Metal engine...")
        let engine = try MarkBaseEngine()
        try engine.compileSource(MetalKernels.combinedSource)
        print("  ✓ Engine ready\n")
        
        print("Step 2: Load 26B-A4B MoE model...")
        
        do {
            let start = Date()
            let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
            let loadTime = Date().timeIntervalSince(start)
            
            print("  ✓ Model loaded successfully in \(String(format: "%.3f", loadTime))s\n")
            
            print("Model parameters:")
            print("  Layers: \(model.numHiddenLayers)")
            print("  Hidden size: \(model.hiddenSize)")
            print("  Vocab size: \(model.vocabSize)")
            
            // Verify 26B-A4B parameters
            XCTAssertEqual(model.numHiddenLayers, 30, "Should have 30 layers")
            XCTAssertEqual(model.hiddenSize, 2816, "Should have hidden_size=2816")
            XCTAssertEqual(model.vocabSize, 262144, "Should have vocab_size=262144")
            
            print("\n═══════════════════════════════════════")
            print("✓ MoE model loading test passed")
            print("═══════════════════════════════════════\n")
            
        } catch {
            print("  ✗ Model loading failed: \(error)\n")
            throw error
        }
    }
    
    func test26BA4BSimpleGeneration() throws {
        print("\n═══════════════════════════════════════")
        print("      26B-A4B MoE Simple Generation")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        let engine = try MarkBaseEngine()
        try engine.compileSource(MetalKernels.combinedSource)
        
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
        
        print("Testing simple generation (5 tokens)...")
        let prompt = "Hello"
        let config = GenerationConfig(maxTokens: 5, temperature: Float(0.7))
        
        let start = Date()
        let response = try generator.generateComplete(prompt: prompt, config: config)
        let genTime = Date().timeIntervalSince(start)
        
        print("  ✓ Generated in \(String(format: "%.3f", genTime))s")
        print("  Response: \"\(response)\"")
        print("  Speed: \(String(format: "%.2f", 5.0 / genTime)) tok/s")
        
        XCTAssertGreaterThan(response.count, 0, "Should generate some text")
        
        print("\n═══════════════════════════════════════")
        print("✓ Simple generation test passed")
        print("═══════════════════════════════════════\n")
    }
}