import XCTest
@testable import MarkBase

final class G12BLoadingTests: XCTestCase {
    
    func testModelConfigLoading() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Model Config Loading Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Loading config from: \(modelDir)")
        let cfg = try ModelConfig.load(from: modelDir)
        
        print("\nConfig loaded successfully:")
        print("  hidden_size: \(cfg.hiddenSize ?? 0)")
        print("  num_layers: \(cfg.numHiddenLayers ?? 0)")
        print("  num_heads: \(cfg.numAttentionHeads ?? 0)")
        print("  num_kv_heads: \(cfg.numKeyValueHeads ?? 0)")
        print("  intermediate_size: \(cfg.intermediateSize ?? 0)")
        print("  head_dim: \(cfg.headDim ?? 0)")
        print("  vocab_size: \(cfg.vocabSize ?? 0)")
        print("  num_kv_shared: \(cfg.numKvSharedLayers ?? 0)")
        print("  sliding_window: \(cfg.slidingWindow ?? 0)")
        print("  max_position: \(cfg.maxPositionEmbeddings ?? 0)")
        
        // Verify 12B parameters
        XCTAssertEqual(cfg.hiddenSize, 3840, "12B should have hidden_size=3840")
        XCTAssertEqual(cfg.numHiddenLayers, 48, "12B should have 48 layers")
        XCTAssertEqual(cfg.numAttentionHeads, 16, "12B should have 16 heads")
        XCTAssertEqual(cfg.numKeyValueHeads, 8, "12B should have 8 kv_heads")
        XCTAssertEqual(cfg.numKvSharedLayers, 0, "12B should have NO KV sharing")
        XCTAssertEqual(cfg.slidingWindow, 1024, "12B sliding_window should be 1024")
        
        print("\n═══════════════════════════════════════")
        print("✓ Config validation passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testShardDetection() throws {
        print("\n═══════════════════════════════════════")
        print("        12B Shard Detection Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        // Check for index file
        let indexPath = "\(modelDir)/model.safetensors.index.json"
        let hasIndex = FileManager.default.fileExists(atPath: indexPath)
        print("Index file exists: \(hasIndex) ✓")
        
        if hasIndex {
            // Load index
            let index = try SafeTensorsIndex(modelDir: modelDir)
            print("\nShard files found:")
            for shard in index.shardFiles.sorted() {
                print("  - \(shard)")
            }
            
            print("\nTotal tensors: \(index.allTensors.count)")
            print("Sample tensors:")
            let sampleTensors = index.allTensors.prefix(5)
            for tensor in sampleTensors {
                if let shard = index.weightMap[tensor] {
                    print("  \(tensor) → \(shard)")
                }
            }
            
            XCTAssertGreaterThan(index.shardFiles.count, 1, "12B should have multiple shards")
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ Shard detection passed")
        print("═══════════════════════════════════════\n")
    }
}