import XCTest
@testable import MarkBase

final class VisionSeparateTest: XCTestCase {
    
    func testE2BVisionLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  E2B Vision Load Test (VisionTowerE2B)")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-e2b-it-4bit/snapshots/2c3e507453b4f218d05fe3cc97bea5c5a654257e"
        
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        let vcfg = loadVisionConfig(modelDir: modelDir)
        guard let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") else {
            print("  ✗ No safetensors file")
            throw NSError(domain: "Vision", code: -1, userInfo: [NSLocalizedDescriptionKey: "No safetensors"])
        }
        
        print("  Loading VisionTowerE2B (bfloat16 weights)...")
        let start = Date()
        let visionTower = try loadVisionTowerE2B(reader: reader, config: vcfg, engine: engine)
        let loadTime = Date().timeIntervalSince(start) * 1000
        
        print("  ✓ Loaded in \(loadTime) ms")
        print("  Type: VisionTowerE2B (\(visionTower.config.numHiddenLayers) layers)")
        print("  Config: hiddenSize=\(visionTower.config.hiddenSize)")
        print("  Layers: \(visionTower.weights.layers.count)")
        print("\n✓ E2B vision test passed\n")
    }
    
    func testE4BVisionLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  E4B Vision Load Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/E4B-MarkBase"
        
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        let vcfg = loadVisionConfig(modelDir: modelDir)
        guard let reader = try? SafeTensorsReader(path: modelDir + "/model.safetensors") else {
            print("  ✗ No safetensors file")
            return
        }
        
        let start = Date()
        let visionTower = try loadVisionTower(reader: reader, config: vcfg, engine: engine)
        let loadTime = Date().timeIntervalSince(start) * 1000
        
        print("  ✓ Loaded in \(loadTime) ms")
        print("  hiddenSize=\(visionTower.config.hiddenSize), layers=\(visionTower.config.numHiddenLayers)")
        print("\n✓ E4B vision test passed\n")
    }
    
    func test12BVisionLoad() throws {
        print("\n═══════════════════════════════════════")
        print("  12B Vision Load Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine created")
        
        let start = Date()
        let visionTower = try VisionTower12B.load(modelDir: modelDir, engine: engine)
        let loadTime = Date().timeIntervalSince(start) * 1000
        
        print("  ✓ Loaded in \(loadTime) ms")
        print("  hiddenDim=\(visionTower.hiddenDim), patchDim=\(visionTower.patchDim)")
        print("\n✓ 12B vision test passed\n")
    }
}