import XCTest
@testable import MarkBase

final class G12BGeneratorTests: XCTestCase {
    
    func testSampler() throws {
        print("\n═══════════════════════════════════════")
        print("      Sampler Test")
        print("═══════════════════════════════════════\n")
        
        print("Step 1: Test greedy sampling...")
        let sampler = Sampler()
        
        // Create test logits
        let logits: [Float] = [1.0, 2.0, 3.0, 0.5, 1.5]
        
        let greedyToken = sampler.greedySample(logits: logits)
        print("  Logits: \(logits)")
        print("  Greedy token: \(greedyToken)")
        XCTAssertEqual(greedyToken, 2, "Greedy should pick index 2 (max=3.0)")
        
        print("\nStep 2: Test temperature sampling...")
        // High temperature = more random
        let highTempToken = sampler.sample(logits: logits, temperature: 2.0)
        print("  Temperature=2.0 token: \(highTempToken)")
        
        // Low temperature = more deterministic
        let lowTempToken = sampler.sample(logits: logits, temperature: 0.1)
        print("  Temperature=0.1 token: \(lowTempToken)")
        XCTAssertEqual(lowTempToken, 2, "Low temperature should be near greedy")
        
        print("\nStep 3: Test top-k sampling...")
        let topKToken = sampler.sample(logits: logits, temperature: 1.0, topK: 2)
        print("  TopK=2 token: \(topKToken)")
        // Should be in top 2 indices (2 or 4)
        XCTAssert(topKToken == 2 || topKToken == 4, "Should be in top-2")
        
        print("\nStep 4: Test top-p sampling...")
        let topPToken = sampler.sample(logits: logits, temperature: 1.0, topP: 0.8)
        print("  TopP=0.8 token: \(topPToken)")
        
        print("\n═══════════════════════════════════════")
        print("✓ Sampler test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testStreamingGenerator() async throws {
        print("\n═══════════════════════════════════════")
        print("      Streaming Generator Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load model and tokenizer...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
        print("  ✓ Generator created")
        
        print("\nStep 2: Test complete generation...")
        let config = GenerationConfig(maxTokens: 5, temperature: 0.8)
        let response = try generator.generateComplete(prompt: "Hello", config: config)
        print("  Prompt: Hello")
        print("  Response: \(response)")
        XCTAssertGreaterThan(response.count, 0, "Should generate text")
        
        print("\nStep 3: Test streaming generation...")
        print("  Streaming tokens:")
        var streamedTokens: [String] = []
        
        let streamConfig = GenerationConfig(maxTokens: 5, temperature: 1.0)
        for await tokenText in generator.generate(prompt: "Test", config: streamConfig) {
            print("    Token: \(tokenText)")
            streamedTokens.append(tokenText)
        }
        
        XCTAssertGreaterThan(streamedTokens.count, 0, "Should stream tokens")
        
        print("\n═══════════════════════════════════════")
        print("✓ Streaming generator test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testGenerationPerformance() throws {
        print("\n═══════════════════════════════════════")
        print("      Generation Performance Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
        
        print("\nStep 2: Generate 20 tokens...")
        let config = GenerationConfig(maxTokens: 20, temperature: 0.8)
        
        let start = Date()
        let response = try generator.generateComplete(prompt: "Hello", config: config)
        let duration = Date().timeIntervalSince(start)
        
        print("  Tokens generated: ~20")
        print("  Time: \(duration) seconds")
        print("  Speed: \(20.0 / duration) tokens/second")
        print("  Response: \(response)")
        
        XCTAssertGreaterThan(response.count, 0)
        
        print("\n═══════════════════════════════════════")
        print("✓ Generation performance test passed")
        print("═══════════════════════════════════════\n")
    }
}