import XCTest
@testable import MarkBase

final class TextModelLoadTest: XCTestCase {
    
    func testE4BTextLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B TEXT Model Load Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Load TEXT model only...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let loadStart = Date()
        let textModel = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        
        print("  ✓ TEXT model loaded in \(loadTime) ms")
        print("  Hidden: \(textModel.hiddenSize)")
        print("  Layers: \(textModel.numHiddenLayers)")
        
        print("\nStep 2: Generate 10 tokens...")
        var tokens: [Int] = [2]  // BOS
        let genStart = Date()
        for _ in 0..<10 {
            let logits = try textModel.forward(tokenId: tokens.last!, position: tokens.count - 1)
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            tokens.append(maxIdx)
        }
        let genTime = Date().timeIntervalSince(genStart) * 1000
        print("  ✓ Generated 10 tokens in \(genTime) ms")
        print("  Tokens: \(tokens.suffix(10))")
        
        print("\n═══════════════════════════════════════")
        print("✓ E4B TEXT test passed")
        print("  Load: \(loadTime) ms, Gen: \(genTime) ms")
        print("═══════════════════════════════════════\n")
    }
    
    func test12BTextLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  12B TEXT Model Load Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load TEXT model only...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let loadStart = Date()
        let textModel = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        let loadTime = Date().timeIntervalSince(loadStart) * 1000
        
        print("  ✓ TEXT model loaded in \(loadTime) ms")
        print("  Hidden: \(textModel.hiddenSize)")
        print("  Layers: \(textModel.numHiddenLayers)")
        
        print("\nStep 2: Generate 10 tokens...")
        var tokens: [Int] = [2]  // BOS
        let genStart = Date()
        for _ in 0..<10 {
            let logits = try textModel.forward(tokenId: tokens.last!, position: tokens.count - 1)
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            tokens.append(maxIdx)
        }
        let genTime = Date().timeIntervalSince(genStart) * 1000
        print("  ✓ Generated 10 tokens in \(genTime) ms")
        print("  Tokens: \(tokens.suffix(10))")
        
        print("\n═══════════════════════════════════════")
        print("✓ 12B TEXT test passed")
        print("  Load: \(loadTime) ms, Gen: \(genTime) ms")
        print("═══════════════════════════════════════\n")
    }
}