import XCTest
@testable import MarkBase

final class G12BBufferTests: XCTestCase {
    
    func testForwardTempsAllocation() throws {
        print("\n═══════════════════════════════════════")
        print("    ForwardTemps Allocation Test")
        print("═══════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine()
        
        // Test E4B parameters (default)
        print("Test 1: E4B parameters (2560 hidden, 8 heads, 2 kv_heads)")
        do {
            let tempsE4B = try ForwardTemps(device: engine.device,
                                            maxHeadDim: 256,
                                            maxIntermediate: 20480,
                                            hiddenSize: 2560,
                                            nHeads: 8,
                                            nKvHeads: 2)
            print("  ✓ E4B ForwardTemps allocated successfully")
            print("    q: \(tempsE4B.q.length) bytes")
            print("    gate: \(tempsE4B.gate.length) bytes")
        } catch {
            print("  ✗ E4B ForwardTemps failed: \(error)")
            throw error
        }
        
        // Test 12B parameters
        print("\nTest 2: 12B parameters (3840 hidden, 16 heads, 8 kv_heads)")
        do {
            let temps12B = try ForwardTemps(device: engine.device,
                                            maxHeadDim: 256,
                                            maxIntermediate: 15360,
                                            hiddenSize: 3840,
                                            nHeads: 16,
                                            nKvHeads: 8)
            print("  ✓ 12B ForwardTemps allocated successfully")
            print("    q: \(temps12B.q.length) bytes")
            print("    gate: \(temps12B.gate.length) bytes")
            
            // Verify sizes
            XCTAssertEqual(temps12B.q.length, 16 * 256 * 4, "q should be 16 heads × 256 dim")
            XCTAssertEqual(temps12B.gate.length, 15360 * 4, "gate should be 15360")
            
        } catch {
            print("  ✗ 12B ForwardTemps failed: \(error)")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ ForwardTemps allocation passed")
        print("═══════════════════════════════════════\n")
    }
    
    func testKVCacheAllocation() throws {
        print("\n═══════════════════════════════════════")
        print("      KVCache Allocation Test")
        print("═══════════════════════════════════════\n")
        
        let engine = try MarkBaseEngine()
        
        // Test single KV cache
        print("Test: Single KV cache (512 positions, 8 kv_heads, 256 head_dim)")
        do {
            let cache = try KVCache(device: engine.device,
                                    isSliding: true,
                                    maxContextLength: 512,
                                    nKvHeads: 8,
                                    headDim: 256)
            print("  ✓ KV cache allocated successfully")
            print("    Buffer size: \(cache.buffer.length) bytes = \(cache.buffer.length / 1024)KB")
        } catch {
            print("  ✗ KV cache failed: \(error)")
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ KVCache allocation passed")
        print("═══════════════════════════════════════\n")
    }
}