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 (
    MODEL_BASE_FOLDER,
    SPACY_MODEL_NAME,
    STFR_BACKEND,
    STFR_CUSTOM_MODELS,
    STFR_DEVICE,
    STFR_MODEL_ARGS,
    STFR_MODEL_NAME,
    STFR_SIMILARITY,
)
from lang_main.errors import LanguageModelNotFoundError
from lang_main.loggers import logger_config as logger
from lang_main.types import (
    LanguageModels,
    Model,
    ModelLoaderMap,
    SpacyModel,
    STFRBackends,
    STFRDeviceTypes,
    STFRModelArgs,
    STFRModelTypes,
)


@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(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())
    logger.info('[MODEL LOADING] Loaded model >>%s<< successfully', model_name)

    return pretrained_model


def _preprocess_STFR_model_name(
    model_name: STFRModelTypes | str,
    backend: STFRBackends,
    force_download: bool = False,
) -> str:
    """utility function to parse specific model names to their
    local file paths per backend
    necessary for models not present on the Hugging Face Hub (like
    own pretrained or optimised models)
    only if chosen model and backend in combination are defined a local
    file path is generated

    Parameters
    ----------
    model_name : STFRModelTypes
        model name given by configuration
    backend: STFRBackends
        backend given by configuration
    force_download: bool
        try to download model even if it is configured as local,
        by default: False

    Returns
    -------
    str
        model name or specific file path if applicable
    """
    combination = (model_name, backend)
    model_name_or_path: str
    if force_download:
        model_name_or_path = model_name
    elif combination in STFR_CUSTOM_MODELS and STFR_CUSTOM_MODELS[combination]:
        # !! defined that each model is placed in a folder with its model name
        # !! without any user names
        folder_name = model_name.split('/')[-1]
        model_path = MODEL_BASE_FOLDER / folder_name
        if not model_path.exists():  # pragma: no cover
            raise FileNotFoundError(
                f'Target model >{model_name}< not found under {model_path}'
            )
        model_name_or_path = str(model_path)
    else:
        model_name_or_path = model_name

    return model_name_or_path


def load_sentence_transformer(
    model_name: STFRModelTypes | str,
    similarity_func: SimilarityFunction = SimilarityFunction.COSINE,
    backend: STFRBackends = STFRBackends.TORCH,
    device: STFRDeviceTypes = STFRDeviceTypes.CPU,
    local_files_only: bool = True,
    trust_remote_code: bool = False,
    model_save_folder: str | None = None,
    model_kwargs: STFRModelArgs | dict[str, Any] | None = None,
    force_download: bool = False,
) -> SentenceTransformer:
    model_name_or_path = _preprocess_STFR_model_name(
        model_name=model_name, backend=backend, force_download=force_download
    )
    model = SentenceTransformer(
        model_name_or_path=model_name_or_path,
        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,
        trust_remote_code=trust_remote_code,
        model_kwargs=model_kwargs,  # type: ignore
    )
    logger.info('[MODEL LOADING] Loaded model >>%s<< successfully', model_name)

    return model


# ** 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,
        },
    },
}
