# Copyright 2023 The KServe Authors.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#    http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.

import asyncio
import base64
from concurrent.futures import ThreadPoolExecutor
import fnmatch
from functools import partial
import glob
import gzip
import json
import mimetypes
import multiprocessing
import os
import re
import shutil
import tarfile
import tempfile
import time
from typing import List, Optional
import zipfile
from pathlib import Path
from typing import Tuple
from urllib.parse import urlparse
import requests

from kserve_storage.logging import logger
from kserve_storage.storage_errors import (
    raise_storage_error,
    check_http_response,
)

MODEL_MOUNT_DIRS = "/mnt/models"

_GCS_PREFIX = "gs://"
_S3_PREFIX = "s3://"
_HDFS_PREFIX = "hdfs://"
_WEBHDFS_PREFIX = "webhdfs://"
_AZURE_BLOB_RE = [
    "https://(.+?).blob.core.windows.net/(.+)",
    "https://(.+?).z[0-9]{1,2}.blob.storage.azure.net/(.+)",
]
_AZURE_FILE_RE = [
    "https://(.+?).file.core.windows.net/(.+)",
    "https://(.+?).z[0-9]{1,2}.file.storage.azure.net/(.+)",
]
_LOCAL_PREFIX = "file://"
_URI_RE = "https?://(.+)/(.+)"
_HTTP_PREFIX = "http(s)://"
_HEADERS_SUFFIX = "-headers"
_PVC_PREFIX = "/mnt/pvc"
_HF_PREFIX = "hf://"
_GIT_RE = r"https://.+\.git"

_HDFS_SECRET_DIRECTORY = "/var/secrets/kserve-hdfscreds"
_HDFS_FILE_SECRETS = ["KERBEROS_KEYTAB", "TLS_CERT", "TLS_KEY", "TLS_CA"]

# S3 parallel download configuration
_S3_MAX_FILE_CONCURRENCY = int(os.getenv("S3_MAX_FILE_CONCURRENCY", "4"))
# Global variable for S3 resource in worker processes
_worker_s3_resource = None
# Azure async download configuration
_AZURE_MAX_FILE_CONCURRENCY = int(os.getenv("AZURE_MAX_FILE_CONCURRENCY", "4"))
_AZURE_MAX_CHUNK_CONCURRENCY = int(os.getenv("AZURE_MAX_CHUNK_CONCURRENCY", "4"))


def _should_download(
    relative_path: str,
    allow_patterns: Optional[List[str]] = None,
    ignore_patterns: Optional[List[str]] = None,
) -> bool:
    """Determine whether a file should be downloaded based on allow/ignore patterns.

    Uses fnmatch semantics consistent with huggingface_hub.snapshot_download.
    Matching is performed against the relative path of the file.

    Args:
        relative_path: The relative path of the file (e.g., "subdir/model.safetensors").
        allow_patterns: If set, only files matching at least one pattern are included.
        ignore_patterns: If set, files matching any pattern are excluded.

    Returns:
        True if the file should be downloaded, False otherwise.
    """
    if allow_patterns is not None and len(allow_patterns) > 0:
        if not any(fnmatch.fnmatch(relative_path, p) for p in allow_patterns):
            return False

    if ignore_patterns is not None and len(ignore_patterns) > 0:
        if any(fnmatch.fnmatch(relative_path, p) for p in ignore_patterns):
            return False

    return True


def _parse_patterns_from_env(env_var_name: str) -> Optional[List[str]]:
    """Parse allow/ignore patterns from an environment variable.

    Supports JSON array format: '["*.safetensors", "*.json"]'
    Falls back to comma-separated: '*.safetensors,*.json'

    Returns None if the env var is not set or empty.
    """
    value = os.environ.get(env_var_name, "").strip()
    if not value:
        return None

    try:
        parsed = json.loads(value)
        if isinstance(parsed, list):
            return [str(p).strip() for p in parsed if str(p).strip()]
        return [str(parsed).strip()] if str(parsed).strip() else None
    except (json.JSONDecodeError, TypeError):
        patterns = [p.strip() for p in value.split(",") if p.strip()]
        return patterns if patterns else None


class Storage(object):
    @staticmethod
    def download_files(
        source_uris: list[str],
        out_dirs: list[str],
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> list[str]:
        download_fn = partial(
            Storage.download,
            allow_patterns=allow_patterns,
            ignore_patterns=ignore_patterns,
        )
        with ThreadPoolExecutor() as executor:
            model_dirs = list(executor.map(download_fn, source_uris, out_dirs))
        return model_dirs

    @staticmethod
    def download(
        uri: str,
        out_dir: Optional[str] = None,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        start = time.monotonic()
        Storage._update_with_storage_spec()

        if allow_patterns is None:
            allow_patterns = _parse_patterns_from_env("STORAGE_ALLOW_PATTERNS")
        if ignore_patterns is None:
            ignore_patterns = _parse_patterns_from_env("STORAGE_IGNORE_PATTERNS")

        logger.info("Copying contents of %s to local", uri)

        if allow_patterns:
            logger.info("Allow patterns: %s", allow_patterns)
        if ignore_patterns:
            logger.info("Ignore patterns: %s", ignore_patterns)

        if uri.startswith(_PVC_PREFIX) and not os.path.exists(uri):
            raise Exception(f"Cannot locate source uri {uri} for PVC")

        is_local = uri.startswith(_LOCAL_PREFIX) or os.path.exists(uri)
        if is_local:
            if out_dir is None:
                # noop if out_dir is not set and the path is local
                model_dir = Storage._download_local(uri)
            else:
                if not os.path.exists(out_dir):
                    os.mkdir(out_dir)
                model_dir = Storage._download_local(
                    uri, out_dir, allow_patterns, ignore_patterns
                )
        else:
            if out_dir is None:
                out_dir = tempfile.mkdtemp()
            elif not os.path.exists(out_dir):
                os.mkdir(out_dir)

            if uri.startswith(MODEL_MOUNT_DIRS):
                # Don't need to download models if this InferenceService is running in the multi-model
                # serving mode. The model agent will download models.
                model_dir = out_dir
            elif uri.startswith(_GCS_PREFIX):
                model_dir = Storage._download_gcs(
                    uri, out_dir, allow_patterns, ignore_patterns
                )
            elif uri.startswith(_S3_PREFIX):
                model_dir = Storage._download_s3(
                    uri, out_dir, allow_patterns, ignore_patterns
                )
            elif uri.startswith(_HDFS_PREFIX) or uri.startswith(_WEBHDFS_PREFIX):
                model_dir = Storage._download_hdfs(
                    uri, out_dir, allow_patterns, ignore_patterns
                )
            elif any(re.search(pattern, uri) for pattern in _AZURE_BLOB_RE):
                model_dir = Storage._download_azure_blob(
                    uri, out_dir, allow_patterns, ignore_patterns
                )
            elif any(re.search(pattern, uri) for pattern in _AZURE_FILE_RE):
                model_dir = Storage._download_azure_file_share(
                    uri, out_dir, allow_patterns, ignore_patterns
                )
            elif uri.startswith(_HF_PREFIX):
                model_dir = Storage._download_hf(
                    uri, out_dir, allow_patterns, ignore_patterns
                )
            elif re.search(_GIT_RE, uri):
                model_dir = Storage._download_git_repo(uri, out_dir)
            # "catch-all" pattern, should always be last
            elif re.search(_URI_RE, uri):
                model_dir = Storage._download_from_uri(uri, out_dir)
            else:
                raise Exception(
                    "Cannot recognize storage type for "
                    + uri
                    + "\n'%s', '%s', '%s', '%s' and '%s' are the current available storage type."
                    % (_GCS_PREFIX, _S3_PREFIX, _LOCAL_PREFIX, _HTTP_PREFIX, _HF_PREFIX)
                )

        logger.info("Successfully copied %s to %s", uri, out_dir)
        logger.info(f"Model downloaded in {time.monotonic() - start} seconds.")
        return model_dir

    @staticmethod
    def _update_with_storage_spec():
        storage_secret_json = json.loads(os.environ.get("STORAGE_CONFIG", "{}"))
        storage_secret_override_params = json.loads(
            os.environ.get("STORAGE_OVERRIDE_CONFIG", "{}")
        )
        if storage_secret_override_params:
            for key, value in storage_secret_override_params.items():
                storage_secret_json[key] = value

        if storage_secret_json.get("type", "") == "s3":
            for env_var, key in (
                ("AWS_ENDPOINT_URL", "endpoint_url"),
                ("AWS_ACCESS_KEY_ID", "access_key_id"),
                ("AWS_SECRET_ACCESS_KEY", "secret_access_key"),
                ("AWS_DEFAULT_REGION", "region"),
                ("AWS_CA_BUNDLE", "ca_bundle"),
                ("S3_VERIFY_SSL", "verify_ssl"),
                ("awsAnonymousCredential", "anonymous"),
            ):
                if key in storage_secret_json:
                    os.environ[env_var] = storage_secret_json.get(key)

        if (
            storage_secret_json.get("type", "") == "hdfs"
            or storage_secret_json.get("type", "") == "webhdfs"
        ):
            temp_dir = tempfile.mkdtemp()
            os.environ["HDFS_SECRET_DIR"] = temp_dir
            for key, value in storage_secret_json.items():
                mode = "w"

                # If the secret is supposed to be a file, then it was base64 encoded in the json
                if key in _HDFS_FILE_SECRETS:
                    value = base64.b64decode(value)
                    mode = "wb"

                with open(f"{temp_dir}/{key}", mode) as f:
                    f.write(value)
                    f.flush()

    @staticmethod
    def get_S3_config():
        from botocore import UNSIGNED
        from botocore.client import Config

        # Configurable timeouts and retries
        connect_timeout = int(os.getenv("S3_CONNECT_TIMEOUT", "15"))
        read_timeout = int(os.getenv("S3_READ_TIMEOUT", "30"))
        max_attempts = int(os.getenv("S3_MAX_ATTEMPTS", "3"))

        # default s3 config with timeouts
        c = Config(
            connect_timeout=connect_timeout,
            read_timeout=read_timeout,
            retries={"max_attempts": max_attempts, "mode": "standard"},
        )

        # anon environment variable defined in s3_secret.go
        anon = "true" == os.getenv("awsAnonymousCredential", "false").lower()
        # S3UseVirtualBucket environment variable defined in s3_secret.go
        # use virtual hosted-style URLs if enabled
        virtual = "true" == os.getenv("S3_USER_VIRTUAL_BUCKET", "false").lower()
        # S3UseAccelerate environment variable defined in s3_secret.go
        # use transfer acceleration if enabled
        accelerate = "true" == os.getenv("S3_USE_ACCELERATE", "false").lower()

        if anon:
            c = c.merge(Config(signature_version=UNSIGNED))
        if virtual:
            c = c.merge(Config(s3={"addressing_style": "virtual"}))
        if accelerate:
            c = c.merge(Config(s3={"use_accelerate_endpoint": accelerate}))

        # NOTE: If endpoint_url provided is legacy ("https://s3.amazonaws.com") and region is not global (us-east-1), set to virtual addressing style
        # So that request would not return PermanentRedirect due to region not in the endpoint url
        # AWS SDK retries under the hood to set the correct region when the valid virtual addressing style endpoint url is provided
        endpoint_url = os.getenv("AWS_ENDPOINT_URL")
        region = os.getenv("AWS_DEFAULT_REGION")
        if endpoint_url == "https://s3.amazonaws.com" and region not in (
            None,
            "us-east-1",
        ):
            c = c.merge(Config(s3={"addressing_style": "virtual"}))

        return c

    @staticmethod
    def _get_s3_client_kwargs():
        """
        Get the standardized boto3 client kwargs for S3 operations.

        Returns:
            dict: kwargs for creating boto3 S3 resource/client
        """
        kwargs = {"config": Storage.get_S3_config()}
        endpoint_url = os.getenv("AWS_ENDPOINT_URL")
        if endpoint_url:
            kwargs.update({"endpoint_url": endpoint_url})
        verify_ssl = os.getenv("S3_VERIFY_SSL")
        if verify_ssl:
            verify_ssl = verify_ssl.lower() not in ["0", "false"]
            kwargs.update({"verify": verify_ssl})
        else:
            verify_ssl = True

        # If verify_ssl is true, then check there is custom ca bundle cert
        # The CA bundle can be any local file in the container under the path
        # set in the AWS_CA_BUNDLE environment variable.
        # It can also be coming from a ConfigMap, in which case the filename
        # is cabundle.crt.
        if verify_ssl:
            global_ca_bundle_configmap = os.getenv("CA_BUNDLE_CONFIGMAP_NAME")
            isvc_aws_ca_bundle_path = os.getenv("AWS_CA_BUNDLE")
            ca_bundle_set = False
            if isvc_aws_ca_bundle_path and isvc_aws_ca_bundle_path != "":
                ca_bundle_set = True
                ca_bundle_full_path = isvc_aws_ca_bundle_path
            elif global_ca_bundle_configmap:
                ca_bundle_set = True
                global_ca_bundle_volume_mount_path = os.getenv(
                    "CA_BUNDLE_VOLUME_MOUNT_POINT"
                )
                ca_bundle_full_path = os.path.join(
                    global_ca_bundle_volume_mount_path, "cabundle.crt"
                )
            if ca_bundle_set:
                if os.path.exists(ca_bundle_full_path):
                    logger.info("ca bundle file(%s) exists." % (ca_bundle_full_path))
                    kwargs.update({"verify": ca_bundle_full_path})
                else:
                    raise RuntimeError(
                        "Failed to find ca bundle file(%s)." % ca_bundle_full_path
                    )
        return kwargs

    @staticmethod
    def _init_s3_worker():
        """
        Initialize S3 resources for worker processes.
        This runs once per worker process to avoid repeated initialization overhead.
        """
        global _worker_s3_resource
        try:
            import boto3

            kwargs = Storage._get_s3_client_kwargs()
            _worker_s3_resource = boto3.resource("s3", **kwargs)
        except Exception as e:
            logger.error(f"Failed to initialize S3 worker: {e}")
            _worker_s3_resource = None

    @staticmethod
    def _download_s3_object(args: Tuple) -> Tuple[bool, str, str]:
        """
        Worker function to download a single S3 object in a separate process.
        Uses pre-initialized S3 resource from _init_s3_worker().

        Args:
            args: Tuple containing (bucket_name:str, obj_key:str, target_path:str)

        Returns:
            Tuple of (success: bool, obj_key: str, error_message: str)
        """
        try:
            if _worker_s3_resource is None:
                return False, args[1], "S3 resource not initialized in worker process"

            bucket_name, obj_key, target_path = args
            bucket = _worker_s3_resource.Bucket(bucket_name)

            # Download the file
            bucket.download_file(obj_key, target_path)

            return True, obj_key, ""

        except Exception as e:
            return False, obj_key, str(e)

    @staticmethod
    def _download_s3(
        uri,
        temp_dir: str,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        import boto3
        from botocore.exceptions import ClientError, NoCredentialsError

        # Get S3 configuration using the shared helper method
        kwargs = Storage._get_s3_client_kwargs()
        parsed = urlparse(uri, scheme="s3")
        bucket_name = parsed.netloc
        bucket_path = parsed.path.lstrip("/")

        try:
            s3 = boto3.resource("s3", **kwargs)
        except ValueError as e:
            endpoint_url = os.getenv("AWS_ENDPOINT_URL", "")
            raise_storage_error("S3", uri, e, endpoint_url)

        # Collect all objects to download
        download_tasks = []
        exact_obj_found = False

        try:
            s3.meta.client.head_bucket(Bucket=bucket_name)
            bucket = s3.Bucket(bucket_name)

            for obj in bucket.objects.filter(Prefix=bucket_path):
                if obj.key.endswith("/") or obj.size == 0:
                    logger.debug("Skipping: %s", obj.key)
                    continue

                logger.info("Found S3 object: %s (%d bytes)", obj.key, obj.size)

                if bucket_path == obj.key:
                    target_key = obj.key.rsplit("/", 1)[-1]
                    exact_obj_found = True
                else:
                    target_key = obj.key.removeprefix(bucket_path).lstrip("/")

                # Apply file filtering (skip for exact object match)
                if not exact_obj_found and not _should_download(
                    target_key, allow_patterns, ignore_patterns
                ):
                    logger.info("Skipping %s due to file pattern filter", obj.key)
                    continue

                target_path = f"{temp_dir}/{target_key}"

                # Create target directory if it doesn't exist
                if not os.path.exists(dir_path := os.path.dirname(target_path)):
                    os.makedirs(dir_path, exist_ok=True)

                download_tasks.append((bucket_name, obj.key, target_path))

                # If the exact object is found, then it is sufficient to download that and break the loop
                if exact_obj_found:
                    break
        except (ClientError, NoCredentialsError) as e:
            raise_storage_error("S3", uri, e, bucket_name)

        if len(download_tasks) == 0:
            raise RuntimeError(
                "Failed to fetch model. No model found in %s." % bucket_path
            )

        num_processes = min(_S3_MAX_FILE_CONCURRENCY, len(download_tasks))

        with multiprocessing.Pool(
            processes=num_processes, initializer=Storage._init_s3_worker
        ) as pool:
            results = list(
                pool.imap_unordered(
                    Storage._download_s3_object, download_tasks, chunksize=1
                )
            )

        # Process results and handle errors
        successful_downloads = []
        failed_downloads = []

        for success, obj_key, error_msg in results:
            if success:
                successful_downloads.append(obj_key)
                logger.info("Downloaded object %s" % obj_key)
            else:
                failed_downloads.append((obj_key, error_msg))
                logger.error("Failed to download object %s: %s" % (obj_key, error_msg))

        if len(failed_downloads) > 0:
            error_details = "; ".join(
                [f"{obj}: {err}" for obj, err in failed_downloads]
            )
            raise RuntimeError(
                f"Failed to download {len(failed_downloads)} files: {error_details}"
            )

        # Unpack compressed file, supports .tgz, tar.gz and zip file formats.
        if len(successful_downloads) == 1:
            target_path = download_tasks[0][
                2
            ]  # target_path is the 3rd element in task_args
            mimetype, _ = mimetypes.guess_type(target_path)
            if mimetype in ["application/x-tar", "application/zip"]:
                temp_dir = Storage._unpack_archive_file(target_path, mimetype, temp_dir)
        return temp_dir

    @staticmethod
    def _download_hf(
        uri,
        temp_dir: str,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        from huggingface_hub import snapshot_download

        from huggingface_hub.utils import (
            RepositoryNotFoundError,
            RevisionNotFoundError,
            GatedRepoError,
            HfHubHTTPError,
        )

        components = uri[len(_HF_PREFIX) :].split("/")

        # Validate that the URI has two parts: repo and model (optional hash)
        if len(components) != 2:
            raise RuntimeError(
                "Invalid Hugging Face URI format. Expected 'hf://owner/model[:revision]', got '%s'"
                % uri
            )

        repo = components[0]
        model_part = components[1]

        if not repo:
            raise RuntimeError(
                "Hugging Face repository owner cannot be empty in URI: %s" % uri
            )
        if not model_part:
            raise RuntimeError(
                "Hugging Face model name cannot be empty in URI: %s" % uri
            )

        model, _, hash_value = model_part.partition(":")
        # Ensure model is non-empty
        if not model:
            raise RuntimeError(
                "Hugging Face model name cannot be empty in URI: %s" % uri
            )

        revision = hash_value if hash_value else None
        repo_id = f"{repo}/{model}"

        try:
            kwargs = dict(repo_id=repo_id, revision=revision, local_dir=temp_dir)
            if allow_patterns:
                kwargs["allow_patterns"] = allow_patterns
            if ignore_patterns:
                kwargs["ignore_patterns"] = ignore_patterns
            snapshot_download(**kwargs)
        except (
            RepositoryNotFoundError,
            RevisionNotFoundError,
            GatedRepoError,
            HfHubHTTPError,
        ) as e:
            raise_storage_error("HuggingFace", uri, e, repo_id)

        return temp_dir

    @staticmethod
    def _download_gcs(
        uri,
        temp_dir: str,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        from google.auth import exceptions as auth_exceptions
        from google.cloud import storage
        from google.api_core import exceptions as api_exceptions
        import copy

        try:
            storage_client = storage.Client()
        except auth_exceptions.DefaultCredentialsError:
            storage_client = storage.Client.create_anonymous_client()

        bucket_args = uri.replace(_GCS_PREFIX, "", 1).split("/", 1)
        bucket_name = bucket_args[0]
        bucket_path = bucket_args[1] if len(bucket_args) > 1 else ""

        try:
            bucket = storage_client.bucket(bucket_name)
            prefix = bucket_path
            if not prefix.endswith("/"):
                prefix = prefix + "/"
            blobs = bucket.list_blobs(prefix=prefix)
            file_count = 0

            # Shallow copy, otherwise Iterator has already started
            shallow_blobs = copy.copy(blobs)
            blob = bucket.blob(bucket_path)
            # checks if the blob is a file or a directory
            if blob.name == bucket_path and len(list(shallow_blobs)) == 0:
                dest_path = os.path.join(temp_dir, os.path.basename(bucket_path))
                logger.info("Downloading single file to: %s", dest_path)
                blob.download_to_filename(dest_path)
                file_count = 1

            else:
                for blob in blobs:
                    # Replace any prefix from the object key with temp_dir
                    subdir_object_key = blob.name.replace(bucket_path, "", 1).lstrip(
                        "/"
                    )
                    # Create necessary subdirectory to store the object locally
                    if "/" in subdir_object_key:
                        local_object_dir = os.path.join(
                            temp_dir, subdir_object_key.rsplit("/", 1)[0]
                        )
                        if not os.path.isdir(local_object_dir):
                            os.makedirs(local_object_dir, exist_ok=True)
                    if (
                        subdir_object_key.strip() != ""
                        and not subdir_object_key.endswith("/")
                    ):
                        if not _should_download(
                            subdir_object_key, allow_patterns, ignore_patterns
                        ):
                            logger.info(
                                "Skipping %s due to file pattern filter", blob.name
                            )
                            continue
                        dest_path = os.path.join(temp_dir, subdir_object_key)
                        logger.info("Downloading: %s", dest_path)
                        blob.download_to_filename(dest_path)
                        file_count += 1
        except (
            api_exceptions.Forbidden,
            api_exceptions.NotFound,
            auth_exceptions.GoogleAuthError,
        ) as e:
            raise_storage_error("GCS", uri, e, bucket_name)

        if file_count == 0:
            raise RuntimeError("Failed to fetch model. No model found in %s." % uri)

        # Unpack compressed file, supports .tgz, tar.gz and zip file formats.
        if file_count == 1:
            mimetype, _ = mimetypes.guess_type(blob.name)
            if mimetype in ["application/x-tar", "application/zip"]:
                temp_dir = Storage._unpack_archive_file(dest_path, mimetype, temp_dir)
        return temp_dir

    @staticmethod
    def _load_hdfs_configuration() -> dict:
        config = {
            "HDFS_NAMENODE": None,
            "USER_PROXY": None,
            "HDFS_ROOTPATH": None,
            "TLS_CERT": None,
            "TLS_KEY": None,
            "TLS_CA": None,
            "TLS_SKIP_VERIFY": "false",
            "HEADERS": None,
            "N_THREADS": "2",
            "KERBEROS_KEYTAB": None,
            "KERBEROS_PRINCIPAL": None,
        }

        secret_dir = _HDFS_SECRET_DIRECTORY
        if os.environ.get("HDFS_SECRET_DIR"):
            secret_dir = os.environ["HDFS_SECRET_DIR"]

        for filename in os.listdir(secret_dir):
            if filename not in config:
                continue

            # We don't read files which are supposed to be files, just save their path
            if filename in _HDFS_FILE_SECRETS:
                config[filename] = f"{secret_dir}/{filename}"
                continue

            # Read file and save value in config dict
            with open(f"{secret_dir}/{filename}") as f:
                config[filename] = f.read()

        return config

    @staticmethod
    def _download_hdfs(
        uri,
        out_dir: str,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        from krbcontext.context import krbContext
        from hdfs.ext.kerberos import Client, KerberosClient
        from hdfs.util import HdfsError

        config = Storage._load_hdfs_configuration()

        logger.info("Using the following hdfs config: %s", config)

        # Remove hdfs:// or webhdfs:// from the uri to get just the path
        # e.g. hdfs://user/me/model -> user/me/model
        if uri.startswith(_HDFS_PREFIX):
            path = uri[len(_HDFS_PREFIX) :]
        else:
            path = uri[len(_WEBHDFS_PREFIX) :]

        if not config["HDFS_ROOTPATH"]:
            path = "/" + path

        s = requests.Session()

        if config["TLS_CERT"]:
            s.cert = (config["TLS_CERT"], config["TLS_KEY"])
        # s.verify = , True, False, or CA PATH
        if config["TLS_CA"]:
            s.verify = config["TLS_CA"]
        if config["TLS_SKIP_VERIFY"].lower() == "true":
            s.verify = False

        if config["HEADERS"]:
            headers = json.loads(config["HEADERS"])
            s.headers.update(headers)

        try:
            if config["KERBEROS_PRINCIPAL"]:
                context = krbContext(
                    using_keytab=True,
                    principal=config["KERBEROS_PRINCIPAL"],
                    keytab_file=config["KERBEROS_KEYTAB"],
                )
                context.init_with_keytab()
                client = KerberosClient(
                    config["HDFS_NAMENODE"],
                    proxy=config["USER_PROXY"],
                    root=config["HDFS_ROOTPATH"],
                    session=s,
                )
            else:
                client = Client(
                    config["HDFS_NAMENODE"],
                    proxy=config["USER_PROXY"],
                    root=config["HDFS_ROOTPATH"],
                    session=s,
                )
            file_count = 0
            dest_file_path = ""

            # Check path exists and get path status
            # Raises HdfsError when path does not exist
            status = client.status(path)

            if status["type"] == "FILE":
                client.download(path, out_dir, n_threads=1)
                file_count += 1
                file_name = path.rsplit("/", 1)[-1]
                dest_file_path = f"{out_dir}/{file_name}"
            else:
                files = client.list(path)
                for f in files:
                    if not _should_download(f, allow_patterns, ignore_patterns):
                        logger.info("Skipping %s due to file pattern filter", f)
                        continue
                    client.download(
                        f"{path}/{f}", out_dir, n_threads=int(config["N_THREADS"])
                    )
                    file_count += 1
                    dest_file_path = f"{out_dir}/{f}"
        except (HdfsError, requests.exceptions.ConnectionError) as e:
            raise_storage_error("HDFS", uri, e, config["HDFS_NAMENODE"])

        if file_count == 1:
            mimetype, _ = mimetypes.guess_type(dest_file_path)
            if mimetype in ["application/x-tar", "application/zip"]:
                out_dir = Storage._unpack_archive_file(
                    dest_file_path, mimetype, out_dir
                )
        return out_dir

    @staticmethod
    async def _download_azure_blob_async(
        uri,
        out_dir: str,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        """Async Azure blob download with chunked streaming and multi-level semaphores"""
        from azure.storage.blob.aio import BlobServiceClient
        from azure.core.exceptions import (
            ClientAuthenticationError,
            ResourceNotFoundError,
            HttpResponseError,
        )

        account_name, account_url, container_name, prefix = Storage._parse_azure_uri(
            uri
        )
        logger.info(
            "Connecting to BLOB account: [%s], container: [%s], prefix: [%s]",
            account_name,
            container_name,
            prefix,
        )

        token = (
            Storage._get_azure_storage_token()
            or Storage._get_azure_storage_access_key()
        )
        if token is None:
            logger.warning(
                "Azure credentials or shared access signature token not found, retrying anonymous access"
            )

        # File-level semaphore to control concurrent file downloads
        file_semaphore = asyncio.Semaphore(_AZURE_MAX_FILE_CONCURRENCY)

        try:
            async with BlobServiceClient(
                account_url, credential=token
            ) as blob_service_client:
                container_client = blob_service_client.get_container_client(
                    container_name
                )

                # Get all blobs using flat listing (no delimiter) to get all files regardless of hierarchy
                blobs = []
                logger.info("Listing blobs with prefix: %s", prefix)
                async for blob in container_client.list_blobs(name_starts_with=prefix):
                    logger.info("Found blob: %s (%d bytes)", blob.name, blob.size)
                    if blob.size > 0:
                        relative_path = blob.name.replace(prefix, "", 1).lstrip("/")
                        if not relative_path:
                            relative_path = os.path.basename(prefix)
                        if not _should_download(
                            relative_path, allow_patterns, ignore_patterns
                        ):
                            logger.info(
                                "Skipping %s due to file pattern filter", blob.name
                            )
                            continue
                        blobs.append(blob)

                if not blobs:
                    raise RuntimeError(
                        "Failed to fetch model. No model found in %s." % uri
                    )

                # Create download tasks with semaphore control
                download_tasks = [
                    Storage._download_single_blob_async(
                        container_client, blob, out_dir, prefix, file_semaphore
                    )
                    for blob in blobs
                ]

                # Execute all downloads concurrently
                results = await asyncio.gather(*download_tasks, return_exceptions=True)

                # Check for exceptions
                file_count = 0
                dest_path = None
                for i, result in enumerate(results):
                    if isinstance(result, Exception):
                        logger.error("Blob %s failed: %s", blobs[i].name, result)
                        raise result
                    else:
                        dest_path = result
                        file_count += 1
        except (
            ClientAuthenticationError,
            ResourceNotFoundError,
            HttpResponseError,
        ) as e:
            raise_storage_error("Azure", uri, e, container_name)

        # Handle single file unpacking
        if file_count == 1:
            mimetype, _ = mimetypes.guess_type(dest_path)
            if mimetype in ["application/x-tar", "application/zip"]:
                out_dir = Storage._unpack_archive_file(dest_path, mimetype, out_dir)

        return out_dir

    @staticmethod
    async def _download_single_blob_async(
        container_client,
        blob,
        out_dir: str,
        prefix: str,
        file_semaphore: asyncio.Semaphore,
    ) -> str:
        """Download a single blob with chunked streaming"""
        async with file_semaphore:  # Limit concurrent file downloads
            file_name = blob.name.replace(prefix, "", 1).lstrip("/")
            if not file_name:
                file_name = os.path.basename(prefix)
            dest_path = os.path.join(out_dir, file_name)
            Path(os.path.dirname(dest_path)).mkdir(parents=True, exist_ok=True)

            logger.info("Downloading: %s to %s", blob.name, dest_path)

            # Get blob client and download with chunked streaming
            blob_client = container_client.get_blob_client(blob.name)

            # Download with streaming chunks to avoid memory overload
            stream = await blob_client.download_blob(
                max_concurrency=_AZURE_MAX_CHUNK_CONCURRENCY
            )

            with open(dest_path, "wb") as f:
                # Stream chunks instead of loading entire file into memory
                async for chunk in stream.chunks():
                    f.write(chunk)

            return dest_path

    @staticmethod
    def _download_azure_file_share(
        uri,
        out_dir: str,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:  # pylint: disable=too-many-locals
        from azure.storage.fileshare import ShareServiceClient
        from azure.core.exceptions import (
            ClientAuthenticationError,
            ResourceNotFoundError,
            HttpResponseError,
        )

        account_name, account_url, share_name, prefix = Storage._parse_azure_uri(uri)
        logger.info(
            "Connecting to file share account: [%s], container: [%s], prefix: [%s]",
            account_name,
            share_name,
            prefix,
        )
        access_key = Storage._get_azure_storage_access_key()
        if access_key is None:
            logger.warning(
                "Azure storage access key not found, retrying anonymous access"
            )

        try:
            share_service_client = ShareServiceClient(
                account_url, credential=access_key
            )
            share_client = share_service_client.get_share_client(share_name)
            file_count = 0
            share_files = []
            max_depth = 5
            stack = [(prefix, max_depth)]
            while stack:
                curr_prefix, depth = stack.pop()
                if depth < 0:
                    continue
                for item in share_client.list_directories_and_files(
                    directory_name=curr_prefix
                ):
                    if item.is_directory:
                        stack.append(
                            ("/".join([curr_prefix, item.name]).strip("/"), depth - 1)
                        )
                    else:
                        share_files.append((curr_prefix, item))
            for prefix, file_item in share_files:
                parts = [prefix] if prefix else []
                parts.append(file_item.name)
                file_path = "/".join(parts).lstrip("/")
                if not _should_download(file_path, allow_patterns, ignore_patterns):
                    logger.info("Skipping %s due to file pattern filter", file_path)
                    continue
                dest_path = os.path.join(out_dir, file_path)
                Path(os.path.dirname(dest_path)).mkdir(parents=True, exist_ok=True)
                logger.info("Downloading: %s to %s", file_item.name, dest_path)
                file_client = share_client.get_file_client(file_path)
                with open(dest_path, "wb+") as f:
                    data = file_client.download_file()
                    data.readinto(f)
                file_count += 1
        except (
            ClientAuthenticationError,
            ResourceNotFoundError,
            HttpResponseError,
        ) as e:
            raise_storage_error("Azure", uri, e, share_name)

        if file_count == 0:
            raise RuntimeError("Failed to fetch model. No model found in %s." % (uri))

        # Unpack compressed file, supports .tgz, tar.gz and zip file formats.
        if file_count == 1:
            mimetype, _ = mimetypes.guess_type(dest_path)
            if mimetype in ["application/x-tar", "application/zip"]:
                out_dir = Storage._unpack_archive_file(dest_path, mimetype, out_dir)
        return out_dir

    @staticmethod
    def _download_azure_blob(
        uri,
        out_dir: str,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        """Wrapper to run async blob download"""
        return asyncio.run(
            Storage._download_azure_blob_async(
                uri, out_dir, allow_patterns, ignore_patterns
            )
        )

    @staticmethod
    def _parse_azure_uri(uri):  # pylint: disable=too-many-locals
        parsed = urlparse(uri)
        account_name = parsed.netloc.split(".")[0]
        account_url = "https://{}{}".format(
            parsed.netloc, "?" + parsed.query if parsed.query else ""
        )
        object_name, prefix = parsed.path.lstrip("/").split("/", 1)
        prefix = prefix.strip("/")
        return account_name, account_url, object_name, prefix

    @staticmethod
    def _get_azure_storage_token():
        tenant_id = os.getenv("AZ_TENANT_ID", "")
        client_id = os.getenv("AZ_CLIENT_ID", "")
        client_secret = os.getenv("AZ_CLIENT_SECRET", "")

        # convert old environment variable to conform Azure defaults
        # see azure/identity/_constants.py
        if tenant_id:
            os.environ["AZURE_TENANT_ID"] = tenant_id
        if client_id:
            os.environ["AZURE_CLIENT_ID"] = client_id
        if client_secret:
            os.environ["AZURE_CLIENT_SECRET"] = client_secret

        client_id = os.getenv("AZURE_CLIENT_ID", "")
        if not client_id:
            return None

        # note the SP must have "Storage Blob Data Owner" perms for this to work
        from azure.identity import DefaultAzureCredential

        token_credential = DefaultAzureCredential()

        logger.info("Retrieved SP token credential for client_id: %s", client_id)
        return token_credential

    @staticmethod
    def _get_azure_storage_access_key():
        return os.getenv("AZURE_STORAGE_ACCESS_KEY")

    @staticmethod
    def _download_git_repo(uri: str, out_dir: str) -> str:
        """
        Supports authentication via:
        - Username in URL: https://username@host/repo.git
        - Username from GIT_USERNAME environment variable
        - Password from GIT_PASSWORD environment variable (from Kubernetes secret)
        """
        from dulwich import porcelain
        from dulwich.errors import GitProtocolError
        from urllib.parse import urlparse, urlunparse

        logger.info("Downloading Git repository %s into %s", uri, out_dir)

        parsed = urlparse(uri)
        username = None
        clean_uri = uri

        # Extract username from URL if present (format: https://username@host/repo.git)
        # Note: If password is in URL (https://user:pass@host/repo.git), it will be ignored
        # as passwords should come from Kubernetes secrets via GIT_PASSWORD env var
        if parsed.username:
            username = parsed.username

            # Reconstruct URI without username/password
            netloc = parsed.hostname
            if parsed.port:
                netloc += f":{parsed.port}"
            clean_parsed = parsed._replace(netloc=netloc)
            clean_uri = urlunparse(clean_parsed)

        if not username:
            username = os.getenv("GIT_USERNAME")

        password = os.getenv("GIT_PASSWORD")

        try:
            clone_kwargs = {"depth": 1}
            if username:
                clone_kwargs["username"] = username
            if password:
                clone_kwargs["password"] = password

            porcelain.clone(clean_uri, out_dir, **clone_kwargs)
            logger.info("git clone successful")

        except GitProtocolError as e:
            raise_storage_error("Git", uri, e, clean_uri)
        except Exception as e:
            raise_storage_error("Git", uri, e, clean_uri)

        return out_dir

    @staticmethod
    def _download_local(
        uri,
        out_dir=None,
        allow_patterns: Optional[List[str]] = None,
        ignore_patterns: Optional[List[str]] = None,
    ) -> str:
        local_path = uri.replace(_LOCAL_PREFIX, "", 1)
        if not os.path.exists(local_path):
            raise RuntimeError("Local path %s does not exist." % (uri))

        if out_dir is None:
            return local_path
        elif not os.path.isdir(out_dir):
            os.makedirs(out_dir, exist_ok=True)

        if os.path.isdir(local_path):
            local_path = os.path.join(local_path, "*")

        file_count = 0
        for src in glob.glob(local_path):
            _, tail = os.path.split(src)
            if not _should_download(tail, allow_patterns, ignore_patterns):
                logger.info("Skipping %s due to file pattern filter", src)
                continue
            dest_path = os.path.join(out_dir, tail)
            logger.info("Linking: %s to %s", src, dest_path)
            if not os.path.exists(dest_path):
                os.symlink(src, dest_path)
            else:
                logger.info("File %s already exist", dest_path)
            file_count += 1
        if file_count == 0:
            raise RuntimeError("Failed to fetch model. No model found in %s." % (uri))
        # Unpack compressed file, supports .tgz, tar.gz and zip file formats.
        if file_count == 1:
            mimetype, _ = mimetypes.guess_type(dest_path)
            if mimetype in ["application/x-tar", "application/zip"]:
                out_dir = Storage._unpack_archive_file(dest_path, mimetype, out_dir)
        return out_dir

    @staticmethod
    def _download_from_uri(uri, out_dir=None) -> str:
        url = urlparse(uri)
        filename = os.path.basename(url.path)
        # Determine if the symbol '?' exists in the path
        if mimetypes.guess_type(url.path)[0] is None and url.query != "":
            mimetype, encoding = mimetypes.guess_type(url.query)
        else:
            mimetype, encoding = mimetypes.guess_type(url.path)
        local_path = os.path.join(out_dir, filename)

        if filename == "":
            raise RuntimeError("No filename contained in URI: %s" % (uri))

        # Get header information from host url
        headers = {}
        host_uri = url.hostname

        headers_json = os.getenv(host_uri + _HEADERS_SUFFIX, "{}")
        headers = json.loads(headers_json)

        try:
            response = requests.get(uri, stream=True, headers=headers, timeout=30)
        except (requests.exceptions.ConnectionError, requests.exceptions.Timeout) as e:
            raise_storage_error("HTTP", uri, e, host_uri)

        with response:
            check_http_response(uri, response)
            zip_content_types = (
                "application/x-zip-compressed",
                "application/zip",
                "application/zip-compressed",
            )
            if mimetype == "application/zip" and not response.headers.get(
                "Content-Type", ""
            ).startswith(zip_content_types):
                raise RuntimeError(
                    "URI: %s did not respond with any of following 'Content-Type': "
                    % uri
                    + ", ".join(zip_content_types)
                )
            tar_content_types = (
                "application/x-tar",
                "application/x-gtar",
                "application/x-gzip",
                "application/gzip",
            )
            if mimetype == "application/x-tar" and not response.headers.get(
                "Content-Type", ""
            ).startswith(tar_content_types):
                raise RuntimeError(
                    "URI: %s did not respond with any of following 'Content-Type': "
                    % uri
                    + ", ".join(tar_content_types)
                )
            if (
                mimetype != "application/zip" and mimetype != "application/x-tar"
            ) and not response.headers.get("Content-Type", "").startswith(
                "application/octet-stream"
            ):
                raise RuntimeError(
                    "URI: %s did not respond with 'Content-Type': 'application/octet-stream'"
                    % uri
                )

            if encoding == "gzip":
                stream = gzip.GzipFile(fileobj=response.raw)
                local_path = os.path.join(out_dir, f"{filename}.tar")
            else:
                stream = response.raw
            with open(local_path, "wb") as out:
                shutil.copyfileobj(stream, out)

        if mimetype in ["application/x-tar", "application/zip"]:
            out_dir = Storage._unpack_archive_file(local_path, mimetype, out_dir)
        return out_dir

    @staticmethod
    def _unpack_archive_file(file_path, mimetype, target_dir=None) -> str:
        if not target_dir:
            target_dir = os.path.dirname(file_path)

        try:
            logger.info("Unpacking: %s", file_path)
            if mimetype == "application/x-tar":
                with tarfile.open(file_path, "r", encoding="utf-8") as archive:
                    archive.extractall(target_dir, filter="data")
            else:
                with zipfile.ZipFile(file_path, "r") as archive:
                    archive.extractall(target_dir)
        except (tarfile.TarError, zipfile.BadZipfile) as e:
            raise RuntimeError(
                "Failed to unpack archive file. The file format is not valid."
            ) from e
        os.remove(file_path)
        return target_dir
