import XCTest
@testable import MarkBase

final class RouterInputBufferSyncTest: XCTestCase {
    
    func testRouterInputBufferSync() throws {
        print("\n═══════════════════════════════════════")
        print("      Router Input Buffer Sync Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        
        print("Loading model...")
        fflush(stdout)
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 32)
        let layer0 = model.layers[0]
        
        print("\nSimulating forward pass timing...")
        fflush(stdout)
        
        let h = model.temps.h
        let ns = model.temps.ns
        let hs = model.hiddenSize
        
        // Create input
        let inputBuffer = engine.device.makeBuffer(length: hs * MemoryLayout<Float>.size)!
        var inputData = [Float](repeating: 0.1, count: hs)
        memcpy(inputBuffer.contents(), &inputData, inputData.count * MemoryLayout<Float>.size)
        
        // Create a command buffer (simulating main cmdBuf)
        let mainCmdBuf = engine.commandQueue.makeCommandBuffer()!
        
        // Step 1: input_layernorm → temps.h (NOT COMMITTED)
        print("  Step 1: input_layernorm → temps.h (NOT COMMITTED)")
        fflush(stdout)
        try layer0.rmsNorm(engine: engine, cmdBuf: mainCmdBuf,
                          input: inputBuffer, weight: layer0.inputLayernorm,
                          output: h, count: hs, eps: 1e-6)
        
        // Check h before commit
        let hBeforeCommit = engine.readFloats(from: h, count: 10)
        print("    h before commit: \(hBeforeCommit)")
        fflush(stdout)
        
        // Step 2: Many other operations on mainCmdBuf (NOT COMMITTED)
        // (simulate attention, etc.)
        print("  Step 2: Other operations (NOT COMMITTED)")
        fflush(stdout)
        
        // Step 3: pre_feedforward_layernorm → temps.ns (NOT COMMITTED)
        print("  Step 3: pre_feedforward_layernorm → temps.ns (NOT COMMITTED)")
        fflush(stdout)
        try layer0.rmsNorm(engine: engine, cmdBuf: mainCmdBuf,
                          input: h, weight: layer0.preFeedforwardLayernorm,
                          output: ns, count: hs, eps: 1e-6)
        
        // Check ns before commit
        let nsBeforeCommit = engine.readFloats(from: ns, count: 10)
        print("    ns before commit: \(nsBeforeCommit)")
        fflush(stdout)
        
        // Step 4: Router matmul (separate routerCmdBuf)
        print("  Step 4: Router matmul (separate routerCmdBuf)")
        fflush(stdout)
        
        guard let router = layer0.routerProj else {
            XCTFail("routerProj is nil")
            return
        }
        
        let routerCmdBuf = engine.commandQueue.makeCommandBuffer()!
        try layer0.quantizedMatmul(engine: engine, cmdBuf: routerCmdBuf,
                                   input: ns, weights: router,
                                   output: model.temps.gate)
        routerCmdBuf.commit()
        routerCmdBuf.waitUntilCompleted()
        
        // Check router output
        let routerOutput = engine.readFloats(from: model.temps.gate, count: router.outDim)
        print("    Router output (using ns before main commit): \(routerOutput[0..<min(10, router.outDim)])")
        print("    Router output max/min: \(routerOutput.max() ?? 0), \(routerOutput.min() ?? 0)")
        fflush(stdout)
        
        // Step 5: NOW commit mainCmdBuf
        print("  Step 5: Commit mainCmdBuf")
        fflush(stdout)
        mainCmdBuf.commit()
        mainCmdBuf.waitUntilCompleted()
        
        // Check ns after commit
        let nsAfterCommit = engine.readFloats(from: ns, count: 10)
        print("    ns after commit: \(nsAfterCommit)")
        fflush(stdout)
        
        // Step 6: Router matmul again (using properly populated ns)
        print("  Step 6: Router matmul again (using populated ns)")
        fflush(stdout)
        
        let routerCmdBuf2 = engine.commandQueue.makeCommandBuffer()!
        try layer0.quantizedMatmul(engine: engine, cmdBuf: routerCmdBuf2,
                                   input: ns, weights: router,
                                   output: model.temps.gate)
        routerCmdBuf2.commit()
        routerCmdBuf2.waitUntilCompleted()
        
        let routerOutput2 = engine.readFloats(from: model.temps.gate, count: router.outDim)
        print("    Router output (using ns after main commit): \(routerOutput2[0..<min(10, router.outDim)])")
        print("    Router output max/min: \(routerOutput2.max() ?? 0), \(routerOutput2.min() ?? 0)")
        fflush(stdout)
        
        print("\n═══════════════════════════════════════")
        print("✓ Router input buffer sync test completed")
        print("═══════════════════════════════════════\n")
        fflush(stdout)
    }
}