import XCTest
@testable import MarkBase

final class G12BCompleteTests: XCTestCase {

    // ── Buffer Pool Tests ──────────────────────────────

    func testBufferPoolAcquireAndRelease() throws {
        let engine = try MarkBaseEngine()
        
        // Acquire a buffer
        let buf1 = engine.acquireBuffer(length: 1024)
        XCTAssertGreaterThan(buf1.length, 0)
        
        // Release back to pool
        engine.releaseBuffer(buf1)
        
        // Acquire again - should reuse
        let buf2 = engine.acquireBuffer(length: 1024)
        XCTAssertEqual(buf1.length, buf2.length)
    }

    func testBufferPoolAlignment() throws {
        let engine = try MarkBaseEngine()
        
        // Request non-aligned size
        let buf = engine.acquireBuffer(length: 100)
        
        // Should be aligned to 256 bytes
        XCTAssertEqual(buf.length % 256, 0)
        XCTAssertEqual(buf.length, 256)
    }

    func testBufferPoolReuse() throws {
        let engine = try MarkBaseEngine()
        
        // Acquire and release multiple buffers
        for _ in 0..<10 {
            let buf = engine.acquireBuffer(length: 2048)
            engine.releaseBuffer(buf)
        }
        
        // Should have reused buffers
        XCTAssertTrue(engine.bufferPool.totalReuses >= 9, 
                      "Expected at least 9 reuses, got \(engine.bufferPool.totalReuses)")
    }

    // ── Batch Dispatch Tests ──────────────────────────

    func testBatchDispatch() throws {
        let engine = try MarkBaseEngine()
        try engine.compileSource(MetalKernels.vectorAdd)
        let pipeline = try engine.pipeline(named: "vector_add")
        
        let n: UInt = 128
        let a = (0..<n).map { Float($0) }
        let b = (0..<n).map { Float($0 * 2) }
        
        let bufA = try engine.makeBuffer(a)
        let bufB = try engine.makeBuffer(b)
        let bufC = try engine.makeBuffer(length: Int(n) * MemoryLayout<Float>.stride)
        let bufN = try engine.makeBuffer([n])
        
        let gridSize = MTLSize(width: Int(n), height: 1, depth: 1)
        
        try engine.batchDispatch([
            (pipeline, [bufA, bufB, bufC, bufN], gridSize)
        ])
        
        let result = engine.readFloats(from: bufC, count: Int(n))
        for i in 0..<Int(n) {
            XCTAssertEqual(result[i], a[i] + b[i], accuracy: 1e-5, "at index \(i)")
        }
    }

    // ── RMS Norm Chunked Tests ────────────────────────

    func testRMSNormChunkedKernel() throws {
        let engine = try MarkBaseEngine(autoCompile: true)
        
        // Test with size larger than threadgroup (256)
        let size = 3840  // 12B hidden size
        
        let input: [Float] = (0..<size).map { Float($0) / Float(size) }
        let weight: [Float] = Array(repeating: 1.0, count: size)
        
        let inputBuf = try engine.makeBuffer(input)
        let weightBuf = try engine.makeBuffer(weight)
        let outputBuf = try engine.makeBuffer(length: size * MemoryLayout<Float>.stride)
        
        let pso = try engine.pipeline(named: "rms_norm_chunked")
        let cmdBuf = engine.commandQueue.makeCommandBuffer()!
        let enc = cmdBuf.makeComputeCommandEncoder()!
        
        enc.setComputePipelineState(pso)
        enc.setBuffer(inputBuf, offset: 0, index: 0)
        enc.setBuffer(weightBuf, offset: 0, index: 1)
        enc.setBuffer(outputBuf, offset: 0, index: 2)
        var N = UInt32(size)
        enc.setBytes(&N, length: MemoryLayout<UInt32>.size, index: 3)
        var eps: Float = 1e-6
        enc.setBytes(&eps, length: MemoryLayout<Float>.size, index: 4)
        enc.setThreadgroupMemoryLength(256 * 4, index: 0)
        
        let tg = MTLSize(width: 256, height: 1, depth: 1)
        enc.dispatchThreads(MTLSize(width: size, height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()
        
        let result = engine.readFloats(from: outputBuf, count: size)
        
        // Verify RMS normalization
        let sumSq = input.reduce(0) { $0 + $1 * $1 }
        let rms = sqrt(sumSq / Float(size) + 1e-6)
        
        for i in 0..<size {
            let expected = input[i] / rms
            XCTAssertEqual(result[i], expected, accuracy: 1e-3, "at index \(i)")
        }
    }

    // ── Layer Config Tests ────────────────────────────

    func test12BLayerConfig() {
        // 12B should use nHeads=16, nKvHeads=8
        let config = E4BLayerConfig.full(
            hiddenSize: 3840,
            headDim: 512,
            intermediateSize: 30720,
            nHeads: 16,
            nKvHeads: 8
        )
        
        XCTAssertEqual(config.nHeads, 16)
        XCTAssertEqual(config.nKvHeads, 8)
        XCTAssertEqual(config.hiddenSize, 3840)
        XCTAssertEqual(config.headDim, 512)
        XCTAssertFalse(config.isSliding)
    }

    func test12BLayerConfigSliding() {
        let config = E4BLayerConfig.sliding(
            hiddenSize: 3840,
            headDim: 256,
            intermediateSize: 15360,
            nHeads: 16,
            nKvHeads: 8,
            windowSize: 512
        )
        
        XCTAssertEqual(config.nHeads, 16)
        XCTAssertEqual(config.nKvHeads, 8)
        XCTAssertTrue(config.isSliding)
        XCTAssertEqual(config.windowSize, 512)
    }

    // ── ForwardTemps 12B Tests ────────────────────────

    func testForwardTemps12B() throws {
        let engine = try MarkBaseEngine()
        
        // Test 12B parameters
        let temps = try ForwardTemps(
            device: engine.device,
            maxHeadDim: 512,
            maxIntermediate: 30720,  // 15360 * 2
            hiddenSize: 3840,
            nHeads: 16,
            nKvHeads: 8
        )
        
        // Verify sizes
        XCTAssertEqual(temps.q.length, 16 * 512 * 4, "q should be 16 heads × 512 dim")
        XCTAssertEqual(temps.gate.length, 30720 * 4, "gate should be 30720")
        XCTAssertEqual(temps.io.length, 3840 * 4, "io should be 3840")
    }
}