import XCTest
@testable import MarkBase

final class G12BOptimizationTests: XCTestCase {

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

    func testBufferPoolAcquireAndRelease() throws {
        let engine = try MarkBaseEngine()
        
        let buf1 = engine.acquireBuffer(length: 1024)
        XCTAssertGreaterThan(buf1.length, 0)
        
        engine.releaseBuffer(buf1)
        
        let buf2 = engine.acquireBuffer(length: 1024)
        XCTAssertEqual(buf1.length, buf2.length)
    }

    func testBufferPoolAlignment() throws {
        let engine = try MarkBaseEngine()
        
        let buf = engine.acquireBuffer(length: 100)
        XCTAssertEqual(buf.length % 256, 0)
        XCTAssertEqual(buf.length, 256)
    }

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

    func testBufferPoolStats() throws {
        let engine = try MarkBaseEngine()
        
        let buf1 = engine.acquireBuffer(length: 512)
        let buf2 = engine.acquireBuffer(length: 1024)
        engine.releaseBuffer(buf1)
        let buf3 = engine.acquireBuffer(length: 512)
        engine.releaseBuffer(buf2)
        engine.releaseBuffer(buf3)
        
        let stats = engine.bufferPool.stats
        XCTAssertTrue(stats.contains("Allocations:"))
    }

    // ── 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)
        }
    }

    func testMakeBatchEncoder() throws {
        let engine = try MarkBaseEngine()
        try engine.compileSource(MetalKernels.vectorAdd)
        let pipeline = try engine.pipeline(named: "vector_add")
        
        let (cmdBuf, enc) = try engine.makeBatchEncoder()
        
        enc.setComputePipelineState(pipeline)
        let n: UInt = 64
        let a = (0..<n).map { Float($0) }
        let b = (0..<n).map { Float($0) }
        
        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])
        
        enc.setBuffer(bufA, offset: 0, index: 0)
        enc.setBuffer(bufB, offset: 0, index: 1)
        enc.setBuffer(bufC, offset: 0, index: 2)
        enc.setBuffer(bufN, offset: 0, index: 3)
        
        let tg = engine.threadgroupSize1D(pipeline, count: Int(n))
        enc.dispatchThreads(MTLSize(width: Int(n), height: 1, depth: 1),
                           threadsPerThreadgroup: tg)
        enc.endEncoding()
        
        cmdBuf.commit()
        cmdBuf.waitUntilCompleted()
        
        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)
        }
    }

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

    func testRMSNormChunkedKernel() throws {
        let engine = try MarkBaseEngine(autoCompile: true)
        
        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)
        
        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)
        }
    }

    // ── HuggingFace Tokenizer Tests ───────────────────

    func testHuggingFaceTokenizerRoundtrip() throws {
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/main"
        let tokenizerPath = modelDir + "/tokenizer.json"
        
        guard FileManager.default.fileExists(atPath: tokenizerPath) else {
            print("Skipping HuggingFaceTokenizer test - tokenizer.json not found")
            return
        }
        
        let tokenizer = try HuggingFaceTokenizer(jsonPath: tokenizerPath)
        
        // Verify vocab
        XCTAssertGreaterThan(tokenizer.vocabSize, 0)
        
        // Test roundtrip
        let text = "Hello, world! This is a test."
        let tokens = tokenizer.encode(text: text)
        XCTAssertGreaterThan(tokens.count, 1)  // At least BOS + tokens
        
        let decoded = tokenizer.decode(tokens: tokens)
        XCTAssertFalse(decoded.isEmpty)
        
        print("Original: \(text)")
        print("Tokens: \(tokens.count)")
        print("Decoded: \(decoded)")
    }

    func testHuggingFaceTokenizerSpecialTokens() throws {
        let modelDir = "/Users/accusys/.cache/huggingface/hub/models--mlx-community--gemma-4-12B-it-4bit/snapshots/main"
        let tokenizerPath = modelDir + "/tokenizer.json"
        
        guard FileManager.default.fileExists(atPath: tokenizerPath) else {
            return
        }
        
        let tokenizer = try HuggingFaceTokenizer(jsonPath: tokenizerPath)
        
        // Verify special token IDs
        XCTAssertGreaterThanOrEqual(tokenizer.bosTokenId, 0)
        XCTAssertGreaterThanOrEqual(tokenizer.eosTokenId, 0)
        XCTAssertGreaterThanOrEqual(tokenizer.padTokenId, 0)
    }

    // ── 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)
        XCTAssertEqual(config.rotatedDim, 128)  // 512 * 0.25
    }

    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")
    }
}