import XCTest
@testable import MarkBase

class AllModelsFinalTest: XCTestCase {
    
    func testAllModelsTextForwardFinal() throws {
        print("\n═══════════════════════════════════════════════════════════════════")
        print("  All Models TEXT Forward Final Test (Zero NaN Verification)")
        print("═══════════════════════════════════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine(autoCompile: true)
        
        let models = [
            ("E2B", "/Users/accusys/MarkBaseEngine/models/gemma-4-e2b-it-4bit"),
            ("26B-Standard", "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-standard"),
            ("31B", "/Users/accusys/MarkBaseEngine/models/gemma-4-31b-it-4bit"),
            ("26B-A4B", "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit")
        ]
        
        var successCount = 0
        var failedModels: [(String, String)] = []
        
        for (modelName, modelPath) in models {
            print("═══════════════════════════════════════════════════════════════════")
            print("Testing: \(modelName)")
            print("═══════════════════════════════════════════════════════════════════")
            
            guard FileManager.default.fileExists(atPath: modelPath) else {
                print("  ⚠ Model not found, skipping")
                failedModels.append((modelName, "Model not found"))
                continue
            }
            
            do {
                print("  Loading model...")
                let model = try E4BModel(modelDir: modelPath, engine: engine, maxContextLength: 128)
                print("  ✓ Loaded")
                print("    Layers: \(model.numHiddenLayers)")
                print("    Hidden: \(model.hiddenSize)")
                
                print("  Testing forward pass...")
                let result = try model.forwardOptimized(tokenId: 2, position: 0)
                
                let nanCount = result.filter { $0.isNaN }.count
                print("  Forward result: NaN=\(nanCount)/\(result.count)")
                
                if nanCount == 0 {
                    print("  ✓✓✓ Zero NaN - Success!")
                    successCount += 1
                } else {
                    print("  ✗ NaN detected")
                    failedModels.append((modelName, "NaN detected"))
                }
                
            } catch {
                print("  ✗ Failed: \(error)")
                failedModels.append((modelName, error.localizedDescription))
            }
            
            print()
        }
        
        print("═══════════════════════════════════════════════════════════════════")
        print("  Summary")
        print("═══════════════════════════════════════════════════════════════════")
        print("  Success: \(successCount)/\(models.count)")
        
        if failedModels.isEmpty {
            print("  ✓✓✓✓✓✓ All models zero NaN!")
        } else {
            print("  Failed models:")
            for (name, reason) in failedModels {
                print("    - \(name): \(reason)")
            }
        }
        print("═══════════════════════════════════════════════════════════════════\n")
        
        XCTAssertGreaterThanOrEqual(successCount, 1, "At least one model should succeed")
    }
}