# noqa: isort:skip_file, yapf: disable
# Copyright (c) Alibaba, Inc. and its affiliates.
# Copyright 2020 The HuggingFace Datasets Authors and the TensorFlow Datasets Authors.

import json
import os
import re
import shutil
import warnings
from contextlib import contextmanager
from functools import partial
from pathlib import Path
from urllib.parse import urljoin, urlparse
import requests

from datasets import config
from datasets.utils.file_utils import hash_url_to_filename, get_authentication_headers_for_url, ftp_head, fsspec_head, \
    http_head, _raise_if_offline_mode_is_enabled, ftp_get, fsspec_get, http_get
from filelock import FileLock

from modelscope.utils.config_ds import MS_DATASETS_CACHE
from modelscope.utils.logger import get_logger
from modelscope.hub.api import ModelScopeConfig

logger = get_logger()


def get_from_cache_ms(
    url,
    cache_dir=None,
    force_download=False,
    proxies=None,
    etag_timeout=100,
    resume_download=False,
    user_agent=None,
    local_files_only=False,
    use_etag=True,
    max_retries=0,
    token=None,
    use_auth_token='deprecated',
    ignore_url_params=False,
    storage_options=None,
    download_desc=None,
) -> str:
    """
    Given a URL, look for the corresponding file in the local cache.
    If it's not there, download it. Then return the path to the cached file.

    Return:
        Local path (string)

    Raises:
        FileNotFoundError: in case of non-recoverable file
            (non-existent or no cache on disk)
        ConnectionError: in case of unreachable url
            and no cache on disk
    """
    if use_auth_token != 'deprecated':
        warnings.warn(
            "'use_auth_token' was deprecated in favor of 'token' in version 2.14.0 and will be removed in 3.0.0.\n"
            f"You can remove this warning by passing 'token={use_auth_token}' instead.",
            FutureWarning,
        )
        token = use_auth_token
    if cache_dir is None:
        cache_dir = MS_DATASETS_CACHE
    if isinstance(cache_dir, Path):
        cache_dir = str(cache_dir)

    os.makedirs(cache_dir, exist_ok=True)

    if ignore_url_params:
        # strip all query parameters and #fragments from the URL
        cached_url = urljoin(url, urlparse(url).path)
    else:
        cached_url = url  # additional parameters may be added to the given URL

    connected = False
    response = None
    cookies = None
    etag = None
    head_error = None
    scheme = None

    # Try a first time to file the file on the local file system without eTag (None)
    # if we don't ask for 'force_download' then we spare a request
    filename = hash_url_to_filename(cached_url, etag=None)
    cache_path = os.path.join(cache_dir, filename)

    if os.path.exists(cache_path) and not force_download and not use_etag:
        return cache_path

    # Prepare headers for authentication
    headers = get_authentication_headers_for_url(url, token=token)
    if user_agent is not None:
        headers['user-agent'] = user_agent

    # We don't have the file locally or we need an eTag
    if not local_files_only:
        scheme = urlparse(url).scheme
        if scheme == 'ftp':
            connected = ftp_head(url)
        elif scheme not in ('http', 'https'):
            response = fsspec_head(url, storage_options=storage_options)
            # s3fs uses "ETag", gcsfs uses "etag"
            etag = (response.get('ETag', None) or response.get('etag', None)) if use_etag else None
            connected = True
        try:
            cookies = ModelScopeConfig.get_cookies()
            response = http_head(
                url,
                allow_redirects=True,
                proxies=proxies,
                timeout=etag_timeout,
                max_retries=max_retries,
                headers=headers,
                cookies=cookies,
            )
            if response.status_code == 200:  # ok
                etag = response.headers.get('ETag') if use_etag else None
                for k, v in response.cookies.items():
                    # In some edge cases, we need to get a confirmation token
                    if k.startswith('download_warning') and 'drive.google.com' in url:
                        url += '&confirm=' + v
                        cookies = response.cookies
                connected = True
                # Fix Google Drive URL to avoid Virus scan warning
                if 'drive.google.com' in url and 'confirm=' not in url:
                    url += '&confirm=t'
            # In some edge cases, head request returns 400 but the connection is actually ok
            elif (
                (response.status_code == 400 and 'firebasestorage.googleapis.com' in url)
                or (response.status_code == 405 and 'drive.google.com' in url)
                or (
                    response.status_code == 403
                    and (
                        re.match(r'^https?://github.com/.*?/.*?/releases/download/.*?/.*?$', url)
                        or re.match(r'^https://.*?s3.*?amazonaws.com/.*?$', response.url)
                    )
                )
                or (response.status_code == 403 and 'ndownloader.figstatic.com' in url)
            ):
                connected = True
                logger.info(f"Couldn't get ETag version for url {url}")
            elif response.status_code == 401 and config.HF_ENDPOINT in url and token is None:
                raise ConnectionError(
                    f'Unauthorized for URL {url}. '
                    f'Please use the parameter `token=True` after logging in with `huggingface-cli login`'
                )
        except (OSError, requests.exceptions.Timeout) as e:
            # not connected
            head_error = e
            pass

    # connected == False = we don't have a connection, or url doesn't exist, or is otherwise inaccessible.
    # try to get the last downloaded one
    if not connected:
        if os.path.exists(cache_path) and not force_download:
            return cache_path
        if local_files_only:
            raise FileNotFoundError(
                f'Cannot find the requested files in the cached path at {cache_path} and outgoing traffic has been'
                " disabled. To enable file online look-ups, set 'local_files_only' to False."
            )
        elif response is not None and response.status_code == 404:
            raise FileNotFoundError(f"Couldn't find file at {url}")
        _raise_if_offline_mode_is_enabled(f'Tried to reach {url}')
        if head_error is not None:
            raise ConnectionError(f"Couldn't reach {url} ({repr(head_error)})")
        elif response is not None:
            raise ConnectionError(f"Couldn't reach {url} (error {response.status_code})")
        else:
            raise ConnectionError(f"Couldn't reach {url}")

    # Try a second time
    filename = hash_url_to_filename(cached_url, etag)
    cache_path = os.path.join(cache_dir, filename)

    if os.path.exists(cache_path) and not force_download:
        return cache_path

    # From now on, connected is True.
    # Prevent parallel downloads of the same file with a lock.
    lock_path = cache_path + '.lock'
    with FileLock(lock_path):
        # Retry in case previously locked processes just enter after the precedent process releases the lock
        if os.path.exists(cache_path) and not force_download:
            return cache_path

        incomplete_path = cache_path + '.incomplete'

        @contextmanager
        def temp_file_manager(mode='w+b'):
            with open(incomplete_path, mode) as f:
                yield f

        resume_size = 0
        if resume_download:
            temp_file_manager = partial(temp_file_manager, mode='a+b')
            if os.path.exists(incomplete_path):
                resume_size = os.stat(incomplete_path).st_size

        # Download to temporary file, then copy to cache path once finished.
        # Otherwise, you get corrupt cache entries if the download gets interrupted.
        with temp_file_manager() as temp_file:
            logger.info(f'Downloading to {temp_file.name}')

            # GET file object
            if scheme == 'ftp':
                ftp_get(url, temp_file)
            elif scheme not in ('http', 'https'):
                fsspec_get(url, temp_file, storage_options=storage_options, desc=download_desc)
            else:
                http_get(
                    url,
                    temp_file=temp_file,
                    proxies=proxies,
                    resume_size=resume_size,
                    headers=headers,
                    cookies=cookies,
                    max_retries=max_retries,
                    desc=download_desc,
                )

        logger.info(f'storing {url} in cache at {cache_path}')
        shutil.move(temp_file.name, cache_path)
        umask = os.umask(0o666)
        os.umask(umask)
        os.chmod(cache_path, 0o666 & ~umask)

        logger.info(f'creating metadata file for {cache_path}')
        meta = {'url': url, 'etag': etag}
        meta_path = cache_path + '.json'
        with open(meta_path, 'w', encoding='utf-8') as meta_file:
            json.dump(meta, meta_file)

    return cache_path
