import XCTest
@testable import MarkBase

final class G12BTokenizerTests: XCTestCase {
    
    func testTokenizerLoading() throws {
        print("\n═══════════════════════════════════════")
        print("      Tokenizer Loading Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Check tokenizer files...")
        let tokenizerJsonPath = modelDir + "/tokenizer.json"
        let tokenizerModelPath = modelDir + "/tokenizer.model"
        
        let hasTokenizerJson = FileManager.default.fileExists(atPath: tokenizerJsonPath)
        let hasTokenizerModel = FileManager.default.fileExists(atPath: tokenizerModelPath)
        
        print("  tokenizer.json: \(hasTokenizerJson ? "✓ Found" : "✗ Not found")")
        print("  tokenizer.model: \(hasTokenizerModel ? "✓ Found" : "✗ Not found")")
        
        XCTAssert(hasTokenizerJson || hasTokenizerModel, "Should have at least one tokenizer file")
        
        print("\nStep 2: Load tokenizer...")
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        print("  ✓ Tokenizer loaded")
        print("  Type: \(type(of: tokenizer))")
        print("  Vocab size: \(tokenizer.vocabSize)")
        print("  BOS token: \(tokenizer.bosTokenId)")
        print("  EOS token: \(tokenizer.eosTokenId)")
        
        XCTAssertGreaterThan(tokenizer.vocabSize, 0, "Vocab size should be > 0")
        
        print("\n═══════════════════════════════════════")
        print("✓ Tokenizer loading test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testEncodingDecoding() throws {
        print("\n═══════════════════════════════════════")
        print("      Encoding/Decoding Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        print("Step 1: Test encode...")
        let testText = "Hello world"
        let tokens = tokenizer.encode(text: testText)
        print("  Input: \(testText)")
        print("  Tokens: \(tokens)")
        print("  Token count: \(tokens.count)")
        
        XCTAssertGreaterThan(tokens.count, 2, "Should have at least BOS + EOS + content")
        XCTAssertEqual(tokens.first, tokenizer.bosTokenId, "First token should be BOS")
        XCTAssertEqual(tokens.last, tokenizer.eosTokenId, "Last token should be EOS")
        
        print("\nStep 2: Test decode...")
        let decodedText = tokenizer.decode(tokens: tokens)
        print("  Tokens: \(tokens)")
        print("  Decoded: \(decodedText)")
        
        // Note: May not be exact match due to simplified implementation
        XCTAssertGreaterThan(decodedText.count, 0, "Decoded text should not be empty")
        
        print("\nStep 3: Test roundtrip...")
        let reencodedTokens = tokenizer.encode(text: decodedText)
        let redecodedText = tokenizer.decode(tokens: reencodedTokens)
        print("  Original: \(testText)")
        print("  Roundtrip: \(redecodedText)")
        
        print("\n═══════════════════════════════════════")
        print("✓ Encoding/decoding test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testSpecialTokens() throws {
        print("\n═══════════════════════════════════════")
        print("      Special Tokens Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        print("Step 1: Verify special token IDs...")
        print("  BOS token ID: \(tokenizer.bosTokenId)")
        print("  EOS token ID: \(tokenizer.eosTokenId)")
        print("  PAD token ID: \(tokenizer.padTokenId)")
        
        XCTAssertGreaterThanOrEqual(tokenizer.bosTokenId, 0, "BOS ID should be valid")
        XCTAssertGreaterThanOrEqual(tokenizer.eosTokenId, 0, "EOS ID should be valid")
        XCTAssertGreaterThanOrEqual(tokenizer.padTokenId, 0, "PAD ID should be valid")
        
        print("\nStep 2: Test with empty text...")
        let emptyTokens = tokenizer.encode(text: "")
        print("  Empty input tokens: \(emptyTokens)")
        XCTAssertEqual(emptyTokens, [tokenizer.bosTokenId, tokenizer.eosTokenId], "Empty text should be BOS + EOS only")
        
        print("\n═══════════════════════════════════════")
        print("✓ Special tokens test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testTokenizerWithModel() throws {
        print("\n═══════════════════════════════════════")
        print("      Tokenizer + Model Integration")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load tokenizer and model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        print("  ✓ Model loaded")
        print("  ✓ Tokenizer loaded")
        print("  Model vocab size: \(model.vocabSize)")
        print("  Tokenizer vocab size: \(tokenizer.vocabSize)")
        
        // Note: vocab sizes should match (both 262144 for Gemma-4)
        XCTAssertEqual(model.vocabSize, tokenizer.vocabSize, "Model and tokenizer vocab should match")
        
        print("\nStep 2: Generate tokens from text...")
        let prompt = "Hello"
        let tokens = tokenizer.encode(text: prompt)
        print("  Prompt: \(prompt)")
        print("  Tokens: \(tokens)")
        
        // Generate from first token
        let firstToken = tokens[1] // Skip BOS
        let logits = try model.forward(tokenId: firstToken, position: 0)
        print("  Generated logits for token \(firstToken)")
        print("  Logits count: \(logits.count)")
        
        // Find next token (greedy)
        var maxLogit = logits[0]
        var nextTokenId = 0
        for i in 1..<logits.count {
            if logits[i] > maxLogit {
                maxLogit = logits[i]
                nextTokenId = i
            }
        }
        
        let nextTokenText = tokenizer.decode(tokens: [nextTokenId])
        print("  Next token: \(nextTokenId)")
        print("  Next token text: \(nextTokenText)")
        
        print("\n═══════════════════════════════════════")
        print("✓ Tokenizer + Model integration passed")
        print("═══════════════════════════════════════\n")
    }
}