from __future__ import annotations

import importlib
from typing import (
    Any,
    Final,
    Literal,
    cast,
    overload,
)

from sentence_transformers import SentenceTransformer, SimilarityFunction

from lang_main.constants import (
    SPACY_MODEL_NAME,
    STFR_BACKEND,
    STFR_DEVICE,
    STFR_MODEL_ARGS,
    STFR_MODEL_NAME,
    STFR_SIMILARITY,
)
from lang_main.errors import LanguageModelNotFoundError
from lang_main.types import (
    LanguageModels,
    Model,
    ModelLoaderMap,
    SpacyModel,
    STFRBackends,
    STFRDeviceTypes,
)


@overload
def instantiate_model(
    model_load_map: ModelLoaderMap,
    model: Literal[LanguageModels.SENTENCE_TRANSFORMER],
) -> SentenceTransformer: ...


@overload
def instantiate_model(
    model_load_map: ModelLoaderMap,
    model: Literal[LanguageModels.SPACY],
) -> SpacyModel: ...


def instantiate_model(
    model_load_map: ModelLoaderMap,
    model: LanguageModels,
) -> Model:
    if model not in model_load_map:
        raise KeyError(f'Model >>{model}<< not known. Choose from: {model_load_map.keys()}')
    builder_func = model_load_map[model]['func']
    func_kwargs = model_load_map[model]['kwargs']

    return builder_func(**func_kwargs)


def load_spacy(
    model_name: str,
) -> SpacyModel:
    try:
        spacy_model_obj = importlib.import_module(SPACY_MODEL_NAME)
    except ModuleNotFoundError:
        raise LanguageModelNotFoundError(
            (
                f'Could not find spaCy model >>{model_name}<<. '
                f'Check if it is installed correctly.'
            )
        )
    pretrained_model = cast(SpacyModel, spacy_model_obj.load())

    return pretrained_model


def load_sentence_transformer(
    model_name: str,
    similarity_func: SimilarityFunction = SimilarityFunction.COSINE,
    backend: STFRBackends = STFRBackends.TORCH,
    device: STFRDeviceTypes = STFRDeviceTypes.CPU,
    local_files_only: bool = False,
    model_save_folder: str | None = None,
    model_kwargs: dict[str, Any] | None = None,
) -> SentenceTransformer:
    return SentenceTransformer(
        model_name_or_path=model_name,
        similarity_fn_name=similarity_func,
        backend=backend,  # type: ignore Literal matches Enum
        device=device,
        cache_folder=model_save_folder,
        local_files_only=local_files_only,
        model_kwargs=model_kwargs,
    )


# ** configured model builder functions
MODEL_LOADER_MAP: Final[ModelLoaderMap] = {
    LanguageModels.SENTENCE_TRANSFORMER: {
        'func': load_sentence_transformer,
        'kwargs': {
            'model_name': STFR_MODEL_NAME,
            'similarity_func': STFR_SIMILARITY,
            'backend': STFR_BACKEND,
            'device': STFR_DEVICE,
            'model_kwargs': STFR_MODEL_ARGS,
        },
    },
    LanguageModels.SPACY: {
        'func': load_spacy,
        'kwargs': {
            'model_name': SPACY_MODEL_NAME,
        },
    },
}
