import XCTest
@testable import MarkBase

class LongContextTest: XCTestCase {
    
    func test26BStandardLongContext() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  26B-Standard Long Context Test (KV Cache Scaling)")
        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)
        
        print("Loading 26B-Standard with maxContext=2048...")
        let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 2048)
        print("✓ Model loaded")
        
        // Test different position ranges
        print("\nTesting position ranges:")
        var currentToken = 2
        
        // Test 0-10 (baseline)
        let start0 = Date()
        for i in 0..<10 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let time0 = Date().timeIntervalSince(start0) * 1000 / 10.0
        print("  Position 0-9: \(String(format: "%.1f", time0))ms/token")
        
        // Test 100-110
        let start100 = Date()
        for i in 100..<110 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let time100 = Date().timeIntervalSince(start100) * 1000 / 10.0
        print("  Position 100-109: \(String(format: "%.1f", time100))ms/token")
        
        // Test 500-510
        let start500 = Date()
        for i in 500..<510 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let time500 = Date().timeIntervalSince(start500) * 1000 / 10.0
        print("  Position 500-509: \(String(format: "%.1f", time500))ms/token")
        
        // Test 1000-1010
        let start1000 = Date()
        for i in 1000..<1010 {
            let result = try model.forwardOptimized(tokenId: currentToken, position: i)
            var maxIdx = 0
            var maxVal = result[0]
            for j in 1..<result.count {
                if result[j] > maxVal {
                    maxVal = result[j]
                    maxIdx = j
                }
            }
            currentToken = maxIdx
        }
        let time1000 = Date().timeIntervalSince(start1000) * 1000 / 10.0
        print("  Position 1000-1009: \(String(format: "%.1f", time1000))ms/token")
        
        // Check performance degradation
        let degradation = ((time1000 - time0) / time0) * 100.0
        print("\nPerformance analysis:")
        print("  Degradation at position 1000: \(String(format: "%.1f", degradation))%")
        
        if degradation < 20 {
            print("✓ KV cache efficient (<20% degradation)")
        } else {
            print("⚠ KV cache needs optimization (>20% degradation)")
        }
        
        XCTAssertLessThan(degradation, 50, "KV cache should not degrade >50%")
        
        print("\n═══════════════════════════════════════════════════════════════════")
    }
}