import XCTest
@testable import MarkBase

final class E2BAudioLoadTest: XCTestCase {
    func testE2BAudioLoad() throws {
        print("\n═══════════════════════════════════════")
        print("      E2B AudioTower Loading Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-e2b-it-4bit/snapshots/2c3e507453b4f218d05fe3cc97bea5c5a654257e"
        
        print("Step 1: Load AudioConfig...")
        let acfg = loadAudioConfig(modelDir: modelDir)
        print("  hiddenSize: \(acfg.hiddenSize)")
        print("  numHiddenLayers: \(acfg.numHiddenLayers)")
        
        print("\nStep 2: Create engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("\nStep 3: Load safetensors...")
        let reader = try SafeTensorsReader(path: modelDir + "/model.safetensors")
        
        print("\nStep 4: Load AudioTowerE2B...")
        do {
            let tower = try loadAudioTowerE2B(reader: reader, config: acfg, engine: engine)
            print("  ✓ AudioTowerE2B loaded: \(tower.config.numHiddenLayers) layers")
            
            print("\nStep 5: Test forward pass...")
            let seqLen = 50
            let nMels = 128
            var melFeatures: [Float] = Array(repeating: 0.1, count: seqLen * nMels)
            for i in 0..<melFeatures.count { melFeatures[i] = Float(i % 100) / 100.0 * 0.5 }
            
            let inputBuffer = engine.device.makeBuffer(bytes: melFeatures, length: melFeatures.count * 4)!
            let outputLen = seqLen / 4 * tower.config.outputProjDims
            let outputBuffer = engine.device.makeBuffer(length: outputLen * 4)!
            
            try tower.forward(inputBuffer: inputBuffer, seqLen: seqLen, outputBuffer: outputBuffer)
            
            let ptr = outputBuffer.contents().assumingMemoryBound(to: Float.self)
            var hasNaN = false
            var maxVal: Float = 0, minVal: Float = 0
            for i in 0..<outputLen {
                let v = ptr[i]
                if v.isNaN { hasNaN = true }
                if v > maxVal { maxVal = v }
                if v < minVal { minVal = v }
            }
            print("  Output shape: [\(seqLen/4), \(tower.config.outputProjDims)]")
            print("  Range: [\(minVal), \(maxVal)]")
            print("  NaN: \(hasNaN)")
            
            XCTAssertFalse(hasNaN, "E2B audio output should not have NaN")
        } catch {
            print("  ✗ Error: \(error)")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ E2B AudioTower test passed")
        print("═══════════════════════════════════════\n")
    }
}