import AVFoundation
import CoreImage
import Foundation
import MarkBase
import Hummingbird

// ─────────────────────────────────────────────────────────────
// OpenAI Compatible API Server with SSE Support
// ─────────────────────────────────────────────────────────────

public final class MarkBaseServer: @unchecked Sendable {
    private let modelDir: String
    private let modelId: String
    private let maxContextLength: Int
    
    // Runtime state
    private let engine: MarkBaseEngine
    private let model: E4BModel
    private let tokenizer: Tokenizer
    private let generator: StreamingGenerator
    private let sampler: Sampler
    private let multimodalModel: MultimodalModel?
    
    public init(modelDir: String, modelId: String = "markbase-12b", maxContextLength: Int = 512) throws {
        self.modelDir = modelDir
        self.modelId = modelId
        self.maxContextLength = maxContextLength
        
        // Validate model path
        try Validator.validateModelPath(modelDir)
        
        print("Loading model from \(modelDir)...")
        engine = try MarkBaseEngine(autoCompile: true)
        model = try E4BModel(modelDir: modelDir, engine: engine, maxContextLength: maxContextLength)
        tokenizer = try TokenizerFactory.load(modelDir: modelDir)
        generator = StreamingGenerator(model: model, tokenizer: tokenizer, engine: engine)
        sampler = Sampler()
        print("✓ Model loaded: \(modelId)")
        // Skip multimodal loading for faster startup (it creates another E4BModel internally)
        multimodalModel = nil
        // multimodalModel = try? MultimodalModel(modelDir: modelDir, engine: engine, maxContextLength: maxContextLength)
        // if multimodalModel != nil { print("  ✓ Multimodal model ready") }
        print("  Layers: \(model.numHiddenLayers)")
        print("  Vocab: \(model.vocabSize)")
    }
    
    public func start(port: Int = 8080) async throws {
        print("  Port: \(port)")
        
        // Create router
        let router = Router()
        
        // Middleware
        router.middlewares.add(LogRequestsMiddleware(.info))
        router.middlewares.add(CORSMiddleware(
            allowOrigin: .all,
            allowHeaders: [.contentType, .authorization],
            allowMethods: [.get, .post, .options]
        ))
        
        // Routes
        router.get("/health") { _, _ in
            return "OK"
        }
        
        router.get("/v1/models") { _, _ in
            let models = [
                ModelDetails(
                    id: self.modelId,
                    capabilities: ModelCapabilities(
                        vision: self.multimodalModel != nil,
                        audio: self.multimodalModel?.audioTowerFull != nil
                    ),
                    parameters: ModelParameters(
                        context_length: self.maxContextLength,
                        num_hidden_layers: self.model.numHiddenLayers,
                        hidden_size: self.model.hiddenSize,
                        vocab_size: self.model.vocabSize,
                        num_attention_heads: 8,
                        num_kv_heads: 2
                    )
                )
            ]
            return models
        }
        
        router.post("/v1/chat/completions") { request, _ in
            let buffer = try await request.body.collect(upTo: .max)
            
            let req = try JSONDecoder().decode(ChatCompletionRequest.self, from: Data(buffer: buffer))
            let response = try self.handleChatCompletion(messages: req.messages, config: req.toGenerationConfig())
            
            return try ByteBuffer(data: JSONEncoder().encode(response))
        }
        
        router.post("/v1/multimodal/chat/completions") { request, _ in
            print("[DEBUG ROUTER] Received multimodal request")
            let buffer = try await request.body.collect(upTo: .max)
            print("[DEBUG ROUTER] Body collected: \(buffer.readableBytes) bytes")
            
            do {
                let req = try JSONDecoder().decode(MultimodalChatCompletionRequest.self, from: Data(buffer: buffer))
                print("[DEBUG ROUTER] Request decoded successfully, messages: \(req.messages.count)")
                let response = try self.handleMultimodalChatCompletion(messages: req.messages, config: req.toGenerationConfig())
                print("[DEBUG ROUTER] Response generated")
                return try ByteBuffer(data: JSONEncoder().encode(response))
            } catch {
                print("[DEBUG ROUTER] Error: \(error)")
                throw error
            }
        }
        
        // Create Hummingbird app
        let app = Application(
            router: router,
            configuration: .init(address: .hostname(port: port))
        )
        
        print("\nEndpoints:")
        print("  GET  /health")
        print("  GET  /v1/models")
        print("  POST /v1/chat/completions")
        print("  POST /v1/multimodal/chat/completions")
        print("\nServer starting on port \(port)...")
        
        try await app.run()
    }
    
    // ─────────────────────────────────────────────────────────────
    // Multimodal Handlers
    // ─────────────────────────────────────────────────────────────
    
    public func handleMultimodalChatCompletion(
        messages: [MultimodalMessage],
        config: GenerationConfig
    ) throws -> ChatCompletionResponse {
        guard multimodalModel != nil else {
            throw MarkBaseError.multimodalNotSupported
        }
        
        print("[DEBUG] handleMultimodalChatCompletion: Processing multimodal request")
        
        // Build multimodal prompt
        var textParts: [String] = []
        var hasImage = false
        var hasAudio = false
        var imageData: Data? = nil
        var audioData: Data? = nil
        
        for message in messages {
            print("[DEBUG] Message role: \(message.role), content parts: \(message.content.count)")
            for part in message.content {
                switch part {
                case .text(let text):
                    print("[DEBUG] Text part: \(text)")
                    textParts.append(text)
                case .imageUrl(let url):
                    print("[DEBUG] Image URL detected: \(url.url.prefix(50))...")
                    hasImage = true
                    // Load image from URL
                    if url.url.hasPrefix("data:image") {
                        // Base64 encoded
                        let parts = url.url.split(separator: ",")
                        if parts.count == 2 {
                            imageData = Data(base64Encoded: String(parts[1]))
                            print("[DEBUG] Base64 decoded, data size: \(imageData?.count ?? 0)")
                        }
                    } else if url.url.hasPrefix("file://") {
                        // Local file
                        imageData = try Data(contentsOf: URL(fileURLWithPath: String(url.url.dropFirst(7))))
                        print("[DEBUG] File loaded, data size: \(imageData?.count ?? 0)")
                    }
                case .audioUrl(let url):
                    print("[DEBUG] Audio URL detected: \(url.url.prefix(50))...")
                    hasAudio = true
                    // Load audio from URL
                    if url.url.hasPrefix("data:audio") {
                        // Base64 encoded
                        let parts = url.url.split(separator: ",")
                        if parts.count == 2 {
                            audioData = Data(base64Encoded: String(parts[1]))
                            print("[DEBUG] Audio base64 decoded, data size: \(audioData?.count ?? 0)")
                        }
                    } else if url.url.hasPrefix("file://") {
                        // Local file
                        audioData = try Data(contentsOf: URL(fileURLWithPath: String(url.url.dropFirst(7))))
                        print("[DEBUG] Audio file loaded, data size: \(audioData?.count ?? 0)")
                    }
                case .videoUrl:
                    // Video handling (future)
                    break
                }
            }
        }
        
        let prompt = textParts.joined(separator: " ")
        let promptTokens = tokenizer.encode(text: prompt)
        
        // Process image if present
        if hasImage, let data = imageData {
            print("[DEBUG] Processing image data: \(data.count) bytes")
            
            // Vision preprocessing
            let visionFeatures = try processImageData(data)
            print("[DEBUG] Vision features created, buffer length: \(visionFeatures.length)")
            
            // Generate with vision conditioning
            let response = try generateWithVision(
                textTokens: promptTokens,
                visionFeatures: visionFeatures,
                maxTokens: config.maxTokens
            )
            print("[DEBUG] Generated response: \(response.prefix(100))")
            
            return ChatCompletionResponse(
                id: generateId("chatcmpl"),
                object: "chat.completion",
                created: Int(Date().timeIntervalSince1970),
                model: modelId,
                choices: [
                    Choice(
                        index: 0,
                        message: ChatMessage(role: "assistant", content: response),
                        finish_reason: "stop"
                    )
                ],
                usage: Usage(
                    promptTokens: promptTokens.count,
                    completionTokens: tokenizer.encode(text: response).count,
                    totalTokens: promptTokens.count + tokenizer.encode(text: response).count
                )
            )
        } else if hasAudio, let data = audioData {
            print("[DEBUG] Processing audio data: \(data.count) bytes")
            
            // Audio preprocessing
            let audioFeatures = try processAudioData(data)
            print("[DEBUG] Audio features created")
            
            // Generate with audio conditioning
            // Note: Need to determine numFrames from audio length
            // For now, use placeholder
            let numFrames = 100  // Placeholder
            
            let response = try generateWithAudio(
                textTokens: promptTokens,
                audioFeatures: audioFeatures,
                numFrames: numFrames,
                maxTokens: config.maxTokens
            )
            print("[DEBUG] Generated audio response: \(response.prefix(100))")
            
            return ChatCompletionResponse(
                id: generateId("chatcmpl"),
                object: "chat.completion",
                created: Int(Date().timeIntervalSince1970),
                model: modelId,
                choices: [
                    Choice(
                        index: 0,
                        message: ChatMessage(role: "assistant", content: response),
                        finish_reason: "stop"
                    )
                ],
                usage: Usage(
                    promptTokens: promptTokens.count,
                    completionTokens: tokenizer.encode(text: response).count,
                    totalTokens: promptTokens.count + tokenizer.encode(text: response).count
                )
            )
        } else {
            // Pure text generation
            return try handleChatCompletion(
                messages: messages.map { ChatMessage(role: $0.role, content: $0.content.map { part in
                    if case .text(let t) = part { return t }
                    return ""
                }.joined()) },
                config: config
            )
        }
    }
    
    private func processImageData(_ data: Data) throws -> MTLBuffer {
        print("[VISION] Processing image data: \(data.count) bytes")
        
        // Create CIImage from data
        guard let ciImage = CIImage(data: data) else {
            print("[VISION] Failed to create CIImage")
            throw MarkBaseError.imageProcessingFailed
        }
        print("[VISION] Image size: \(ciImage.extent.width) x \(ciImage.extent.height)")
        
        // Resize to 224x224
        let resizeFilter = CIFilter(name: "CILanczosScaleTransform")!
        resizeFilter.setValue(ciImage, forKey: kCIInputImageKey)
        let scale = 224.0 / max(ciImage.extent.width, ciImage.extent.height)
        resizeFilter.setValue(scale, forKey: kCIInputScaleKey)
        resizeFilter.setValue(1.0, forKey: kCIInputAspectRatioKey)
        
        guard let resized = resizeFilter.outputImage else {
            print("[VISION] Failed to resize image")
            throw MarkBaseError.imageProcessingFailed
        }
        print("[VISION] Resized to: \(resized.extent.width) x \(resized.extent.height)")
        
        // Convert to pixel data
        let context = CIContext()
        guard let cgImage = context.createCGImage(resized, from: resized.extent) else {
            print("[VISION] Failed to create CGImage")
            throw MarkBaseError.imageProcessingFailed
        }
        
        // Extract RGB pixels
        let dataProvider = cgImage.dataProvider!
        let pixelData = dataProvider.data!
        let ptr = CFDataGetBytePtr(pixelData)!
        
        // Create patch embeddings (16x16 patches)
        let patchSize = 16
        let numPatches = 14 * 14  // 224/16 = 14
        let hiddenSize = 768
        
        var patchEmbeddings = [Float](repeating: 0, count: numPatches * hiddenSize)
        
        for patchIdx in 0..<numPatches {
            let patchRow = patchIdx / 14
            let patchCol = patchIdx % 14
            
            for y in 0..<patchSize {
                for x in 0..<patchSize {
                    let globalY = patchRow * patchSize + y
                    let globalX = patchCol * patchSize + x
                    
                    // Clamp to bounds
                    if globalY >= 224 || globalX >= 224 { continue }
                    
                    let pixelIdx = globalY * 224 + globalX
                    let offset = pixelIdx * 4  // RGBA
                    
                    // Normalize to [0, 1]
                    let r = Float(ptr[offset]) / 255.0
                    let g = Float(ptr[offset + 1]) / 255.0
                    let b = Float(ptr[offset + 2]) / 255.0
                    
                    let embedIdx = patchIdx * hiddenSize + (y * patchSize + x) * 3
                    if embedIdx + 2 < patchEmbeddings.count {
                        patchEmbeddings[embedIdx] = r
                        patchEmbeddings[embedIdx + 1] = g
                        patchEmbeddings[embedIdx + 2] = b
                    }
                }
            }
        }
        
        // Debug: Show first patch stats
        var firstPatchSum: Float = 0
        for i in 0..<768 {
            firstPatchSum += patchEmbeddings[i]
        }
        print("[VISION] First patch RGB mean: R=\(patchEmbeddings[0]), G=\(patchEmbeddings[1]), B=\(patchEmbeddings[2]), sum=\(firstPatchSum/768)")
        
        // Create buffer
        let buffer = engine.device.makeBuffer(bytes: patchEmbeddings, length: patchEmbeddings.count * 4)!
        print("[VISION] Patch embeddings buffer created: \(buffer.length) bytes")
        return buffer
    }
    
    private func generateWithVision(
        textTokens: [Int],
        visionFeatures: MTLBuffer,
        maxTokens: Int
    ) throws -> String {
        print("[VISION GEN] Starting vision-guided generation")
        print("[VISION GEN] Text tokens: \(textTokens.count), max tokens: \(maxTokens)")
        
        guard let mm = multimodalModel,
              let tower = mm.visionTowerFull else {
            print("[VISION GEN] Multimodal model not available")
            throw MarkBaseError.multimodalNotSupported
        }
        print("[VISION GEN] Vision tower available")
        
        let numPatches = 196
        let hiddenSize = 2560
        
        // Vision tower forward
        print("[VISION GEN] Running vision tower forward pass...")
        let visionOutputBuffer = engine.device.makeBuffer(length: numPatches * hiddenSize * 4)!
        try tower.forward(patchEmbeddings: visionFeatures, numPatches: numPatches, outputBuffer: visionOutputBuffer)
        print("[VISION GEN] Vision tower forward completed")
        
        // Pool embeddings
        let visionPtr = visionOutputBuffer.contents().assumingMemoryBound(to: Float.self)
        
        // Debug: Check vision output stats
        var outputMag: Float = 0
        for p in 0..<min(5, numPatches) {
            var patchMag: Float = 0
            for i in 0..<hiddenSize {
                let val = visionPtr[p * hiddenSize + i]
                patchMag += val * val
            }
            outputMag += sqrt(patchMag)
        }
        print("[VISION GEN] Vision output magnitude (avg 5 patches): \(outputMag/5)")
        
        var pooled = [Float](repeating: 0, count: hiddenSize)
        
        for i in 0..<hiddenSize {
            var sum: Float = 0
            for p in 0..<numPatches {
                sum += visionPtr[p * hiddenSize + i]
            }
            pooled[i] = sum / Float(numPatches)
        }
        
        // Normalize to match text embeddings (magnitude ~5)
        let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        print("[VISION GEN] Pooled embedding magnitude before norm: \(mag)")
        
        let scale: Float = 5.0 / mag
        for i in 0..<hiddenSize {
            pooled[i] *= scale
        }
        
        let magAfterNorm = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        print("[VISION GEN] Pooled embedding magnitude after norm: \(magAfterNorm)")
        
        // Create inference
        let inference = try MultimodalInference(model: mm)
        let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
        
        // Generate
        print("[VISION GEN] Starting multimodal inference...")
        let generatedTokens = try inference.generate(
            textTokens: textTokens,
            precomputedVisionEmbedding: pooledBuffer,
            maxTokens: maxTokens
        )
        print("[VISION GEN] Generated \(generatedTokens.count) tokens")
        
        // Decode response (skip prompt + BOI + IMAGE + EOI)
        let responseStart = textTokens.count + 4
        if generatedTokens.count > responseStart {
            let responseTokens = Array(generatedTokens[responseStart...])
            return tokenizer.decode(tokens: responseTokens)
        }
        
        return ""
    }
    
    // ─────────────────────────────────────────────────────────────
    // Audio Handlers
    // ─────────────────────────────────────────────────────────────
    
    private func processAudioData(_ data: Data) throws -> MTLBuffer {
        print("[AUDIO] Processing audio data: \(data.count) bytes")
        
        // Create audio extractor
        let extractor = AudioFeatureExtractor(
            sampleRate: 16000,
            nMels: 128,
            nFft: 400,
            hopLength: 160,
            fMin: 0,
            fMax: 8000
        )
        
        // Save data to temp file
        let tempFile = "/tmp/audio_input.wav"
        try data.write(to: URL(fileURLWithPath: tempFile))
        
        // Load audio file
        let audioSamples = try extractor.loadAudioFile(url: URL(fileURLWithPath: tempFile))
        print("[AUDIO] Audio samples: \(audioSamples.count)")
        
        // Extract mel spectrogram
        let melSpec = extractor.extractMelSpectrogram(from: audioSamples)
        print("[AUDIO] Mel spectrogram: \(melSpec.count) frames x \(melSpec[0].count) mels")
        
        // Flatten to [frames, 128]
        let numFrames = melSpec.count
        let melDim = 128
        var audioFeatures = [Float](repeating: 0, count: numFrames * melDim)
        
        for frameIdx in 0..<numFrames {
            for melIdx in 0..<melDim {
                audioFeatures[frameIdx * melDim + melIdx] = melSpec[frameIdx][melIdx]
            }
        }
        
        // Normalize
        let mean = audioFeatures.reduce(0, +) / Float(audioFeatures.count)
        let std = sqrt(audioFeatures.map { ($0 - mean) * ($0 - mean) }.reduce(0, +) / Float(audioFeatures.count))
        for i in 0..<audioFeatures.count {
            audioFeatures[i] = (audioFeatures[i] - mean) / max(std, 1e-6)
        }
        
        print("[AUDIO] Audio features normalized: mean=\(mean), std=\(std)")
        
        // Create buffer
        let buffer = engine.device.makeBuffer(bytes: audioFeatures, length: audioFeatures.count * 4)!
        print("[AUDIO] Audio buffer created: \(buffer.length) bytes")
        
        return buffer
    }
    
    private func generateWithAudio(
        textTokens: [Int],
        audioFeatures: MTLBuffer,
        numFrames: Int,
        maxTokens: Int
    ) throws -> String {
        print("[AUDIO GEN] Starting audio-guided generation")
        print("[AUDIO GEN] Text tokens: \(textTokens.count), frames: \(numFrames), max tokens: \(maxTokens)")
        
        guard let mm = multimodalModel else {
            print("[AUDIO GEN] Multimodal model not available")
            throw MarkBaseError.audioProcessingFailed
        }
        
        // Check if audio tower is available (either AudioTower or AudioTower12B)
        let hasAudioTower = mm.audioTowerFull != nil || mm.audioTower != nil
        if !hasAudioTower {
            print("[AUDIO GEN] Audio tower not available")
            throw MarkBaseError.audioProcessingFailed
        }
        
        print("[AUDIO GEN] Audio tower available")
        
        // Audio tower forward pass (simplified)
        // Note: Full audio tower implementation would use audioOutputBuffer
        // For now, pool directly from input features
        
        // Pool audio features
        let audioPtr = audioFeatures.contents().assumingMemoryBound(to: Float.self)
        let audioLength = audioFeatures.length / 4  // Float size
        var pooled = [Float](repeating: 0, count: 2560)
        
        // Simple pooling: average across frames
        let melDim = 128
        for i in 0..<min(2560, melDim * 20) {  // Take first 20 mel bands
            var sum: Float = 0
            for frame in 0..<numFrames {
                let idx = frame * melDim + i
                if idx < audioLength {
                    sum += audioPtr[idx]
                }
            }
            pooled[i] = sum / Float(numFrames)
        }
        
        // Pad remaining
        for i in (melDim * 20)..<2560 {
            pooled[i] = 0
        }
        
        // Normalize to magnitude ~5
        let mag = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        print("[AUDIO GEN] Pooled audio magnitude: \(mag)")
        
        let scale: Float = 5.0 / max(mag, 1e-6)
        for i in 0..<2560 {
            pooled[i] *= scale
        }
        
        let magAfter = sqrt(pooled.reduce(0) { $0 + $1 * $1 })
        print("[AUDIO GEN] Normalized magnitude: \(magAfter)")
        
        // Multimodal inference
        let inference = try MultimodalInference(model: mm)
        let pooledBuffer = engine.device.makeBuffer(bytes: pooled, length: pooled.count * 4)!
        
        // Note: MultimodalInference needs audio support
        // For now, reuse vision path with audio embedding
        let generatedTokens = try inference.generate(
            textTokens: textTokens,
            precomputedVisionEmbedding: pooledBuffer,  // Use audio embedding
            maxTokens: maxTokens
        )
        
        print("[AUDIO GEN] Generated \(generatedTokens.count) tokens")
        
        // Decode response
        let responseStart = textTokens.count + 4
        if generatedTokens.count > responseStart {
            let responseTokens = Array(generatedTokens[responseStart...])
            return tokenizer.decode(tokens: responseTokens)
        }
        
        return ""
    }
    
    public func handleChatCompletion(
        messages: [ChatMessage],
        config: GenerationConfig
    ) throws -> ChatCompletionResponse {
        let prompt = buildChatPrompt(messages: messages)
        let response = try generator.generateComplete(prompt: prompt, config: config)
        
        return ChatCompletionResponse(
            id: generateId("chatcmpl"),
            object: "chat.completion",
            created: Int(Date().timeIntervalSince1970),
            model: modelId,
            choices: [
                Choice(
                    index: 0,
                    message: ChatMessage(role: "assistant", content: response),
                    finish_reason: "stop"
                )
            ],
            usage: Usage(
                promptTokens: tokenizer.encode(text: prompt).count,
                completionTokens: tokenizer.encode(text: response).count,
                totalTokens: tokenizer.encode(text: prompt + response).count
            )
        )
    }
    
    public func handleStreamChatCompletion(
        messages: [ChatMessage],
        config: GenerationConfig
    ) throws -> [SSEEvent] {
        let prompt = buildChatPrompt(messages: messages)
        let id = generateId("chatcmpl")
        
        // Get tokens for streaming
        let promptTokens = tokenizer.encode(text: prompt)
        var generatedTokens: [Int] = []
        var position = promptTokens.count
        var lastLogits: [Float] = []
        
        // Pre-fill
        for (i, tokenId) in promptTokens.enumerated() {
            lastLogits = try model.forward(tokenId: tokenId, position: i)
        }
        
        var events: [SSEEvent] = []
        var streamDecoder = StreamingDecoder(tokenizer: tokenizer)
        
        // Stream tokens
        for _ in 0..<config.maxTokens {
            let nextToken = sampler.sample(
                logits: lastLogits,
                temperature: config.temperature,
                topK: config.topK,
                topP: config.topP
            )
            
            if tokenizer.eosTokenIds.contains(nextToken) {
                break
            }
            
            generatedTokens.append(nextToken)
            let tokenText = streamDecoder.consume(tokenId: nextToken)
            
            if !tokenText.isEmpty {
                events.append(SSEStream.chatChunk(
                    id: id,
                    model: modelId,
                    content: tokenText
                ))
            }
            
            lastLogits = try model.forward(tokenId: nextToken, position: position)
            position += 1
        }
        
        // Final chunk
        events.append(SSEStream.chatChunk(
            id: id,
            model: modelId,
            finishReason: "stop"
        ))
        
        // Done
        events.append(SSEStream.done())
        
        return events
    }
    
    // ─────────────────────────────────────────────────────────────
    // Text Completion Handlers
    // ─────────────────────────────────────────────────────────────
    
    public func handleCompletion(
        prompt: String,
        config: GenerationConfig
    ) throws -> CompletionResponse {
        let response = try generator.generateComplete(prompt: prompt, config: config)
        
        return CompletionResponse(
            id: generateId("cmpl"),
            object: "text_completion",
            created: Int(Date().timeIntervalSince1970),
            model: modelId,
            choices: [
                CompletionChoice(
                    index: 0,
                    text: response,
                    finishReason: "stop"
                )
            ],
            usage: Usage(
                promptTokens: tokenizer.encode(text: prompt).count,
                completionTokens: tokenizer.encode(text: response).count,
                totalTokens: tokenizer.encode(text: prompt + response).count
            )
        )
    }
    
    public func handleStreamCompletion(
        prompt: String,
        config: GenerationConfig
    ) throws -> [SSEEvent] {
        let id = generateId("cmpl")
        
        // Get tokens
        let promptTokens = tokenizer.encode(text: prompt)
        var generatedTokens: [Int] = []
        var position = promptTokens.count
        var lastLogits: [Float] = []
        
        // Pre-fill
        for (i, tokenId) in promptTokens.enumerated() {
            lastLogits = try model.forward(tokenId: tokenId, position: i)
        }
        
        var events: [SSEEvent] = []
        var streamDecoder = StreamingDecoder(tokenizer: tokenizer)
        
        // Stream tokens
        for _ in 0..<config.maxTokens {
            let nextToken = sampler.sample(
                logits: lastLogits,
                temperature: config.temperature,
                topK: config.topK,
                topP: config.topP
            )
            
            if tokenizer.eosTokenIds.contains(nextToken) {
                break
            }
            
            generatedTokens.append(nextToken)
            let tokenText = streamDecoder.consume(tokenId: nextToken)
            
            if !tokenText.isEmpty {
                events.append(SSEStream.textChunk(
                    id: id,
                    model: modelId,
                    text: tokenText
                ))
            }
            
            lastLogits = try model.forward(tokenId: nextToken, position: position)
            position += 1
        }
        
        // Final chunk with finish reason
        events.append(SSEStream.textChunk(
            id: id,
            model: modelId,
            text: "",
            finishReason: "stop"
        ))
        
        // Done
        events.append(SSEStream.done())
        
        return events
    }
    
    // ─────────────────────────────────────────────────────────────
    // Multimodal Handlers
    // ─────────────────────────────────────────────────────────────
    
    public func handleMultimodalChat(
        messages: [MultimodalMessage],
        config: GenerationConfig
    ) throws -> ChatCompletionResponse {
        // Validate messages
        try Validator.validateMultimodalMessages(messages)
        
        // Build prompt with multimodal content
        let prompt = try buildMultimodalPrompt(messages: messages)
        
        // Process images/audio if present
        for message in messages {
            for imageUrl in message.imageUrls {
                let _ = try MediaProcessor.loadImage(from: imageUrl.url)
                // TODO: Pass to vision tower
            }
            for audioUrl in message.audioUrls {
                let _ = try MediaProcessor.loadAudio(from: audioUrl.url)
                // TODO: Pass to audio tower
            }
        }
        
        let response = try generator.generateComplete(prompt: prompt, config: config)
        
        return ChatCompletionResponse(
            id: generateId("chatcmpl"),
            object: "chat.completion",
            created: Int(Date().timeIntervalSince1970),
            model: modelId,
            choices: [
                Choice(
                    index: 0,
                    message: ChatMessage(role: "assistant", content: response),
                    finish_reason: "stop"
                )
            ],
            usage: Usage(
                promptTokens: tokenizer.encode(text: prompt).count,
                completionTokens: tokenizer.encode(text: response).count,
                totalTokens: tokenizer.encode(text: prompt + response).count
            )
        )
    }
    
    // ─────────────────────────────────────────────────────────────
    // Utilities
    // ─────────────────────────────────────────────────────────────
    
    private func buildMultimodalPrompt(messages: [MultimodalMessage]) throws -> String {
        var prompt = ""
        
        for message in messages {
            switch message.role {
            case "system":
                prompt += "<start_of_turn>user\nSystem: \(message.textContent)<end_of_turn>\n"
            case "user":
                prompt += "<start_of_turn>user\n"
                
                // Add image tags
                for _ in message.imageUrls {
                    prompt += "<image>\n"
                }
                
                // Add audio tags
                for _ in message.audioUrls {
                    prompt += "<audio>\n"
                }
                
                // Add video tags
                for _ in message.videoUrls {
                    prompt += "<video>\n"
                }
                
                // Add text
                if !message.textContent.isEmpty {
                    prompt += "\(message.textContent)\n"
                }
                
                prompt += "<end_of_turn>\n"
            case "assistant":
                prompt += "<start_of_turn>model\n\(message.textContent)<end_of_turn>\n"
            default:
                prompt += "\(message.textContent)\n"
            }
        }
        
        prompt += "<start_of_turn>model\n"
        return prompt
    }
    
    private func buildChatPrompt(messages: [ChatMessage]) -> String {
        var prompt = ""
        for message in messages {
            let role = message.role == "assistant" ? "model" : message.role
            prompt += "<|turn>\(role)\n\(message.content ?? "")<turn|>\n"
        }
        prompt += "<|turn>model\n"
        return prompt
    }
    
    private func generateId(_ prefix: String) -> String {
        let uuid = UUID().uuidString.replacingOccurrences(of: "-", with: "")
        return "\(prefix)-\(uuid.prefix(29))"
    }
}

// ─────────────────────────────────────────────────────────────
// Request Models
// ─────────────────────────────────────────────────────────────

public struct ChatMessage: Codable, Sendable {
    public let role: String
    public let content: String?
    public let tool_calls: [ToolCall]?
    public let name: String?
    
    public init(
        role: String,
        content: String? = nil,
        tool_calls: [ToolCall]? = nil,
        name: String? = nil
    ) {
        self.role = role
        self.content = content
        self.tool_calls = tool_calls
        self.name = name
    }
}

// ─────────────────────────────────────────────────────────────
// Response Models
// ─────────────────────────────────────────────────────────────

public struct ChatCompletionResponse: Codable, Sendable {
    public let id: String
    public let object: String
    public let created: Int
    public let model: String
    public let choices: [Choice]
    public let usage: Usage
}

public struct CompletionResponse: Codable, Sendable {
    public let id: String
    public let object: String
    public let created: Int
    public let model: String
    public let choices: [CompletionChoice]
    public let usage: Usage
}

public struct Choice: Codable, Sendable {
    public let index: Int
    public let message: ChatMessage
    public let finish_reason: String
    
    public init(
        index: Int,
        message: ChatMessage,
        finish_reason: String
    ) {
        self.index = index
        self.message = message
        self.finish_reason = finish_reason
    }
}

public struct CompletionChoice: Codable, Sendable {
    public let index: Int
    public let text: String
    public let finishReason: String
}

public struct Usage: Codable, Sendable {
    public let promptTokens: Int
    public let completionTokens: Int
    public let totalTokens: Int
}

// ─────────────────────────────────────────────────────────────
// Embeddings Handler
// ─────────────────────────────────────────────────────────────

extension MarkBaseServer {
    /// Get model details
    public func getModelDetails(modelId: String) -> ModelDetails {
        ModelDetails(
            id: modelId,
            capabilities: ModelCapabilities(
                text: true,
                vision: true,
                audio: true,
                embeddings: true,
                streaming: true
            ),
            parameters: ModelParameters(
                context_length: maxContextLength,
                num_hidden_layers: model.numHiddenLayers,
                hidden_size: model.hiddenSize,
                vocab_size: model.vocabSize,
                num_attention_heads: 16,
                num_kv_heads: 8
            )
        )
    }
    
    /// Handle embeddings request
    public func handleEmbeddings(
        request: EmbeddingsRequest
    ) throws -> EmbeddingsResponse {
        var embeddings: [EmbeddingData] = []
        var totalTokens = 0
        
        // Process input
        switch request.input {
        case .string(let text):
            let embedding = try generateEmbedding(text: text)
            let tokens = tokenizer.encode(text: text)
            totalTokens += tokens.count
            embeddings.append(EmbeddingData(index: 0, embedding: embedding))
            
        case .strings(let texts):
            for (index, text) in texts.enumerated() {
                let embedding = try generateEmbedding(text: text)
                let tokens = tokenizer.encode(text: text)
                totalTokens += tokens.count
                embeddings.append(EmbeddingData(index: index, embedding: embedding))
            }
            
        case .tokens(let tokens):
            let text = tokenizer.decode(tokens: tokens)
            let embedding = try generateEmbedding(text: text)
            totalTokens += tokens.count
            embeddings.append(EmbeddingData(index: 0, embedding: embedding))
            
        case .tokensList(let tokensList):
            for (index, tokens) in tokensList.enumerated() {
                let text = tokenizer.decode(tokens: tokens)
                let embedding = try generateEmbedding(text: text)
                totalTokens += tokens.count
                embeddings.append(EmbeddingData(index: index, embedding: embedding))
            }
        }
        
        return EmbeddingsResponse(
            data: embeddings,
            model: modelId,
            usage: EmbeddingUsage(
                prompt_tokens: totalTokens,
                total_tokens: totalTokens
            )
        )
    }
    
    /// Generate embedding for text
    private func generateEmbedding(text: String) throws -> [Float] {
        // Tokenize
        let tokens = tokenizer.encode(text: text)
        
        // Run forward pass for all tokens
        var lastHidden: [Float] = []
        for (position, tokenId) in tokens.enumerated() {
            lastHidden = try model.forward(tokenId: tokenId, position: position)
        }
        
        // Use final hidden state as embedding
        // For a proper embedding, we should use the mean of all token embeddings
        // but for simplicity, we use the last hidden state
        return lastHidden
    }
}

// ─────────────────────────────────────────────────────────────
// Video Analysis Handler
// ─────────────────────────────────────────────────────────────

extension MarkBaseServer {
    /// Analyze a video file with frame pooling: frames → per-frame VisionTower → mean pool → single IMAGE token.
    public func handleVideoAnalysis(videoURL: URL, config: GenerationConfig) async throws -> ChatCompletionResponse {
        // 1. Extract frames and audio
        print("Video: Processing \(videoURL.lastPathComponent)...")
        let videoConfig = VideoProcessor.Config(maxFrames: 32, sceneThreshold: 0.15)
        let videoData = try await VideoProcessor.process(url: videoURL, config: videoConfig)
        print("  Frames: \(videoData.frames.count), Audio: \(videoData.audioSamples.count) samples")

        guard let mm = multimodalModel else {
            throw MarkBaseError.invalidParameter(parameter: "video", message: "Multimodal model not loaded")
        }
        guard mm.visionTowerFull != nil || mm.visionTower != nil else {
            throw MarkBaseError.invalidParameter(parameter: "video", message: "Vision tower not loaded")
        }

        let hiddenSize = model.hiddenSize
        let patchSize = 16
        let targetSize = 224
        let device = engine.device

        // 2. Per-frame: patchify → VisionTower → mean pool
        var pooledFrames: [Float] = []

        for (i, frame) in videoData.frames.enumerated() {
            guard let resized = VideoProcessor.resizePixelBuffer(frame.pixelBuffer,
                                                                  targetWidth: targetSize,
                                                                  targetHeight: targetSize) else { continue }

            let (embeddings, numPatches, _) = VideoProcessor.frameToPatchEmbeddings(resized, patchSize: patchSize)
            guard numPatches > 0 else { continue }

            // Create GPU buffer and run vision tower
            let inputBuf = device.makeBuffer(bytes: embeddings,
                                             length: embeddings.count * MemoryLayout<Float>.stride)!
            let outputBuf = device.makeBuffer(length: numPatches * hiddenSize * MemoryLayout<Float>.stride)!
            if let vt = mm.visionTowerFull {
                try vt.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: outputBuf)
            } else if let vt = mm.visionTower {
                try vt.forward(patchEmbeddings: inputBuf, numPatches: numPatches, outputBuffer: outputBuf)
            }

            // Read back and mean pool across patches
            let outputPtr = outputBuf.contents().assumingMemoryBound(to: Float.self)
            let outputFloats = Array(UnsafeBufferPointer(start: outputPtr, count: numPatches * hiddenSize))
            var frameEmbedding = [Float](repeating: 0, count: hiddenSize)
            for p in 0..<numPatches {
                for h in 0..<hiddenSize {
                    frameEmbedding[h] += outputFloats[p * hiddenSize + h]
                }
            }
            for h in 0..<hiddenSize {
                frameEmbedding[h] /= Float(numPatches)
            }
            pooledFrames.append(contentsOf: frameEmbedding)
            print("  Frame \(i): \(numPatches) patches → pooled [\(hiddenSize)]")
        }

        guard !pooledFrames.isEmpty else {
            throw MarkBaseError.invalidParameter(parameter: "video", message: "No frames processed")
        }

        // 3. Cross-frame mean pool → single embedding
        let numFrames = videoData.frames.count
        var videoEmbedding = [Float](repeating: 0, count: hiddenSize)
        for f in 0..<numFrames {
            for h in 0..<hiddenSize {
                videoEmbedding[h] += pooledFrames[f * hiddenSize + h]
            }
        }
        for h in 0..<hiddenSize {
            videoEmbedding[h] /= Float(numFrames)
        }

        let precomputedBuf = device.makeBuffer(bytes: videoEmbedding,
                                               length: hiddenSize * MemoryLayout<Float>.stride)!
        print("  ✓ Pooled \(numFrames) frames → 1×\(hiddenSize) embedding")

        // 4. Extract audio mel spectrogram
        var audioFeatures: [[Float]] = []
        if !videoData.audioSamples.isEmpty {
            let extractor = AudioFeatureExtractor()
            audioFeatures = extractor.extractMelSpectrogram(from: videoData.audioSamples)
            print("  Audio frames (mel): \(audioFeatures.count)")
        }

        // 5. Tokenize prompt
        let prompt = "Describe this video in detail, including visual content and audio."
        let promptTokens = tokenizer.encode(text: prompt)

        // 6. Run multimodal inference with precomputed embedding
        let inference = try MultimodalInference(model: mm)
        let generatedTokens = try inference.generate(
            textTokens: promptTokens,
            audioFeatures: audioFeatures.isEmpty ? nil : audioFeatures,
            precomputedVisionEmbedding: precomputedBuf,
            maxTokens: config.maxTokens
        )

        // 7. Decode
        let generatedText = tokenizer.decode(tokens: Array(generatedTokens[promptTokens.count...]))

        return ChatCompletionResponse(
            id: generateId("chatcmpl"),
            object: "chat.completion",
            created: Int(Date().timeIntervalSince1970),
            model: modelId,
            choices: [
                Choice(
                    index: 0,
                    message: ChatMessage(role: "assistant", content: generatedText),
                    finish_reason: "stop"
                )
            ],
            usage: Usage(
                promptTokens: promptTokens.count,
                completionTokens: generatedTokens.count - promptTokens.count,
                totalTokens: generatedTokens.count
            )
        )
    }
}