# SPDX-License-Identifier: Apache-2.0
# SPDX-FileCopyrightText: Copyright contributors to the vLLM project
"""Utilities for model repo interaction."""

import fnmatch
import json
import os
import time
from collections.abc import Callable
from functools import cache
from pathlib import Path
from typing import Any, TypeVar

import huggingface_hub
from huggingface_hub import HfApi, try_to_load_from_cache
from huggingface_hub.utils import (
    EntryNotFoundError,
    HfHubHTTPError,
    LocalEntryNotFoundError,
    RepositoryNotFoundError,
    RevisionNotFoundError,
)

from vllm import envs
from vllm.logger import init_logger
from vllm.version import __version__ as VLLM_VERSION

logger = init_logger(__name__)

_hf_api: HfApi | None = None


def hf_api() -> HfApi:
    """Return a shared HfApi instance tagged with vLLM's library info."""
    global _hf_api
    if _hf_api is None:
        _hf_api = HfApi(
            library_name="vllm",
            library_version=VLLM_VERSION,
        )
    return _hf_api


def hf_fs() -> "huggingface_hub.HfFileSystem":
    """Return a fresh HfFileSystem tagged with vLLM's library info."""
    return huggingface_hub.HfFileSystem(
        library_name="vllm",
        library_version=VLLM_VERSION,
    )


_R = TypeVar("_R")


def with_retry(
    func: Callable[[], _R],
    log_msg: str,
    max_retries: int = 2,
    retry_delay: int = 2,
) -> _R:
    for attempt in range(max_retries):
        try:
            return func()
        except Exception as e:
            if attempt == max_retries - 1:
                logger.error("%s: %s", log_msg, e)
                raise
            logger.error(
                "%s: %s, retrying %d of %d", log_msg, e, attempt + 1, max_retries
            )
            time.sleep(retry_delay)
            retry_delay *= 2

    raise AssertionError("Should not be reached")


# @cache doesn't cache exceptions
@cache
def list_repo_files(
    repo_id: str,
    *,
    revision: str | None = None,
    repo_type: str | None = None,
    token: str | bool | None = None,
) -> list[str]:
    def lookup_files() -> list[str]:
        # directly list files if model is local
        if (local_path := Path(repo_id)).exists():
            return [
                str(file.relative_to(local_path))
                for file in local_path.rglob("*")
                if file.is_file()
            ]
        # if model is remote, use hf_hub api to list files
        try:
            if envs.VLLM_USE_MODELSCOPE:
                from vllm.transformers_utils.utils import modelscope_list_repo_files

                return modelscope_list_repo_files(
                    repo_id,
                    revision=revision,
                    token=os.getenv("MODELSCOPE_API_TOKEN", None),
                )
            return hf_api().list_repo_files(
                repo_id, revision=revision, repo_type=repo_type, token=token
            )
        except huggingface_hub.errors.OfflineModeIsEnabled:
            # Don't raise in offline mode,
            # all we know is that we don't have this
            # file cached.
            return []

    return with_retry(lookup_files, "Error retrieving file list")


def list_filtered_repo_files(
    model_name_or_path: str,
    allow_patterns: list[str],
    revision: str | None = None,
    repo_type: str | None = None,
    token: str | bool | None = None,
) -> list[str]:
    try:
        all_files = list_repo_files(
            repo_id=model_name_or_path,
            revision=revision,
            token=token,
            repo_type=repo_type,
        )
    except Exception:
        logger.error(
            "Error retrieving file list. Please ensure your `model_name_or_path`"
            "`repo_type`, `token` and `revision` arguments are correctly set. "
            "Returning an empty list."
        )
        return []

    file_list = []
    # Filter patterns on filenames
    for pattern in allow_patterns:
        file_list.extend(
            [
                file
                for file in all_files
                if fnmatch.fnmatch(os.path.basename(file), pattern)
            ]
        )
    return file_list


def any_pattern_in_repo_files(
    model_name_or_path: str,
    allow_patterns: list[str],
    revision: str | None = None,
    repo_type: str | None = None,
    token: str | bool | None = None,
):
    return (
        len(
            list_filtered_repo_files(
                model_name_or_path=model_name_or_path,
                allow_patterns=allow_patterns,
                revision=revision,
                repo_type=repo_type,
                token=token,
            )
        )
        > 0
    )


def is_mistral_model_repo(
    model_name_or_path: str,
    revision: str | None = None,
    repo_type: str | None = None,
    token: str | bool | None = None,
) -> bool:
    return any_pattern_in_repo_files(
        model_name_or_path=model_name_or_path,
        allow_patterns=["consolidated*.safetensors"],
        revision=revision,
        repo_type=repo_type,
        token=token,
    )


def file_exists(
    repo_id: str,
    file_name: str,
    *,
    repo_type: str | None = None,
    revision: str | None = None,
    token: str | bool | None = None,
) -> bool:
    # `list_repo_files` is cached and retried on error, so this is more efficient than
    # huggingface_hub.file_exists default implementation when looking for multiple files
    file_list = list_repo_files(
        repo_id, repo_type=repo_type, revision=revision, token=token
    )
    return file_name in file_list


# In offline mode the result can be a false negative
def file_or_path_exists(
    model: str | Path, config_name: str, revision: str | None
) -> bool:
    if (local_path := Path(model)).exists():
        return (local_path / config_name).is_file()

    # Offline mode support: Check if config file is cached already
    cached_filepath = try_to_load_from_cache(
        repo_id=model, filename=config_name, revision=revision
    )
    if isinstance(cached_filepath, str):
        # The config file exists in cache - we can continue trying to load
        return True

    # NB: file_exists will only check for the existence of the config file on
    # hf_hub. This will fail in offline mode.

    if cached_filepath is None:
        # The config file is not cached - check if it exists on hf_hub
        return file_exists(str(model), config_name, revision=revision)
    # The config file is known to not exist in cache - we can return False
    return False


def get_model_path(model: str | Path, revision: str | None = None):
    if os.path.exists(model):
        return model
    assert huggingface_hub.constants.HF_HUB_OFFLINE
    common_kwargs = dict(
        local_files_only=huggingface_hub.constants.HF_HUB_OFFLINE,
        ignore_patterns="*",
        revision=revision,
    )

    if envs.VLLM_USE_MODELSCOPE:
        from modelscope.hub.snapshot_download import snapshot_download

        return snapshot_download(model_id=model, **common_kwargs)

    return hf_api().snapshot_download(
        repo_id=model,
        **common_kwargs,
    )


def _try_download_from_hf_hub(
    model: str | Path, file_name: str, revision: str | None
) -> Path | None:
    """Try to download a file from HuggingFace Hub.

    Returns the local path on success, None on failure.
    Skips download if model is a local directory.
    """
    if Path(model).is_dir():
        return None
    try:
        return Path(
            hf_api().hf_hub_download(
                model,
                file_name,
                revision=revision,
            )
        )
    except huggingface_hub.errors.OfflineModeIsEnabled:
        return None
    except (
        RepositoryNotFoundError,
        RevisionNotFoundError,
        EntryNotFoundError,
        LocalEntryNotFoundError,
    ) as e:
        logger.debug("File or repository not found in hf_hub_download: %s", e)
        return None
    except HfHubHTTPError as e:
        logger.warning(
            "Cannot connect to Hugging Face Hub. Skipping file download for '%s':",
            file_name,
            exc_info=e,
        )
        return None


def get_hf_file_bytes(
    file_name: str, model: str | Path, revision: str | None = "main"
) -> bytes | None:
    """Get file contents from HuggingFace repository as bytes."""
    file_path = try_get_local_file(model=model, file_name=file_name, revision=revision)

    if file_path is None:
        file_path = _try_download_from_hf_hub(model, file_name, revision)

    if isinstance(file_path, Path) and file_path.is_file():
        with open(file_path, "rb") as file:
            return file.read()

    return None


def try_get_local_file(
    model: str | Path, file_name: str, revision: str | None = "main"
) -> Path | Any | None:
    """
    Try to get a local file from the HuggingFace repository.

    The possible return values are:

    - A `Path` object if the local file is found
    - The `huggingface_hub._CACHED_NO_EXIST` sentinel if the file is known to not exist
    - `None` if the file is not found and we cannot determine if it exists or not

    Callers of this method should handle the `_CACHED_NO_EXIST` sentinel appropriately.
    Checking if the return value `is not None` is not sufficient because it does not
    distinguish between the file not existing and the file not being found.
    """
    file_path = Path(model) / file_name
    if file_path.is_file():
        return file_path
    else:
        try:
            cached_filepath = try_to_load_from_cache(
                repo_id=model, filename=file_name, revision=revision
            )
            if isinstance(cached_filepath, str):
                return Path(cached_filepath)
            return cached_filepath
        except ValueError:
            ...
    return None


def get_hf_file_to_dict(
    file_name: str, model: str | Path, revision: str | None = "main"
):
    """
    Downloads a file from the Hugging Face Hub and returns
    its contents as a dictionary.

    Parameters:
    - file_name (str): The name of the file to download.
    - model (str): The name of the model on the Hugging Face Hub.
    - revision (str): The specific version of the model.

    Returns:
    - config_dict (dict): A dictionary containing
    the contents of the downloaded file.
    """

    file_path = try_get_local_file(model=model, file_name=file_name, revision=revision)

    if file_path is None:
        file_path = _try_download_from_hf_hub(model, file_name, revision)

    if isinstance(file_path, Path) and file_path.is_file():
        with open(file_path) as file:
            return json.load(file)

    return None