//
//  InferenceEngine.swift
//  Zyquo MLX
//
//  Author: Simon-Pierre Boucher
//  Mail: contact@spboucher.ai
//

import Foundation
import MLX
import MLXEmbedders
import MLXLMCommon

/// The in-process inference engine. An actor: one loaded model at a time,
/// streaming generation for LLM/VLM, vector embedding, and verifiable
/// memory release on unload (charter Phase 2).
actor InferenceEngine {

    static let shared = InferenceEngine()

    private var adapter: LoadedModelAdapter?
    private(set) var loadedModel: LocalModel?

    // MARK: - Load / unload

    /// Load a local model (replacing any previously loaded one).
    func load(model: LocalModel) async throws {
        try await unload()
        adapter = try await ModelTypeAdapters.load(model: model)
        loadedModel = model
    }

    /// Unload the current model and free GPU buffer cache. Returns the
    /// active-memory delta so callers can verify the release.
    @discardableResult
    func unload() async throws -> Int64 {
        let before = Memory.snapshot().activeMemory
        adapter = nil
        loadedModel = nil
        Memory.clearCache()
        let after = Memory.snapshot().activeMemory
        return Int64(before - after)
    }

    func memorySnapshot() -> (active: Int, cache: Int, peak: Int) {
        let snap = Memory.snapshot()
        return (snap.activeMemory, snap.cacheMemory, snap.peakMemory)
    }

    // MARK: - Text / chat generation (LLM & VLM)

    /// Stream a response for a chat conversation. Images/videos attach to the
    /// user message for VLMs. Cancellation: cancel the surrounding Task.
    func generate(
        messages: [Chat.Message],
        params: GenerationParams
    ) async throws -> AsyncThrowingStream<InferenceEvent, Error> {
        guard let adapter else { throw InferenceEngineError.noModelLoaded }
        guard case .language(let container) = adapter else {
            throw InferenceEngineError.wrongModelKind(expected: "a language model")
        }

        let userInput = UserInput(chat: messages)
        let mlxParams = params.asMLX

        return AsyncThrowingStream { continuation in
            let task = Task {
                do {
                    try await container.perform(nonSendable: userInput) { (context: ModelContext, userInput: UserInput) in
                        let input = try await context.processor.prepare(input: userInput)
                        let stream = try MLXLMCommon.generate(
                            input: input, parameters: mlxParams, context: context)
                        for await generation in stream {
                            if Task.isCancelled { break }
                            switch generation {
                            case .chunk(let text):
                                continuation.yield(.chunk(text))
                            case .info(let info):
                                continuation.yield(
                                    .finished(
                                        InferenceStats(
                                            promptTokens: info.promptTokenCount,
                                            generatedTokens: info.generationTokenCount,
                                            ttft: info.promptTime,
                                            generateTime: info.generateTime,
                                            stopReason: String(describing: info.stopReason)
                                        )))
                            case .toolCall:
                                break // tool use surfaces in Phase 6
                            }
                        }
                        continuation.finish()
                    }
                } catch {
                    continuation.finish(throwing: error)
                }
            }
            continuation.onTermination = { _ in task.cancel() }
        }
    }

    // MARK: - Embeddings

    /// Embed a batch of texts into L2-normalized vectors using the model's
    /// configured pooling strategy (docs/MLX-RESEARCH.md §4.7).
    func embed(texts: [String]) async throws -> [[Float]] {
        guard let adapter else { throw InferenceEngineError.noModelLoaded }
        guard case .embedder(let container) = adapter else {
            throw InferenceEngineError.wrongModelKind(expected: "an embedding model")
        }

        return await container.perform { (context: EmbedderModelContext) in
            var vectors: [[Float]] = []
            for text in texts {
                let tokens = context.tokenizer.encode(text: text, addSpecialTokens: true)
                let ids = MLXArray(tokens)[.newAxis, 0...]
                let output = context.model(
                    ids, positionIds: nil, tokenTypeIds: nil, attentionMask: nil)
                let pooled = context.pooling(output, normalize: true)
                eval(pooled)
                vectors.append(pooled.squeezed().asArray(Float.self))
            }
            return vectors
        }
    }

    /// Cosine similarity between two embedding vectors (both L2-normalized →
    /// plain dot product).
    static func cosineSimilarity(_ a: [Float], _ b: [Float]) -> Float {
        guard a.count == b.count, !a.isEmpty else { return 0 }
        return zip(a, b).reduce(0) { $0 + $1.0 * $1.1 }
    }
}