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

import Foundation
import MLX
import MLXHuggingFace
import MLXLLM
import MLXLMCommon
import Tokenizers

/// Events streamed by the engine during generation.
enum GenerationEvent: Sendable {
    case token(String)
    case stats(GenerationStats)
    case finished(GenerationFinishReason)
}

enum GenerationFinishReason: String, Sendable {
    case stop
    case length
    case cancelled
}

/// Final per-response statistics.
struct GenerationStats: Sendable {
    var timeToFirstToken: TimeInterval
    var tokensPerSecond: Double
    var promptTokenCount: Int
    var generationTokenCount: Int
    var peakMemoryBytes: Int
}

/// Engine lifecycle state.
enum EngineState: Sendable, Equatable {
    case unloaded
    case loading(repoID: String)
    case ready(repoID: String)
    case generating(repoID: String)
}

/// Errors surfaced by the engine, with human-readable messages.
enum EngineError: LocalizedError {
    case noModelLoaded
    case modelDirectoryInvalid(URL)
    case loadFailed(String, underlying: String)
    case unsupportedArchitecture(String)
    case outOfMemory(String)

    var errorDescription: String? {
        switch self {
        case .noModelLoaded:
            "No model is loaded. Load a model from the Library first."
        case .modelDirectoryInvalid(let url):
            "The model folder is missing required files (config.json, weights, tokenizer): \(url.path)"
        case .loadFailed(let repo, let underlying):
            "Could not load \(repo): \(underlying)"
        case .unsupportedArchitecture(let arch):
            "This model's architecture “\(arch)” is not supported by the MLX engine yet."
        case .outOfMemory(let repo):
            "Not enough memory to load \(repo). Try a smaller quantization (e.g. 4-bit) or a smaller model."
        }
    }
}

/// The single owner of all inference. One model loaded at a time (v1).
/// States: unloaded → loading → ready ⇄ generating.
actor InferenceEngine {
    private(set) var state: EngineState = .unloaded

    private var container: ModelContainer?
    private var loadedModel: LocalModel?
    private var session: ChatSession?
    private var generationTask: Task<Void, Never>?

    /// Context window of the loaded model (from config.json), defaulting
    /// conservatively when absent.
    private(set) var contextWindow: Int = 4096

    // MARK: - Load / unload

    /// Loads a model from its local directory, replacing any loaded model.
    func load(model: LocalModel) async throws {
        await stopGeneration()
        unloadInternal()
        state = .loading(repoID: model.repoID)
        do {
            guard FileManager.default.fileExists(
                atPath: model.directory.appendingPathComponent("config.json").path)
            else {
                throw EngineError.modelDirectoryInvalid(model.directory)
            }
            contextWindow = Self.readContextWindow(directory: model.directory) ?? 4096
            let container = try await loadModelContainer(
                from: model.directory,
                using: #huggingFaceTokenizerLoader()
            )
            self.container = container
            self.loadedModel = model
            state = .ready(repoID: model.repoID)
        } catch let error as EngineError {
            state = .unloaded
            throw error
        } catch {
            state = .unloaded
            let text = String(describing: error)
            if text.localizedCaseInsensitiveContains("Unsupported model type") {
                throw EngineError.unsupportedArchitecture(model.architecture ?? "unknown")
            }
            if text.localizedCaseInsensitiveContains("memory") {
                throw EngineError.outOfMemory(model.repoID)
            }
            throw EngineError.loadFailed(model.repoID, underlying: error.localizedDescription)
        }
    }

    /// Unloads the current model and verifiably frees memory.
    func unload() async {
        await stopGeneration()
        unloadInternal()
    }

    private func unloadInternal() {
        session = nil
        container = nil
        loadedModel = nil
        MemoryAdvisor.reclaimMemory()
        state = .unloaded
    }

    var currentModel: LocalModel? { loadedModel }

    // MARK: - Sessions

    /// Binds (or rebinds) the engine to a conversation. Rebuilding the session
    /// drops the KV cache, so this is only done when switching conversations,
    /// or when the system prompt changed.
    func startSession(conversation: Conversation) throws {
        guard let container else { throw EngineError.noModelLoaded }
        session = ChatSession(
            container: container,
            conversation: conversation,
            contextWindow: contextWindow,
            params: conversation.params
        )
    }

    /// Ensures the active session matches the conversation, preserving the
    /// KV cache when it does.
    func ensureSession(conversation: Conversation) throws {
        if session?.conversationID != conversation.id {
            try startSession(conversation: conversation)
        }
    }

    var contextUsage: (used: Int, window: Int)? {
        session?.estimatedContextUsage
    }

    // MARK: - Generation

    /// Streams a response to `prompt` inside the bound session.
    /// Cancellation: cancel the consuming task or call `stopGeneration()` —
    /// the underlying generation loop actually stops.
    func generate(prompt: String, params: GenerationParams) throws
        -> AsyncThrowingStream<GenerationEvent, Error>
    {
        guard let session else { throw EngineError.noModelLoaded }
        guard let repoID = loadedModel?.repoID else { throw EngineError.noModelLoaded }

        let (stream, continuation) = AsyncThrowingStream.makeStream(of: GenerationEvent.self)
        state = .generating(repoID: repoID)
        MemoryAdvisor.resetPeakMemory()

        let task = Task {
            let start = ContinuousClock.now
            var firstTokenTime: TimeInterval?
            var response = ""
            var finished: GenerationFinishReason = .cancelled
            do {
                for try await generation in session.stream(prompt: prompt, params: params) {
                    if Task.isCancelled { break }
                    switch generation {
                    case .chunk(let text):
                        if firstTokenTime == nil {
                            firstTokenTime = Self.seconds(since: start)
                        }
                        response += text
                        continuation.yield(.token(text))
                    case .info(let info):
                        let stats = GenerationStats(
                            timeToFirstToken: firstTokenTime ?? info.promptTime,
                            tokensPerSecond: info.tokensPerSecond,
                            promptTokenCount: info.promptTokenCount,
                            generationTokenCount: info.generationTokenCount,
                            peakMemoryBytes: MemoryAdvisor.peakMemoryBytes
                        )
                        finished = switch info.stopReason {
                        case .stop: .stop
                        case .length: .length
                        case .cancelled: .cancelled
                        }
                        continuation.yield(.stats(stats))
                    case .toolCall:
                        break
                    }
                }
                session.noteResponse(response)
                continuation.yield(.finished(Task.isCancelled ? .cancelled : finished))
                continuation.finish()
            } catch {
                continuation.finish(throwing: error)
            }
            self.finishGeneration()
        }
        generationTask = task
        continuation.onTermination = { termination in
            if case .cancelled = termination { task.cancel() }
        }
        return stream
    }

    /// Stops any in-flight generation and waits for it to wind down.
    func stopGeneration() async {
        generationTask?.cancel()
        _ = await generationTask?.value
        generationTask = nil
    }

    private func finishGeneration() {
        if let repoID = loadedModel?.repoID {
            state = .ready(repoID: repoID)
        }
        generationTask = nil
    }

    // MARK: - Helpers

    private static func seconds(since start: ContinuousClock.Instant) -> TimeInterval {
        let duration = start.duration(to: .now)
        return Double(duration.components.seconds)
            + Double(duration.components.attoseconds) * 1e-18
    }

    /// Reads max_position_embeddings from config.json.
    static func readContextWindow(directory: URL) -> Int? {
        let url = directory.appendingPathComponent("config.json")
        guard let data = try? Data(contentsOf: url),
            let json = try? JSONSerialization.jsonObject(with: data) as? [String: Any]
        else { return nil }
        if let value = json["max_position_embeddings"] as? Int { return value }
        if let text = json["text_config"] as? [String: Any],
            let value = text["max_position_embeddings"] as? Int
        {
            return value
        }
        return nil
    }
}