import XCTest
@testable import MarkBase

class LongTextLimitTest: XCTestCase {
    
    func testLongTextLimit() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Long Text Limit Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Test different context lengths
        let contextLengths = [128, 256, 512, 1024, 2048]
        
        for maxLen in contextLengths {
            print("\n───────────────────────────────────────────────────────────────────")
            print("Testing maxContextLength: \(maxLen)")
            print("───────────────────────────────────────────────────────────────────")
            
            do {
                let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: maxLen)
                print("✓ Model loaded with maxContextLength=\(maxLen)")
                print("  Layers: \(model.numHiddenLayers)")
                print("  Hidden: \(model.hiddenSize)")
                
                // Test forward pass at position 0
                let result0 = try model.forwardOptimized(tokenId: 2, position: 0)
                let nan0 = result0.filter { $0.isNaN }.count
                print("  Forward at pos 0: NaN=\(nan0)/\(result0.count)")
                
                // Test forward pass at max position
                let maxPos = maxLen - 1
                if maxPos < 1000 {  // Limit test time
                    let resultMax = try model.forwardOptimized(tokenId: 2, position: maxPos)
                    let nanMax = resultMax.filter { $0.isNaN }.count
                    print("  Forward at pos \(maxPos): NaN=\(nanMax)/\(resultMax.count)")
                    
                    if nanMax == 0 {
                        print("  ✓✓✓ Position \(maxPos) OK")
                    } else {
                        print("  ✗ NaN at position \(maxPos)")
                    }
                }
                
            } catch {
                print("✗ Failed: \(error)")
            }
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Test completed")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
    
    func testMemoryUsage() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Memory Usage Test")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let modelPath = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"
        
        guard FileManager.default.fileExists(atPath: modelPath) else {
            print("⚠ Model not found")
            return
        }
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Calculate theoretical memory usage
        let configs = [
            ("128 context", 128),
            ("512 context", 512),
            ("1024 context", 1024),
            ("2048 context", 2048),
            ("4096 context", 4096),
            ("8192 context", 8192)
        ]
        
        print("\nTheoretical KV Cache Memory Usage:")
        print("Model: 26B-Standard (30 layers, hidden=2816, nKvHeads=4, headDim=256)")
        
        for (name, maxLen) in configs {
            // KV cache memory: 2 * numLayers * maxLen * nKvHeads * headDim * 4 bytes
            let kvMemory = 2 * 30 * maxLen * 4 * 256 * 4  // K + V, 30 layers, Float32
            let kvMemoryMB = Double(kvMemory) / 1024.0 / 1024.0
            
            print("  \(name): \(String(format: "%.2f", kvMemoryMB)) MB")
            
            if kvMemoryMB < 100 {
                print("    ✓ Safe for 128GB unified memory")
            } else if kvMemoryMB < 500 {
                print("    ⚠ Moderate usage")
            } else {
                print("    ✗ High memory usage")
            }
        }
        
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  Test completed")
        print("═══════════════════════════════════════════════════════════════════\n")
    }
}