"""This module provides mixin functionality for the Albumentations library.
It includes utility functions and classes to enhance the core capabilities.
"""

import functools
import logging
from collections.abc import Callable
from pathlib import Path
from typing import Any

from albumentations.core.serialization import load as load_transform
from albumentations.core.serialization import save as save_transform

try:
    from huggingface_hub import HfApi, hf_hub_download
    from huggingface_hub.utils import HfHubHTTPError, SoftTemporaryDirectory

    is_huggingface_hub_available = True
except ImportError:
    is_huggingface_hub_available = False

logger = logging.getLogger(__name__)


def require_huggingface_hub(func: Callable[..., Any]) -> Callable[..., Any]:
    """Decorator that raises ImportError if huggingface_hub is not installed. Use it on Hub
    methods that require the package. Required for push_to_hub.

    This decorator ensures that the `huggingface_hub` package is installed before
    executing the decorated function. If the package is not installed, it raises
    an ImportError with instructions on how to install it.

    Args:
        func (Callable[..., Any]): The function to decorate.

    Returns:
        Callable[..., Any]: The decorated function.

    """

    @functools.wraps(func)
    def wrapper(*args: Any, **kwargs: Any) -> Any:
        if not is_huggingface_hub_available:
            raise ImportError(
                f"You need to install `huggingface_hub` to use {func.__name__}. "
                "Run `pip install huggingface_hub`, or `pip install albumentations[hub]`.",
            )
        return func(*args, **kwargs)

    return wrapper


class HubMixin:
    """Mixin for saving and loading transforms to/from the Hugging Face Hub. Provides
    save_pretrained, from_pretrained, push_to_hub.

    This class provides functionality for saving and loading transforms to/from
    the Hugging Face Hub. It enables serialization, deserialization, and sharing
    of transform configurations.

    Args:
        _CONFIG_KEYS (tuple[str, ...]): Keys used for configuration files.
        _CONFIG_FILE_NAME_TEMPLATE (str): Template for configuration filenames.

    """

    _CONFIG_KEYS = ("train", "eval")
    _CONFIG_FILE_NAME_TEMPLATE = "albumentations_config_{}.json"

    def _save_pretrained(self, save_directory: str | Path, filename: str) -> Path:
        """Save the transform to a specified directory as a JSON config file. Creates directory
        if needed; returns Path to the saved file.

        Args:
            save_directory (str | Path):
                Directory where the transform will be saved.
            filename (str):
                Name of the file to save the transform.

        Returns:
            Path: Path to the saved transform file.

        """
        # create save directory and path
        save_directory = Path(save_directory)
        save_directory.mkdir(parents=True, exist_ok=True)
        save_path = save_directory / filename

        # save transforms
        save_transform(self, save_path, data_format="json")  # type: ignore[arg-type]

        return save_path

    @classmethod
    def _from_pretrained(cls, save_directory: str | Path, filename: str) -> object:
        """Load a transform from a local directory (save_directory/filename). Returns
        deserialized Compose or transform instance. For local loading only.

        Args:
            save_directory (str | Path):
                Directory from where the transform will be loaded.
            filename (str):
                Name of the file to load the transform from.

        Returns:
            object: Loaded transform (Compose or single transform instance).

        """
        save_path = Path(save_directory) / filename
        return load_transform(save_path, data_format="json")

    def save_pretrained(
        self,
        save_directory: str | Path,
        *,
        key: str = "eval",
        allow_custom_keys: bool = False,
        repo_id: str | None = None,
        push_to_hub: bool = False,
        **push_to_hub_kwargs: Any,
    ) -> str | None:
        """Save the transform to a directory and optionally push to the Huggingface Hub. Uses
        key (train/eval) for config filename.

        Args:
            save_directory (str | Path):
                Path to directory in which the transform configuration will be saved.
            key (str):
                Key to identify the configuration type, one of ["train", "eval"]. Defaults to "eval".
            allow_custom_keys (bool):
                Allow custom keys for the configuration. Defaults to False.
            push_to_hub (bool):
                Whether or not to push your transform to the Huggingface Hub after saving it.
            repo_id (str | None):
                ID of your repository on the Hub. Used only if push_to_hub=True. Will default to the folder name if
                not provided.
            push_to_hub_kwargs (Any):
                Additional key word arguments passed along to the push_to_hub method.

        Returns:
            str | None: URL of the commit on the Hub if push_to_hub=True, None otherwise.

        """
        if not allow_custom_keys and key not in self._CONFIG_KEYS:
            raise ValueError(
                f"Invalid key: `{key}`. Please use key from {self._CONFIG_KEYS} keys for upload. "
                "If you want to use a custom key, set `allow_custom_keys=True`.",
            )

        # save model transforms
        filename = self._CONFIG_FILE_NAME_TEMPLATE.format(key)
        self._save_pretrained(save_directory, filename)

        # push to the Hub if required
        if push_to_hub:
            kwargs = push_to_hub_kwargs.copy()  # soft-copy to avoid mutating input
            if repo_id is None:
                repo_id = Path(save_directory).name  # Defaults to `save_directory` name
            return self.push_to_hub(repo_id=repo_id, key=key, **kwargs)
        return None

    @classmethod
    def from_pretrained(
        cls: Any,
        directory_or_repo_id: str | Path,
        *,
        key: str = "eval",
        force_download: bool = False,
        proxies: dict[str, str] | None = None,
        token: str | bool | None = None,
        cache_dir: str | Path | None = None,
        local_files_only: bool = False,
        revision: str | None = None,
    ) -> object:
        """Load a transform from the Huggingface Hub or a local directory. Use key (train/eval)
        to pick which config file to load. Returns Compose or transform.

        Args:
            directory_or_repo_id (str | Path):
                Either the repo_id (string) of a repo with hosted transform on the Hub, e.g. qubvel-hf/albu,
                or a path to a directory containing transform config saved using save_pretrained.
            key (str):
                Key to identify the configuration type, one of ["train", "eval"]. Defaults to "eval".
            revision (str | None):
                Revision of the repo on the Hub. Can be a branch name, a git tag or any commit id.
                Defaults to the latest commit on main branch.
            force_download (bool):
                Whether to force (re-)downloading the transform configuration files from the Hub, overriding
                the existing cache.
            proxies (dict[str, str] | None):
                A dictionary of proxy servers to use by protocol or endpoint.
            token (str | bool | None):
                The token to use as HTTP bearer authorization for remote files.
            cache_dir (str | Path | None):
                Path to the folder where cached files are stored.
            local_files_only (bool):
                If True, avoid downloading the file and return the path to the local cached file if it exists.

        """
        filename = cls._CONFIG_FILE_NAME_TEMPLATE.format(key)
        directory_or_repo_id = Path(directory_or_repo_id)
        transform = None

        # check if the file is already present locally
        if directory_or_repo_id.is_dir():
            if filename in [f.name for f in directory_or_repo_id.iterdir()]:
                transform = cls._from_pretrained(save_directory=directory_or_repo_id, filename=filename)
            elif is_huggingface_hub_available:
                logger.info(
                    f"{filename} not found in {Path(directory_or_repo_id).resolve()}, trying to load from the Hub.",
                )
            else:
                raise FileNotFoundError(
                    f"{filename} not found in {Path(directory_or_repo_id).resolve()}."
                    " Please install `huggingface_hub` to load from the Hub.",
                )
        if transform is not None:
            return transform

        # download the file from the Hub
        try:
            config_file = hf_hub_download(
                repo_id=str(directory_or_repo_id).replace("\\", "/"),
                filename=filename,
                revision=revision,
                cache_dir=cache_dir,
                force_download=force_download,
                proxies=proxies,
                token=token,
                local_files_only=local_files_only,
            )
            directory, filename = Path(config_file).parent, Path(config_file).name
            return cls._from_pretrained(save_directory=directory, filename=filename)

        except HfHubHTTPError as e:
            raise HfHubHTTPError(f"{filename} not found on the HuggingFace Hub") from e

    @require_huggingface_hub
    def push_to_hub(
        self,
        repo_id: str,
        *,
        key: str = "eval",
        allow_custom_keys: bool = False,
        commit_message: str = "Push transform using huggingface_hub.",
        private: bool = False,
        token: str | None = None,
        branch: str | None = None,
        create_pr: bool | None = None,
    ) -> str:
        """Push the saved transform to the Huggingface Hub. Requires repo_id; optional branch,
        token, commit_message. Returns commit URL.

        Use `allow_patterns` and `ignore_patterns` to precisely filter which files should be
        pushed to the hub. Use `delete_patterns` to delete existing remote files in the same
        commit. See [`upload_folder`] reference for more
        details.

        Args:
            repo_id (str):
                ID of the repository to push to (e.g. username/my-model).
            key (str):
                Key to identify the configuration type, one of ["train", "eval"]. Defaults to "eval".
            allow_custom_keys (bool):
                Allow custom keys for the configuration. Defaults to False.
            commit_message (str):
                Message to commit while pushing.
            private (bool):
                Whether the repository created should be private.
            token (str | None):
                The token to use as HTTP bearer authorization for remote files.
            branch (str | None):
                The git branch on which to push the transform. Defaults to "main".
            create_pr (bool | None):
                Whether to create a Pull Request from branch with that commit. Defaults to False.

        Returns:
            str: The URL of the commit of your transform in the given repository.

        """
        if not allow_custom_keys and key not in self._CONFIG_KEYS:
            raise ValueError(
                f"Invalid key: `{key}`. Please use key from {self._CONFIG_KEYS} keys for upload. "
                "If you still want to use a custom key, set `allow_custom_keys=True`.",
            )

        api = HfApi(token=token)
        repo_id = api.create_repo(repo_id=repo_id, private=private, exist_ok=True).repo_id

        # Push the files to the repo in a single commit
        with SoftTemporaryDirectory() as tmp:
            save_directory = Path(tmp) / repo_id
            filename = self._CONFIG_FILE_NAME_TEMPLATE.format(key)
            save_path = self._save_pretrained(save_directory, filename=filename)
            return api.upload_file(
                path_or_fileobj=save_path,
                path_in_repo=filename,
                repo_id=repo_id,
                commit_message=commit_message,
                revision=branch,
                create_pr=create_pr,
            )
