import XCTest
@testable import MarkBase

final class E4BAudioMultimodalTest: XCTestCase {
    func testAudioMultimodalGeneration() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B Audio Multimodal Generation Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        print("Step 1: Load multimodal model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 256)
        print("  ✓ Model loaded: hidden=\(mmModel.textModel.hiddenSize), layers=\(mmModel.textModel.numHiddenLayers)")
        print("  Audio tower: \(mmModel.audioTowerFull != nil ? "Full" : "12B")")
        
        print("\nStep 2: Create fake audio features (mel spectrogram)...")
        let seqLen = 98
        let nMels = 128
        var melFeatures: [[Float]] = []
        for _ in 0..<seqLen {
            var frame: [Float] = []
            for m in 0..<nMels {
                frame.append(Float(m) / 128.0 * 0.5)
            }
            melFeatures.append(frame)
        }
        print("  Mel shape: [\(seqLen), \(nMels)]")
        
        print("\nStep 3: Process audio through AudioTower...")
        let audioEmbeds = try mmModel.processAudio(audioFeatures: melFeatures)
        
        let audioOutputDim: Int
        if let tower = mmModel.audioTowerFull {
            audioOutputDim = tower.config.outputProjDims
        } else if let tower = mmModel.audioTowerE2B {
            audioOutputDim = tower.config.outputProjDims
        } else {
            audioOutputDim = mmModel.textModel.hiddenSize
        }
        
        print("  Output shape: [\(audioEmbeds.count / audioOutputDim), \(audioOutputDim)]")
        print("  Range: [\(audioEmbeds.min() ?? 0), \(audioEmbeds.max() ?? 0)]")
        print("  NaN check: \(audioEmbeds.contains { $0.isNaN })")
        
        XCTAssertFalse(audioEmbeds.contains { $0.isNaN }, "Audio embeddings should not have NaN")
        
        print("\nStep 4: Generate text with audio context...")
        // BOS + audio placeholder tokens
        let audioTokenId = mmModel.audioTokenId
        let boaTokenId = mmModel.boaTokenId
        let eoaTokenId = mmModel.eoaTokenId
        
        // Build prompt: BOS + <|audio|> + <boa> + ... + <eoa>
        var tokens: [Int] = [2]  // BOS
        tokens.append(audioTokenId)
        tokens.append(boaTokenId)
        // Audio embeds occupy seqLen/4 = 24 positions
        for _ in 0..<seqLen / 4 {
            tokens.append(audioTokenId)
        }
        tokens.append(eoaTokenId)
        tokens.append(boaTokenId)  // Start response
        print("  Prompt tokens: \(tokens.count)")
        
        // Generate 10 tokens
        var generated: [Int] = tokens
        let startTime = Date()
        for _ in 0..<10 {
            let pos = generated.count - 1
            let logits = try mmModel.textModel.forward(tokenId: generated.last!, position: pos)
            var maxIdx = 0
            var maxLogit = logits[0]
            for i in 1..<logits.count {
                if logits[i] > maxLogit {
                    maxLogit = logits[i]
                    maxIdx = i
                }
            }
            generated.append(maxIdx)
        }
        let elapsed = Date().timeIntervalSince(startTime)
        print("  Generated: \(generated.suffix(10))")
        print("  Speed: \(10.0 / elapsed) tok/s")
        
        print("\n═══════════════════════════════════════")
        print("✓ Audio multimodal test passed")
        print("═══════════════════════════════════════\n")
    }
}