import Foundation
import MarkBase

// Simple test for 26B-A4B MoE model

let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"

print("=====================================")
print("26B-A4B MoE Simple Test")
print("=====================================")

do {
    print("\n[1] Loading model...")
    let start1 = Date()
    let engine = try E4BEngine(autoCompile: true)
    let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 512)
    let loadTime = Date().timeIntervalSince(start1)
    print("✓ Model loaded in \(String(format: "%.3f", loadTime))s")
    print("  Layers: \(model.numHiddenLayers)")
    print("  Hidden: \(model.hiddenSize)")
    print("  Vocab: \(model.vocabSize)")
    
    print("\n[2] Loading tokenizer...")
    let tokenizer = try TokenizerFactory.load(modelDir: modelPath)
    print("✓ Tokenizer loaded")
    
    print("\n[3] Testing forward pass...")
    let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
    let sampler = Sampler()
    
    let prompt = "Hello"
    let tokens = tokenizer.encode(text: prompt)
    print("  Prompt: \"\(prompt)\"")
    print("  Tokens: \(tokens)")
    
    print("\n[4] Testing token generation (temperature=0.7, max_tokens=10)...")
    let config = GenerationConfig(maxTokens: 10, temperature: Float(0.7))
    
    let start2 = Date()
    let response = try generator.generateComplete(prompt: prompt, config: config)
    let genTime = Date().timeIntervalSince(start2)
    
    print("✓ Generated in \(String(format: "%.3f", genTime))s")
    print("  Response: \"\(response)\"")
    print("  Speed: \(String(format: "%.2f", 10.0 / genTime)) tok/s")
    
    print("\n✅ All tests passed!")
    
} catch {
    print("❌ Error: \(error)")
    exit(1)
}