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

import Foundation
import MLX
import MLXEmbedders
import MLXHuggingFace
import MLXLLM
import MLXLMCommon
import MLXVLM
import Tokenizers  // required by the #huggingFaceTokenizerLoader macro expansion

/// A loaded model held by the engine, dispatched by `ModelType`.
///
/// LLM and VLM share `ModelContainer` (both come from `ModelFactory`
/// implementations); embeddings use `EmbedderModelContainer`. Speech and
/// image generation run through the Python bridge (Phase 3+) and never hold
/// in-process state here.
enum LoadedModelAdapter {
    case language(ModelContainer)   // .llm and .vlm
    case embedder(EmbedderModelContainer)
}

/// Events streamed to callers during generation.
enum InferenceEvent: Sendable {
    case chunk(String)
    case finished(InferenceStats)
}

enum InferenceEngineError: LocalizedError {
    case unsupportedType(ModelType)
    case noModelLoaded
    case wrongModelKind(expected: String)

    var errorDescription: String? {
        switch self {
        case .unsupportedType(let type):
            "\(type.displayName) models run through the Python pipelines — not the in-process engine."
        case .noModelLoaded:
            "No model is loaded."
        case .wrongModelKind(let expected):
            "The loaded model is not \(expected)."
        }
    }
}

enum ModelTypeAdapters {

    /// Load a local model directory into the right container for its type.
    /// Tokenizers come from swift-transformers via the MLXHuggingFace macro
    /// (docs/MLX-RESEARCH.md §4.1/§4.4).
    static func load(model: LocalModel) async throws -> LoadedModelAdapter {
        switch model.type {
        case .llm:
            let container = try await LLMModelFactory.shared.loadContainer(
                from: model.directory, using: #huggingFaceTokenizerLoader())
            return .language(container)
        case .vlm:
            let container = try await VLMModelFactory.shared.loadContainer(
                from: model.directory, using: #huggingFaceTokenizerLoader())
            return .language(container)
        case .embedding:
            let container = try await EmbedderModelFactory.shared.loadContainer(
                from: model.directory, using: #huggingFaceTokenizerLoader())
            return .embedder(container)
        case .speech, .imageGeneration:
            throw InferenceEngineError.unsupportedType(model.type)
        }
    }
}