import XCTest
@testable import MarkBase

final class G12BMultimodalTests: XCTestCase {
    
    func testMultimodalStructure() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Multimodal Structure Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Load multimodal model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        print("  ✓ Multimodal model loaded\n")
        
        print("Step 2: Check multimodal token IDs...")
        print("  audio_token_id: \(mmModel.audioTokenId)")
        print("  boa_token_id: \(mmModel.boaTokenId)")
        print("  eoa_token_id: \(mmModel.eoaTokenId)")
        print("  image_token_id: \(mmModel.imageTokenId)")
        print("  boi_token_id: \(mmModel.boiTokenId)")
        print("  eoi_token_id: \(mmModel.eoiTokenId)")
        
        // Verify token IDs match 12B config
        XCTAssertEqual(mmModel.audioTokenId, 258881, "Audio token ID should be 258881")
        XCTAssertEqual(mmModel.imageTokenId, 258882, "Image token ID should be 258882")
        print("  ✓ Token IDs correct\n")
        
        print("Step 3: Test text generation...")
        let inputTokens = [1, 2, 3]
        let generated = try mmModel.generateText(tokens: inputTokens, maxTokens: 5)
        print("  Input: \(inputTokens)")
        print("  Generated: \(generated)")
        XCTAssertGreaterThan(generated.count, inputTokens.count, "Should generate tokens")
        print("  ✓ Text generation working\n")
        
        print("═══════════════════════════════════════")
        print("✓ Multimodal structure test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testE4BVisionTowerLoading() throws {
        print("\n═══════════════════════════════════════")
        print("      E4B VisionTower Loading Test")
        print("═══════════════════════════════════════\n")

        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"

        print("Step 1: Load multimodal model (E4B-MarkBase)...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        print("  ✓ Multimodal model loaded\n")

        print("Step 2: Verify tower loading...")
        if let vt = mmModel.visionTowerFull {
            print("  ✓ Full VisionTower loaded: \(vt.config.numHiddenLayers) layers, hidden=\(vt.config.hiddenSize)")
        } else {
            print("  ✗ Full VisionTower NOT loaded")
        }
        if let vt12 = mmModel.visionTower {
            print("  ✓ VisionTower12B loaded: hidden=\(vt12.config.hiddenDim)")
        }
        if let at = mmModel.audioTower {
            print("  ✓ AudioTower12B loaded: outputDim=\(at.config.outputDim)")
        }

        print("  VisionTowerFull loaded: \(mmModel.visionTowerFull != nil)")
        print("  Text model: hiddenSize=\(mmModel.textModel.hiddenSize), layers=\(mmModel.textModel.numHiddenLayers)")

        XCTAssertNotNil(mmModel.visionTowerFull, "Full VisionTower should be loaded from E4B-MarkBase")
        XCTAssertEqual(mmModel.textModel.hiddenSize, 2560, "E4B should have hidden_size=2560")
        print("\n═══════════════════════════════════════")
        print("✓ E4B VisionTower loading test passed")
        print("═══════════════════════════════════════\n")
    }

    func testMultimodalPlaceholder() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Multimodal Placeholder Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Note: 12B multimodal implementation is simplified placeholder")
        print("  - Vision: vision_embedder + embed_vision.embedding_projection")
        print("  - Audio: embed_audio.embedding_projection only")
        print("  - Full tower implementation pending\n")
        
        print("Step 1: Load model...")
        let engine = try MarkBaseEngine(autoCompile: true)
        let mmModel = try MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        print("  ✓ Model loaded\n")
        
        print("Step 2: Verify multimodal structure...")
        print("  Hidden size: \(mmModel.textModel.hiddenSize)")
        print("  Vocab size: \(mmModel.textModel.vocabSize)")
        print("  Layers: \(mmModel.textModel.numHiddenLayers)")
        print("  ✓ Structure verified\n")
        
        print("═══════════════════════════════════════")
        print("✓ Placeholder test passed")
        print("Note: Audio/vision processing requires future work")
        print("═══════════════════════════════════════\n")
    }
}