import XCTest
@testable import MarkBase

final class G12BTextGenerationTests: XCTestCase {
    
    func testTextGeneration() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Text Generation Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load tokenizer...")
        // Simple tokenizer test - check if tokenizer.json exists
        let tokenizerPath = modelDir + "/tokenizer.json"
        let tokenizerExists = FileManager.default.fileExists(atPath: tokenizerPath)
        print("  Tokenizer exists: \(tokenizerExists)")
        if !tokenizerExists {
            print("  ⚠️ No tokenizer found - using simple token IDs\n")
        }
        
        print("Step 2: Initialize engine and model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        print("  ✓ Model loaded\n")
        
        print("Step 3: Test generation with simple prompts...")
        
        // Test 1: Generate from token 0 (usually special token)
        print("\nTest 1: Start token (ID=0)")
        var generatedTokens: [Int] = []
        let startToken = 0
        
        for position in 0..<10 {
            let logits = try model.forward(tokenId: startToken, position: position)
            
            // Find max logit (greedy sampling)
            var maxLogit = logits[0]
            var maxIdx = 0
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            generatedTokens.append(maxIdx)
            print("  Position \(position): token=\(maxIdx), logit=\(maxLogit)")
        }
        
        print("  Generated tokens: \(generatedTokens)")
        
        // Test 2: Different start token
        print("\nTest 2: Token ID=1")
        generatedTokens = []
        for position in 0..<10 {
            let logits = try model.forward(tokenId: 1, position: position)
            var maxLogit = logits[0]
            var maxIdx = 0
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            generatedTokens.append(maxIdx)
        }
        print("  Generated tokens: \(generatedTokens)")
        
        // Test 3: Sequence generation
        print("\nTest 3: Sequence generation (different tokens per position)")
        var sequence: [Int] = [1, 2, 3, 4, 5] // Start sequence
        print("  Input sequence: \(sequence)")
        
        for i in 0..<5 {
            let logits = try model.forward(tokenId: sequence[i], position: i)
            var maxLogit = logits[0]
            var maxIdx = 0
            for j in 1..<logits.count {
                if logits[j] > maxLogit {
                    maxLogit = logits[j]
                    maxIdx = j
                }
            }
            sequence.append(maxIdx)
        }
        print("  Extended sequence: \(sequence)")
        
        print("\n═══════════════════════════════════════")
        print("✓ Text generation test passed")
        print("═══════════════════════════════════════\n")
        
        // Verify diversity - tokens should not all be the same
        let uniqueTokens = Set(generatedTokens)
        XCTAssertGreaterThan(uniqueTokens.count, 1, "Generated tokens should have diversity")
    }
    
    func testE4BInference() throws {
        print("\n═══════════════════════════════════════")
        print("      E4B-MarkBase Inference Test")
        print("═══════════════════════════════════════\n")

        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"

        print("Step 1: Load engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine initialized\n")

        print("Step 2: Load tokenizer...")
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        print("  ✓ Tokenizer loaded: vocabSize=\(tokenizer.vocabSize)\n")

        print("Step 3: Load model...")
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        print("  ✓ Model loaded: hiddenSize=\(model.hiddenSize), layers=\(model.numHiddenLayers)\n")
        
        // Debug: check logits at position 0
        print("Debug: Single token test (BOS at pos 0)")
        let logits0 = try model.forward(tokenId: 2, position: 0)
        print("  logits count: \(logits0.count)")
        print("  min: \(logits0.min() ?? 0), max: \(logits0.max() ?? 0)")
        let sorted0 = logits0.enumerated().sorted { $0.element > $1.element }
        print("  Top 10 tokens:")
        for (i, (idx, val)) in sorted0.prefix(10).enumerated() {
            let text = tokenizer.decode(tokens: [idx])
            print("    \(i+1). token \(idx) '\(text)': \(val)")
        }
        print("")

        print("Step 4: Create generator...")
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)

        print("Step 5: Text completion inference...")
        let prompt = "The capital of France is"
        print("  Prompt: \"\(prompt)\"")
        let config = GenerationConfig(maxTokens: 30, temperature: 1.0, topK: 40, topP: 0.9)
        let response = try generator.generateComplete(prompt: prompt, config: config)
        print("  Response: \"\(response)\"")

        XCTAssertFalse(response.isEmpty, "Response should not be empty")
        print("  ✓ Text generation complete\n")

        print("Step 6: Chat-style inference...")
        let chatPrompt = "<|turn>user\nWhat is 2+2?<turn|>\n<|turn>model\n"
        print("  Prompt: \"What is 2+2?\"")
        let chatResponse = try generator.generateComplete(prompt: chatPrompt, config: config)
        print("  Response: \"\(chatResponse)\"")

        XCTAssertFalse(chatResponse.isEmpty, "Chat response should not be empty")
        print("  ✓ Chat generation complete\n")

        print("Step 7: Chinese generation test...")
        let chinesePrompt = "<|turn>user\n用中文說你好<turn|>\n<|turn>model\n"
        print("  Prompt: \"用中文說你好\"")
        let chineseResponse = try generator.generateComplete(prompt: chinesePrompt, config: config)
        print("  Response: \"\(chineseResponse)\"")

        print("\n═══════════════════════════════════════")
        print("✓ E4B inference test passed")
        print("═══════════════════════════════════════\n")
    }

    func testGenerationQuality() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Generation Quality 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)
        print("  ✓ Model loaded\n")
        
        print("Step 2: Check logits distribution...")
        
        // Generate logits for several tokens
        var allLogits: [[Float]] = []
        for tokenId in [0, 1, 100, 1000, 10000] {
            let logits = try model.forward(tokenId: tokenId, position: 0)
            allLogits.append(logits)
            
            // Stats
            let slice = logits[0..<1000]
            let minVal = slice.min() ?? 0
            let maxVal = slice.max() ?? 0
            let mean = slice.reduce(0, +) / Float(slice.count)
            
            print("  Token \(tokenId): min=\(minVal), max=\(maxVal), mean=\(mean)")
        }
        
        print("\nStep 3: Verify generation diversity...")
        
        // Generate 20 tokens and check if they're diverse
        var tokens: [Int] = []
        for i in 0..<20 {
            let logits = try model.forward(tokenId: i, position: i)
            
            // Top-5 sampling (not just greedy)
            var sortedLogits = logits.enumerated().sorted { $0.element > $1.element }
            let top5 = sortedLogits[0..<5].map { $0.offset }
            
            // Pick random from top-5 (simulate sampling)
            let selectedToken = top5[i % 5] // Deterministic for test
            tokens.append(selectedToken)
        }
        
        print("  Generated tokens: \(tokens)")
        let uniqueTokens = Set(tokens)
        print("  Unique tokens: \(uniqueTokens.count)")
        
        XCTAssertGreaterThan(uniqueTokens.count, 5, "Should have reasonable diversity")
        
        print("\n═══════════════════════════════════════")
        print("✓ Generation quality test passed")
        print("═══════════════════════════════════════\n")
    }
}