import Foundation
import MarkBase

// ─────────────────────────────────────────────────────────────
// Performance Benchmarking Tool
// ─────────────────────────────────────────────────────────────

public struct PerformanceBenchmark {
    private let modelDir: String
    private let modelName: String
    
    public init(modelDir: String, modelName: String = "markbase") {
        self.modelDir = modelDir
        self.modelName = modelName
    }
    
    /// Run all benchmarks
    public mutating func run() async throws {
        print("""
        
        ╔══════════════════════════════════════╗
        ║   Performance Benchmark              ║
        ║   Model: \(modelName)                   ║
        ╚══════════════════════════════════════╝
        
        """)
        
        try await benchmarkModelLoading()
        try await benchmarkTokenGeneration()
        try await benchmarkTokenizer()
        try await benchmarkBufferPool()
        
        print("\n✅ All benchmarks completed!")
    }
    
    // ─────────────────────────────────────────────────────────────
    // Model Loading Benchmark
    // ─────────────────────────────────────────────────────────────
    
    private mutating func benchmarkModelLoading() async throws {
        print("\n📊 Model Loading Benchmark")
        print(String(repeating: "─", count: 40))
        
        let start = Date()
        let engine = try MarkBaseEngine(autoCompile: true)
        let loadEngineTime = Date().timeIntervalSince(start)
        print("  Engine initialization: \(String(format: "%.3f", loadEngineTime))s")
        
        let start2 = Date()
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let loadModelTime = Date().timeIntervalSince(start2)
        print("  Model loading: \(String(format: "%.3f", loadModelTime))s")
        
        print("  Total: \(String(format: "%.3f", loadEngineTime + loadModelTime))s")
        print("  Layers: \(model.numHiddenLayers)")
        print("  Vocab: \(model.vocabSize)")
        print("  Hidden: \(model.hiddenSize)")
    }
    
    // ─────────────────────────────────────────────────────────────
    // Token Generation Benchmark
    // ─────────────────────────────────────────────────────────────
    
    private mutating func benchmarkTokenGeneration() async throws {
        print("\n📊 Token Generation Benchmark")
        print(String(repeating: "─", count: 40))
        
        let engine = try MarkBaseEngine(autoCompile: true)
        let model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: 512)
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        let generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
        let sampler = Sampler()
        
        // Test different temperatures
        print("\nTesting with different temperatures:")
        
        for temp in [Float(0.0), Float(0.7), Float(1.0)] {
            let prompt = "Hello, how are you?"
            let config = GenerationConfig(maxTokens: 20, temperature: temp)
            
            print("\n  Temperature \(temp):")
            let response = try generator.generateComplete(prompt: prompt, config: config)
            print("    Generated: \"\(response)\"")
        }
        
        // Now run benchmark with temperature=0.7
        let prompt = "Hello, how are you?"
        let config = GenerationConfig(maxTokens: 20, temperature: Float(0.7))
        
        // Benchmark
        let numRuns = 3
        var totalTokens = 0
        var totalTime: TimeInterval = 0
        
        for i in 1...numRuns {
            let start = Date()
            let response = try generator.generateComplete(prompt: prompt, config: config)
            let elapsed = Date().timeIntervalSince(start)
            
            let tokens = tokenizer.encode(text: response).count
            totalTokens += tokens
            totalTime += elapsed
            
            print("  Run \(i): \(tokens) tokens in \(String(format: "%.3f", elapsed))s (\(String(format: "%.1f", Double(tokens) / elapsed)) tok/s)")
            if i == 1 {
                print("  Generated text: \"\(response)\"")
            }
        }
        
        let avgTokens = totalTokens / numRuns
        let avgTime = totalTime / Double(numRuns)
        let avgSpeed = Double(avgTokens) / avgTime
        
        print("\n  Average: \(avgTokens) tokens in \(String(format: "%.3f", avgTime))s")
        print("  Speed: \(String(format: "%.1f", avgSpeed)) tok/s")
    }
    
    // ─────────────────────────────────────────────────────────────
    // Tokenizer Benchmark
    // ─────────────────────────────────────────────────────────────
    
    private func benchmarkTokenizer() throws {
        print("\n📊 Tokenizer Benchmark")
        print(String(repeating: "─", count: 40))
        
        let tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        
        let testTexts = [
            "Hello, world!",
            "The quick brown fox jumps over the lazy dog.",
            "Swift is a powerful and intuitive programming language for macOS, iOS, watchOS, and tvOS.",
            "Artificial intelligence is transforming the world in unprecedented ways."
        ]
        
        // Encode benchmark
        let numIterations = 100
        var totalEncodeTime: TimeInterval = 0
        
        for text in testTexts {
            let start = Date()
            for _ in 0..<numIterations {
                _ = tokenizer.encode(text: text)
            }
            let elapsed = Date().timeIntervalSince(start)
            totalEncodeTime += elapsed
            
            let tokens = tokenizer.encode(text: text).count
            print("  Encode: \"\(text.prefix(30))...\" -> \(tokens) tokens in \(String(format: "%.3f", elapsed / Double(numIterations) * 1000))ms")
        }
        
        // Decode benchmark
        var totalDecodeTime: TimeInterval = 0
        
        for text in testTexts {
            let tokens = tokenizer.encode(text: text)
            let start = Date()
            for _ in 0..<numIterations {
                _ = tokenizer.decode(tokens: tokens)
            }
            let elapsed = Date().timeIntervalSince(start)
            totalDecodeTime += elapsed
            
            print("  Decode: \(tokens.count) tokens -> \"\(text.prefix(30))...\" in \(String(format: "%.3f", elapsed / Double(numIterations) * 1000))ms")
        }
        
        print("\n  Total encode time: \(String(format: "%.3f", totalEncodeTime))s")
        print("  Total decode time: \(String(format: "%.3f", totalDecodeTime))s")
    }
    
    // ─────────────────────────────────────────────────────────────
    // Buffer Pool Benchmark
    // ─────────────────────────────────────────────────────────────
    
    private func benchmarkBufferPool() throws {
        print("\n📊 Buffer Pool Benchmark")
        print(String(repeating: "─", count: 40))
        
        let engine = try MarkBaseEngine()
        let pool = engine.bufferPool
        
        let bufferSize = 1024
        let numIterations = 1000
        
        // Without pool
        let start1 = Date()
        for _ in 0..<numIterations {
            let _ = try engine.makeBuffer(length: bufferSize)
        }
        let withoutPoolTime = Date().timeIntervalSince(start1)
        print("  Without pool: \(String(format: "%.3f", withoutPoolTime))s (\(String(format: "%.0f", Double(numIterations) / withoutPoolTime)) allocs/s)")
        
        // With pool
        let start2 = Date()
        for _ in 0..<numIterations {
            let buf = engine.acquireBuffer(length: bufferSize)
            engine.releaseBuffer(buf)
        }
        let withPoolTime = Date().timeIntervalSince(start2)
        print("  With pool: \(String(format: "%.3f", withPoolTime))s (\(String(format: "%.0f", Double(numIterations) / withPoolTime)) allocs/s)")
        
        let speedup = withoutPoolTime / withPoolTime
        print("\n  Speedup: \(String(format: "%.1f", speedup))x")
        print("  Pool stats:")
        print("  \(pool.stats)")
    }
}