import XCTest
@testable import MarkBase

final class MoETests: XCTestCase {
    
    func test26BA4BLoading() async throws {
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("\n=====================================")
        print("Testing 26B-A4B MoE Model Loading")
        print("=====================================")
        
        let start = Date()
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let loadTime = Date().timeIntervalSince(start)
        
        print("✓ Model loaded in \(String(format: "%.3f", loadTime))s")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Hidden: \(model.hiddenSize)")
        print("  Vocab: \(model.vocabSize)")
        
        XCTAssertGreaterThan(model.numHiddenLayers, 0)
        XCTAssertEqual(model.hiddenSize, 2816)
        XCTAssertEqual(model.vocabSize, 262144)
    }
    
    func test26BA4BGeneration() async throws {
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("\n=====================================")
        print("Testing 26B-A4B MoE Token Generation")
        print("=====================================")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
        
        let prompt = "Hello"
        let config = GenerationConfig(maxTokens: 10, temperature: Float(0.7))
        
        print("Testing generation...")
        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", 10.0 / genTime)) tok/s")
        
        XCTAssertGreaterThan(response.count, 0)
    }
}