import XCTest
@testable import MarkBase

final class MoEDebugTests: XCTestCase {
    
    func test26BA4BRouterStructure() throws {
        print("\n═══════════════════════════════════════")
        print("      26B-A4B MoE Router Structure Test")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        let engine = try MarkBaseEngine(autoCompile: true)
        
        print("Step 1: Load model...")
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 128)
        print("  ✓ Model loaded (30 layers, MoE enabled)\n")
        
        print("Step 2: Check Layer 0 router structure...")
        let layer0 = model.layers[0]
        
        print("  Layer 0 MoE status:")
        print("    useMoE: \(layer0.useMoE)")
        
        if !layer0.useMoE {
            XCTFail("Layer 0 should use MoE")
            return
        }
        
        print("    ✓ MoE is enabled")
        
        // Check router components
        print("\n  Router components:")
        print("    routerProj: \(layer0.routerProj != nil ? "present" : "missing")")
        print("    routerScale: \(layer0.routerScale)")
        print("    perExpertScale: \(layer0.perExpertScale != nil ? "present" : "missing")")
        print("    topK: \(layer0.topK)")
        
        // Check expert components
        print("\n  Expert components:")
        print("    expertGate: \(layer0.expertGate != nil ? "present" : "missing")")
        print("    expertUp: \(layer0.expertUp != nil ? "present" : "missing")")
        print("    expertDown: \(layer0.expertDown != nil ? "present" : "missing")")
        
        if let gate = layer0.expertGate {
            print("\n  ExpertGate details:")
            print("    numExperts: \(gate.numExperts)")
            print("    expertOutDim: \(gate.expertOutDim)")
            print("    expertInDim: \(gate.expertInDim)")
            print("    bits: \(gate.bits)")
            print("    weightStride: \(gate.weightStride)")
            print("    scalesStride: \(gate.scalesStride)")
            
            XCTAssertEqual(gate.numExperts, 128, "Should have 128 experts")
            XCTAssertEqual(gate.expertOutDim, 704, "MoE intermediate should be 704")
            XCTAssertEqual(gate.expertInDim, 2816, "Expert input should be hidden_size")
        }
        
        if let router = layer0.routerProj {
            print("\n  RouterProj details:")
            print("    bits: \(router.bits)")
            print("    inDim: \(router.inDim)")
            print("    outDim: \(router.outDim)")
            print("    groupSize: \(router.groupSize)")
            
            XCTAssertEqual(router.outDim, 128, "Router should output 128 scores")
            XCTAssertEqual(router.inDim, 2816, "Router input should be hidden_size")
        }
        
        print("\n  ✓ All router and expert components present")
        
        print("\n═══════════════════════════════════════")
        print("✓ Router structure test passed")
        print("═══════════════════════════════════════\n")
    }
    
    func test26BA4BSimpleGenerationDebug() throws {
        print("\n═══════════════════════════════════════")
        print("      26B-A4B Simple Generation Debug")
        print("═══════════════════════════════════════\n")
        
        let modelDir = "/Users/accusys/MarkBaseEngine/models/gemma-4-26b-a4b-it-4bit"
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 64)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        print("Model and tokenizer loaded")
        
        // Create minimal generator
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
        
        print("\nAttempting generation with 1 token...")
        let prompt = "Hi"
        let config = GenerationConfig(maxTokens: 1, temperature: Float(0.7))
        
        print("  Prompt: \"\(prompt)\"")
        print("  Max tokens: 1")
        print("  Temperature: 0.7")
        
        let start = Date()
        
        do {
            let response = try generator.generateComplete(prompt: prompt, config: config)
            let elapsed = Date().timeIntervalSince(start)
            
            print("\n  ✓ Generation completed in \(String(format: "%.3f", elapsed))s")
            print("  Response: \"\(response)\"")
            print("  Speed: \(String(format: "%.2f", 1.0 / elapsed)) tok/s")
            
            XCTAssertGreaterThan(response.count, 0, "Should generate something")
            
        } catch {
            print("\n  ❌ Generation failed with error: \(error)")
            print("  Error type: \(type(of: error))")
            
            // Print more details
            if let e4bError = error as? E4BError {
                print("  E4BError description: \(e4bError.localizedDescription)")
            }
            
            throw error
        }
        
        print("\n═══════════════════════════════════════")
        print("✓ Simple generation test completed")
        print("═══════════════════════════════════════\n")
    }
}