import XCTest
@testable import MarkBase

final class G12BGenerationTests: XCTestCase {
    
    func testSimpleForwardPass() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Forward Pass Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Initialize Metal engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine ready (shaders compiled)\n")
        
        print("Step 2: Load 12B model...")
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        print("  ✓ Model loaded\n")
        
        print("Step 3: Create input buffer...")
        // Create simple input buffer with all zeros (dummy token embedding)
        guard let inputBuffer = engine.device.makeBuffer(
            length: model.hiddenSize * MemoryLayout<Float>.stride,
            options: .storageModeShared
        ) else {
            throw E4BError.bufferCreationFailed
        }
        // Initialize with small values to avoid numerical issues
        let inputPtr = inputBuffer.contents().bindMemory(to: Float.self, capacity: model.hiddenSize)
        for i in 0..<model.hiddenSize {
            inputPtr[i] = 0.001 // Small non-zero value
        }
        print("  ✓ Input buffer created (size: \(model.hiddenSize))\n")
        
        print("Step 4: Run single forward pass...")
        do {
            // Try forward pass through layer 0
            let layer = model.layers[0]
            let kvCache = model.kvCaches[0]
            
            print("  Running layer 0 forward...")
            try layer.forward(
                input: inputBuffer,
                position: 0,
                kvCache: kvCache,
                shouldStoreKV: true,
                temps: model.temps,
                engine: engine
            )
            print("  ✓ Layer 0 forward passed\n")
            
            print("═══════════════════════════════════════")
            print("✓ Forward pass test passed")
            print("═══════════════════════════════════════\n")
            
        } catch {
            print("  ✗ Forward pass failed: \(error)\n")
            throw error
        }
    }
    
    func testFullModelForward() throws {
        print("\n═══════════════════════════════════════")
        print("      12B Full Model Forward Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/73bcf09092aa277861d5a191b989b666f7f32e8f"
        
        print("Step 1: Initialize Metal engine...")
        let engine = try MarkBaseEngine(autoCompile: true)
        print("  ✓ Engine ready (shaders compiled)\n")
        
        print("Step 2: Load 12B model...")
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        print("  ✓ Model loaded\n")
        
        print("Step 3: Create input buffer...")
        guard let inputBuffer = engine.device.makeBuffer(
            length: model.hiddenSize * MemoryLayout<Float>.stride,
            options: .storageModeShared
        ) else {
            throw E4BError.bufferCreationFailed
        }
        let inputPtr = inputBuffer.contents().bindMemory(to: Float.self, capacity: model.hiddenSize)
        for i in 0..<model.hiddenSize {
            inputPtr[i] = 0.001
        }
        print("  ✓ Input buffer created\n")
        
        print("Step 4: Run full model forward pass...")
        do {
            // Use model.forward with token ID
            print("  Testing model.forward method...")
            let logits = try model.forward(tokenId: 0, position: 0)
            print("  ✓ Full model forward passed")
            print("  Logits count: \(logits.count)\n")
            
            // Verify logits are not all zeros
            var maxLogit: Float = -Float.infinity
            var minLogit: Float = Float.infinity
            for i in 0..<min(100, logits.count) {
                let val = logits[i]
                maxLogit = max(maxLogit, val)
                minLogit = min(minLogit, val)
            }
            print("  Logits stats (first 100): min=\(minLogit), max=\(maxLogit)\n")
            
            XCTAssertFalse(maxLogit == minLogit, "Logits should not all be the same")
            
            print("═══════════════════════════════════════")
            print("✓ Full model forward test passed")
            print("═══════════════════════════════════════\n")
            
        } catch {
            print("  ✗ Full model forward failed: \(error)\n")
            throw error
        }
    }
}