from pathlib import Path
from typing import cast

from clearml import InputModel
from clearml import Task
from transformers import Any
from transformers import AutoModel
from transformers import AutoTokenizer
from transformers import PreTrainedModel
from transformers import PreTrainedTokenizerBase

from kostyl.ml.lightning.extenstions.pretrained_model import (
    LightningCheckpointLoaderMixin,
)


def get_tokenizer_from_clearml(
    model_id: str, task: Task | None = None, ignore_remote_overrides: bool = True
) -> PreTrainedTokenizerBase:
    """
    Retrieve a Hugging Face tokenizer stored in a ClearML.

    Args:
        model_id (str): The ClearML InputModel identifier that holds the tokenizer artifacts.
        task (Task | None, optional): An optional ClearML Task used to associate and sync
            the model. Defaults to None.
        ignore_remote_overrides (bool, optional): Whether to ignore remote hyperparameter
            overrides when connecting the ClearML task. Defaults to True.

    Returns:
        PreTrainedTokenizerBase: The instantiated tokenizer loaded from the local copy
            of the referenced ClearML InputModel.

    """
    clearml_tokenizer = InputModel(model_id=model_id)
    if task is not None:
        clearml_tokenizer.connect(task, ignore_remote_overrides=ignore_remote_overrides)

    tokenizer = AutoTokenizer.from_pretrained(
        clearml_tokenizer.get_local_copy(raise_on_error=True)
    )
    return tokenizer


def get_model_from_clearml[
    TModel: PreTrainedModel | LightningCheckpointLoaderMixin | AutoModel
](
    model_id: str,
    model: type[TModel],
    task: Task | None = None,
    ignore_remote_overrides: bool = True,
    **kwargs: Any,
) -> TModel:
    """
    Retrieve a pretrained model from ClearML and instantiate it using the appropriate loader.

    Args:
        model_id: Identifier of the ClearML input model to retrieve.
        model: The model class that implements either PreTrainedModel or LightningCheckpointLoaderMixin.
        task: Optional ClearML task used to resolve the input model reference. If provided, the input model
            will be connected to this task, with remote overrides optionally ignored.
        ignore_remote_overrides: When connecting the input model to the provided task, determines whether
            remote configuration overrides should be ignored.
        **kwargs: Additional keyword arguments to pass to the model loading method.

    Returns:
        An instantiated model loaded either from a ClearML package directory or a Lightning checkpoint.

    """
    input_model = InputModel(model_id=model_id)

    if task is not None:
        input_model.connect(task, ignore_remote_overrides=ignore_remote_overrides)

    local_path = Path(input_model.get_local_copy(raise_on_error=True))

    if local_path.is_dir() and input_model._is_package():
        model_instance = model.from_pretrained(local_path, **kwargs)
    elif local_path.suffix == ".ckpt":
        if not issubclass(model, LightningCheckpointLoaderMixin):
            raise ValueError(
                f"Model class {model.__name__} is not compatible with Lightning checkpoints."
            )
        model_instance = model.from_lighting_checkpoint(local_path, **kwargs)
    else:
        raise ValueError(
            f"Unsupported model format for path: {local_path}. "
            "Expected a ClearML package directory or a .ckpt file."
        )
    model_instance = cast(TModel, model_instance)
    return model_instance
