_storage.py

from __future__ import annotations

import functools
import os
import pathlib
import random
import tempfile
import typing
from typing import AsyncGenerator, Optional
from uuid import UUID

import fsspec
import obstore
from fsspec.asyn import AsyncFileSystem
from fsspec.utils import get_protocol
from obstore.exceptions import GenericError
from obstore.fsspec import FsspecStore, register
from obstore.store import ObjectStore

from flyte._initialize import get_storage
from flyte._logging import logger
from flyte.errors import InitializationError, OnlyAsyncIOSupportedError

from ._remote_fs import FlyteFS

if typing.TYPE_CHECKING:
    from obstore import AsyncReadableFile, AsyncWritableFile

_OBSTORE_SUPPORTED_PROTOCOLS = ["s3", "gs", "abfs", "abfss"]
BATCH_SIZE = int(os.getenv("FLYTE_IO_BATCH_SIZE", str(32)))

# Multipart upload sizing. GCS and S3 both hard-cap a multipart upload at 10,000 parts, so a fixed
# part size imposes a hard file-size ceiling (obstore's 5 MiB default => ~48.8 GiB). We instead
# auto-scale the part size with the file size so uploads never run out of part numbers, while small
# files keep small parts. Nothing here is user-tunable -- it just works.
#
# The floor and concurrency mirror obstore's own put/put_async defaults (chunk_size=5242880,
# max_concurrency=12), defined here:
# https://github.com/developmentseed/obstore/blob/c6279ce19f73eef89321179d55930d0dd7ba6a62/obstore/python/obstore/store.py#L500-L501
_UPLOAD_CHUNK_FLOOR = 5 * 2**20  # 5 MiB -- obstore's own default; floor for small/multipart files
_UPLOAD_MAX_CONCURRENCY = 12  # obstore's own default
_MAX_SAFE_PARTS = 9000  # margin below the cloud 10,000-part hard limit


def _compute_upload_chunk_size(file_size: int | None) -> int:
    """
    Pick a multipart part size that keeps the part count under the cloud 10,000-part limit.

    Returns the 5 MiB floor for small/unknown files, and scales up just enough (ceil division) for
    large files so `file_size / chunk_size <= _MAX_SAFE_PARTS`.
    """
    if not file_size:
        return _UPLOAD_CHUNK_FLOOR
    needed = -(-file_size // _MAX_SAFE_PARTS)  # integer ceil division
    return max(_UPLOAD_CHUNK_FLOOR, needed)


def _is_obstore_supported_protocol(protocol: str | tuple[str, ...]) -> bool:
    """
    Check if the given protocol is supported by obstore.

    Args:
        protocol: Protocol to check.

    Returns:
        True if the protocol is supported, False otherwise.
    """
    return protocol in _OBSTORE_SUPPORTED_PROTOCOLS


def is_remote(path: typing.Union[pathlib.Path | str]) -> bool:
    """
    Let's find a replacement
    """
    protocol = get_protocol(str(path))
    if protocol is None:
        return False
    return protocol != "file"


def strip_file_header(path: str) -> str:
    """
    Drops file:// if it exists from the file
    """
    if path.startswith("file://"):
        return path.replace("file://", "", 1)
    return path


def get_random_local_path(file_path_or_file_name: pathlib.Path | str | None = None) -> pathlib.Path:
    """
    Use file_path_or_file_name, when you want a random directory, but want to preserve the leaf file name
    """
    local_tmp = pathlib.Path(tempfile.mkdtemp(prefix="flyte-tmp-"))
    key = UUID(int=random.getrandbits(128)).hex
    tmp_folder = local_tmp / key
    tail = ""
    if file_path_or_file_name:
        _, tail = os.path.split(file_path_or_file_name)
    if tail:
        tmp_folder.mkdir(parents=True, exist_ok=True)
        return tmp_folder / tail
    local_tmp.mkdir(parents=True, exist_ok=True)
    return tmp_folder


def get_random_local_directory() -> pathlib.Path:
    """
    Returns:
        pathlib.Path: a random directory
    """
    _dir = get_random_local_path(None)
    pathlib.Path(_dir).mkdir(parents=True, exist_ok=True)
    return _dir


def get_configured_fsspec_kwargs(
    protocol: typing.Optional[str] = None, anonymous: bool = False
) -> typing.Dict[str, typing.Any]:
    if protocol:
        # Try to get storage config safely - may not be initialized for local operations
        try:
            storage_config = get_storage()
        except InitializationError:
            storage_config = None

        match protocol:
            case "s3":
                # If the protocol is s3, we can use the s3 filesystem
                from flyte.storage import S3

                if storage_config and isinstance(storage_config, S3):
                    return storage_config.get_fsspec_kwargs(anonymous=anonymous)

                return S3.auto().get_fsspec_kwargs(anonymous=anonymous)
            case "gs":
                # If the protocol is gs, we can use the gs filesystem
                from flyte.storage import GCS

                if storage_config and isinstance(storage_config, GCS):
                    return storage_config.get_fsspec_kwargs(anonymous=anonymous)

                return GCS.auto().get_fsspec_kwargs(anonymous=anonymous)
            case "abfs" | "abfss":
                # If the protocol is abfs or abfss, we can use the abfs filesystem
                from flyte.storage import ABFS

                if storage_config and isinstance(storage_config, ABFS):
                    return storage_config.get_fsspec_kwargs(anonymous=anonymous)

                return ABFS.auto().get_fsspec_kwargs(anonymous=anonymous)
            case _:
                return {}

    # If no protocol, return args from storage config if set
    storage_config = get_storage()
    if storage_config:
        return storage_config.get_fsspec_kwargs(anonymous)

    return {}


@functools.lru_cache(maxsize=None)
def _obstore_filesystem_class(protocol: str) -> type[FsspecStore]:
    """
    Build (and cache) a per-protocol subclass of obstore's ``FsspecStore``, mirroring what
    ``obstore.fsspec.register(protocol, asynchronous=True)`` would register -- but without
    touching the global fsspec registry.

    Subclasses of fsspec's ``AbstractFileSystem`` keep fsspec's instance caching (the
    ``_Cached`` metaclass) even when instantiated directly, so constructing these classes
    behaves the same as resolving them through ``fsspec.filesystem(...)``.
    """
    return type(
        f"FlyteFsspecStore_{protocol}",
        (FsspecStore,),
        {"protocol": protocol, "asynchronous": True},
    )


def get_underlying_filesystem(
    protocol: typing.Optional[str] = None,
    anonymous: bool = False,
    path: typing.Optional[str] = None,
    **kwargs,
) -> fsspec.AbstractFileSystem:
    if protocol is None:
        # If protocol is None, get it from the path
        protocol = get_protocol(typing.cast(str, path))

    configured_kwargs = get_configured_fsspec_kwargs(protocol, anonymous=anonymous)
    configured_kwargs.update(kwargs)

    if _is_obstore_supported_protocol(protocol):
        # Instantiate flyte's own obstore-backed filesystem directly instead of resolving the
        # protocol through the global fsspec registry: the configured kwargs are obstore-shaped
        # (retry_config, client_options, config), and another SDK sharing the process (e.g.
        # flyte v1 / flytekit) may have registered s3fs/gcsfs/adlfs for these protocols.
        return _obstore_filesystem_class(protocol)(**configured_kwargs)

    return fsspec.filesystem(protocol, **configured_kwargs)


def _get_anonymous_filesystem(from_path):
    """Get the anonymous file system if needed."""
    return get_underlying_filesystem(get_protocol(from_path), anonymous=True, asynchronous=True)


async def _get_obstore_bypass(
    from_path: str,
    to_path: str | pathlib.Path,
    recursive: bool = False,
    exclude: list[str] | None = None,
    **kwargs,
) -> str:
    from flyte.storage._parallel_reader import ObstoreParallelReader

    # The obstore-backed filesystem is duck-typed (callers guard on hasattr for _split_path/_construct_store).
    fs = typing.cast(typing.Any, get_underlying_filesystem(path=from_path))
    bucket, prefix = fs._split_path(from_path)  # pylint: disable=W0212
    store: ObjectStore = fs._construct_store(bucket)

    download_kwargs = {}
    if "chunk_size" in kwargs:
        download_kwargs["chunk_size"] = kwargs["chunk_size"]
    if "max_concurrency" in kwargs:
        download_kwargs["max_concurrency"] = kwargs["max_concurrency"]

    reader = ObstoreParallelReader(store, **download_kwargs)
    target_path = pathlib.Path(to_path) if isinstance(to_path, str) else to_path

    # if recursive, just download the prefix to the target path
    if recursive:
        logger.debug(f"Downloading recursively {prefix=} to {target_path=}")
        await reader.download_files(
            prefix,
            target_path,
            exclude=exclude,
        )
        return str(to_path)

    # if not recursive, we need to split out the file name from the prefix
    else:
        path_for_reader = pathlib.Path(prefix).name
        final_prefix = pathlib.Path(prefix).parent
        logger.debug(f"Downloading single file {final_prefix=}, {path_for_reader=} to {target_path=}")
        await reader.download_files(
            final_prefix,
            target_path.parent,
            path_for_reader,
            destination_file_name=target_path.name,
        )
        return str(target_path)


async def get(from_path: str, to_path: Optional[str | pathlib.Path] = None, recursive: bool = False, **kwargs) -> str:
    if not to_path:
        name = pathlib.Path(from_path).name  # may need to be adjusted for windows
        to_path = get_random_local_path(file_path_or_file_name=name)
        logger.debug(f"Storing file from {from_path} to {to_path}")
    else:
        # Only apply directory logic for single files (not recursive)
        if not recursive:
            to_path_str = str(to_path)
            # Check for trailing separator BEFORE converting to Path (which normalizes and removes it)
            ends_with_sep = to_path_str.endswith(os.sep)
            to_path_obj = pathlib.Path(to_path)

            # If path ends with os.sep or is an existing directory, append source filename
            if ends_with_sep or (to_path_obj.exists() and to_path_obj.is_dir()):
                source_filename = pathlib.Path(from_path).name  # may need to be adjusted for windows
                to_path = to_path_obj / source_filename
        # For recursive=True, keep to_path as-is (it's the destination directory for contents)

    file_system = get_underlying_filesystem(path=from_path)

    # Check if we should use obstore bypass
    if (
        _is_obstore_supported_protocol(file_system.protocol)
        and hasattr(file_system, "_split_path")
        and hasattr(file_system, "_construct_store")
    ):
        return await _get_obstore_bypass(from_path, to_path, recursive, **kwargs)

    try:
        return await _get_from_filesystem(file_system, from_path, to_path, recursive=recursive, **kwargs)
    except (OSError, GenericError) as oe:
        logger.debug(f"Error in getting {from_path} to {to_path}, recursive: {recursive}, error: {oe}")
        if isinstance(file_system, AsyncFileSystem):
            try:
                exists = await file_system._exists(from_path)  # pylint: disable=W0212
            except GenericError:
                # for obstore, as it does not raise FileNotFoundError in fsspec but GenericError
                # force it to try get_filesystem(anonymous=True)
                exists = True
        else:
            exists = file_system.exists(from_path)
        if not exists:
            raise AssertionError(f"Unable to load data from {from_path}")
        file_system = _get_anonymous_filesystem(from_path)
        logger.debug(f"Attempting anonymous get with {file_system}")
        return await _get_from_filesystem(file_system, from_path, to_path, recursive=recursive, **kwargs)


async def _get_from_filesystem(
    file_system: fsspec.AbstractFileSystem,
    from_path: str | pathlib.Path,
    to_path: str | pathlib.Path,
    recursive: bool,
    **kwargs,
):
    if isinstance(file_system, AsyncFileSystem):
        dst = await file_system._get(str(from_path), str(to_path), recursive=recursive, **kwargs)  # pylint: disable=W0212
    else:
        dst = file_system.get(str(from_path), str(to_path), recursive=recursive, **kwargs)

    if isinstance(dst, (str, pathlib.Path)):
        return dst
    return str(to_path)


async def _put_single_obstore(store: ObjectStore, local_path: str | pathlib.Path, remote_key: str) -> None:
    """Upload a single local file to `remote_key` with an auto-sized multipart part size."""
    chunk_size = _compute_upload_chunk_size(os.path.getsize(local_path))
    # Passing a pathlib.Path streams the file from disk rather than reading it all into memory.
    await store.put_async(
        remote_key,
        pathlib.Path(local_path),
        chunk_size=chunk_size,
        max_concurrency=_UPLOAD_MAX_CONCURRENCY,
    )


async def _put_obstore_bypass(from_path: str, to_path: str, recursive: bool = False) -> str:
    """
    Upload via the obstore native API so we can control multipart part sizing.

    fsspec's obstore backend (`_put_file`) calls `store.put_async` without a `chunk_size`,
    leaving it pinned to obstore's 5 MiB default and thus a ~48.8 GiB hard ceiling (10,000 parts).
    This bypass computes the part size from each file's size so uploads of any size succeed. Mirrors
    the download-side `_get_obstore_bypass` / `ObstoreParallelReader` pattern.
    """
    import asyncio

    # The obstore-backed filesystem is duck-typed (callers guard on hasattr for _split_path/_construct_store).
    fs = typing.cast(typing.Any, get_underlying_filesystem(path=to_path))
    bucket, prefix = fs._split_path(to_path)  # pylint: disable=W0212
    store: ObjectStore = fs._construct_store(bucket)

    if not recursive:
        logger.debug(f"Uploading single file {from_path} to {to_path}")
        await _put_single_obstore(store, from_path, prefix)
        return to_path

    # Recursive: walk the local tree and upload each file, preserving its path relative to from_path.
    source_root = pathlib.Path(from_path)
    files = [p for p in source_root.rglob("*") if p.is_file()]
    logger.debug(f"Uploading {len(files)} files recursively from {from_path} to {to_path}")

    prefix_path = pathlib.PurePosixPath(prefix)
    semaphore = asyncio.Semaphore(BATCH_SIZE)

    async def _upload(local_file: pathlib.Path) -> None:
        relative = local_file.relative_to(source_root).as_posix()
        remote_key = str(prefix_path / relative)
        async with semaphore:
            await _put_single_obstore(store, local_file, remote_key)

    await asyncio.gather(*(_upload(f) for f in files))
    return to_path


async def put(
    from_path: str,
    to_path: Optional[str] = None,
    recursive: bool = False,
    batch_size: Optional[int] = None,
    **kwargs,
) -> str:
    if not to_path:
        from flyte._context import internal_ctx

        ctx = internal_ctx()
        name = pathlib.Path(from_path).name
        to_path = ctx.raw_data.get_random_remote_path(file_name=name)

    if not batch_size:
        batch_size = BATCH_SIZE

    file_system = get_underlying_filesystem(path=to_path)
    from_path = strip_file_header(from_path)

    # Check if we should use obstore bypass (to control multipart part sizing -- see
    # _put_obstore_bypass). Guarded identically to the download bypass in get().
    if (
        _is_obstore_supported_protocol(file_system.protocol)
        and hasattr(file_system, "_split_path")
        and hasattr(file_system, "_construct_store")
    ):
        return await _put_obstore_bypass(from_path, to_path, recursive)

    if isinstance(file_system, AsyncFileSystem):
        dst = await file_system._put(from_path, to_path, recursive=recursive, batch_size=batch_size, **kwargs)  # pylint: disable=W0212
    else:
        dst = file_system.put(from_path, to_path, recursive=recursive, **kwargs)
    if isinstance(dst, (str, pathlib.Path)):
        return str(dst)
    else:
        return to_path


async def _open_obstore_bypass(path: str, mode: str = "rb", **kwargs) -> AsyncReadableFile | AsyncWritableFile:
    """
    Simple obstore bypass for opening files. No fallbacks, obstore only.
    """

    # The obstore-backed filesystem is duck-typed (callers guard on hasattr for _split_path/_construct_store).
    fs = typing.cast(typing.Any, get_underlying_filesystem(path=path))
    bucket, file_path = fs._split_path(path)  # pylint: disable=W0212
    store: ObjectStore = fs._construct_store(bucket)

    file_handle: AsyncReadableFile | AsyncWritableFile

    if "w" in mode:
        attributes = kwargs.pop("attributes", {})
        buffer_size = 10 * 2**20
        if "buffer_size" in kwargs:
            buffer_size = kwargs.pop("buffer_size")
        if "chunk_size" in kwargs:
            buffer_size = kwargs.pop("chunk_size")
        file_handle = obstore.open_writer_async(store, file_path, attributes=attributes, buffer_size=buffer_size)
    else:  # read mode
        buffer_size = kwargs.pop("buffer_size", 10 * 2**20)
        file_handle = await obstore.open_reader_async(store, file_path, buffer_size=buffer_size)
    return file_handle


async def open(path: str, mode: str = "rb", **kwargs) -> AsyncReadableFile | AsyncWritableFile:
    """
    Asynchronously open a file and return an async context manager.
    This function checks if the underlying filesystem supports obstore bypass.
    If it does, it uses obstore to open the file. Otherwise, it falls back to
    the standard _open function which uses AsyncFileSystem.

    It will raise NotImplementedError if neither obstore nor AsyncFileSystem is supported.
    """
    fs = get_underlying_filesystem(path=path)

    # Check if we should use obstore bypass
    if _is_obstore_supported_protocol(fs.protocol) and hasattr(fs, "_split_path") and hasattr(fs, "_construct_store"):
        return await _open_obstore_bypass(path, mode, **kwargs)

    # Fallback to normal open
    if isinstance(fs, AsyncFileSystem):
        return await fs.open_async(path, mode, **kwargs)

    raise OnlyAsyncIOSupportedError(f"Filesystem {fs} does not support async operations")


async def put_stream(
    data_iterable: typing.AsyncIterable[bytes] | bytes, *, name: str | None = None, to_path: str | None = None, **kwargs
) -> str:
    """
    Put a stream of data to a remote location. This is useful for streaming data to a remote location.
    Example usage:
    ```python
    import flyte.storage as storage
    storage.put_stream(iter([b'hello']), name="my_file.txt")
    OR
    storage.put_stream(iter([b'hello']), to_path="s3://my_bucket/my_file.txt")
    ```

    Args:
        data_iterable: Iterable of bytes to be streamed.
        name: Name of the file to be created. If not provided, a random name will be generated.
        to_path: Path to the remote location where the data will be stored.
        kwargs: Additional arguments to be passed to the underlying filesystem.

    Returns:
        str: The path to the remote location where the data was stored.
    """
    if not to_path:
        from flyte._context import internal_ctx

        ctx = internal_ctx()
        to_path = ctx.raw_data.get_random_remote_path(file_name=name)

    # Check if we should use obstore bypass
    fs = get_underlying_filesystem(path=to_path)
    try:
        file_handle = typing.cast("AsyncWritableFile", await open(to_path, "wb", **kwargs))
        if isinstance(data_iterable, bytes):
            await file_handle.write(data_iterable)
        else:
            async for data in data_iterable:
                await file_handle.write(data)
        await file_handle.close()
        return str(to_path)
    except OnlyAsyncIOSupportedError:
        pass

    # Fallback to normal open
    file_handle_io: typing.IO = fs.open(to_path, mode="wb", **kwargs)
    if isinstance(data_iterable, bytes):
        file_handle_io.write(data_iterable)
    else:
        async for data in data_iterable:
            file_handle_io.write(data)
    file_handle_io.close()

    return str(to_path)


async def get_stream(path: str, chunk_size=10 * 2**20, **kwargs) -> AsyncGenerator[bytes, None]:
    """
    Get a stream of data from a remote location.
    This is useful for downloading streaming data from a remote location.
    Example usage:
    ```python
    import flyte.storage as storage
    async for chunk in storage.get_stream(path="s3://my_bucket/my_file.txt"):
        process(chunk)
    ```

    Args:
        path: Path to the remote location where the data will be downloaded.
        kwargs: Additional arguments to be passed to the underlying filesystem.
        chunk_size: Size of each chunk to be read from the file.

    Returns:
        An async iterator that yields chunks of bytes.
    """
    # Check if we should use obstore bypass
    fs = get_underlying_filesystem(path=path)
    if _is_obstore_supported_protocol(fs.protocol) and hasattr(fs, "_split_path") and hasattr(fs, "_construct_store"):
        # Set buffer_size for obstore if chunk_size is provided
        if "buffer_size" not in kwargs:
            kwargs["buffer_size"] = chunk_size
        file_handle = typing.cast("AsyncReadableFile", await _open_obstore_bypass(path, "rb", **kwargs))
        while chunk := await file_handle.read():
            yield bytes(chunk)
        return

    # Fallback to normal open
    if "block_size" not in kwargs:
        kwargs["block_size"] = chunk_size

    if isinstance(fs, AsyncFileSystem):
        file_handle = await fs.open_async(path, "rb", **kwargs)
        while chunk := await file_handle.read():
            yield chunk
        await file_handle.close()
        return

    file_handle = fs.open(path, "rb", **kwargs)
    while chunk := file_handle.read():
        yield chunk
    file_handle.close()


def join(*paths: str) -> str:
    """
    Join multiple paths together. This is a wrapper around os.path.join.
    # TODO replace with proper join with fsspec root etc

    Args:
        paths: Paths to be joined.
    """
    return str(os.path.join(*paths))


async def exists(path: str, **kwargs) -> bool:
    """
    Check if a path exists.

    Args:
        path: Path to be checked.
        kwargs: Additional arguments to be passed to the underlying filesystem.

    Returns:
        True if the path exists, False otherwise.
    """
    try:
        fs = get_underlying_filesystem(path=path, **kwargs)
        if isinstance(fs, AsyncFileSystem):
            _ = await fs._info(path)
            return True
        _ = fs.info(path)
        return True
    except FileNotFoundError:
        return False


def exists_sync(path: str, **kwargs) -> bool:
    try:
        fs = get_underlying_filesystem(path=path, **kwargs)
        _ = fs.info(path)
        return True
    except FileNotFoundError:
        return False


def get_credentials_error(uri: str, protocol: str) -> str:
    # Check for common credential issues
    if protocol == "s3":
        return (
            f"Failed to download data from {uri}. "
            f"S3 credentials are required to access the data at {uri}. "
            "Please set the following environment variables: AWS_ACCESS_KEY_ID, AWS_SECRET_ACCESS_KEY"
        )
    elif protocol in ("gs", "gcs"):
        return (
            f"Failed to download data from {uri}. "
            f"GCS credentials are required to access the data at {uri}. "
            f"Please set the following environment variable: GOOGLE_APPLICATION_CREDENTIALS"
        )
    elif protocol in ("abfs", "abfss"):
        return (
            f"Failed to download data from {uri}. "
            f"Azure credentials are required to access the data at {uri}. "
            "Please set the following environment variables: AZURE_STORAGE_ACCOUNT_NAME, "
            "AZURE_STORAGE_ACCOUNT_KEY"
        )
    raise ValueError(f"Unsupported protocol: {protocol}")


def _register_obstore_for_missing_protocols() -> None:
    """
    Register obstore with fsspec only for protocols that have no other usable implementation.

    flyte v1 (flytekit) may share the process (e.g. a v1 task importing ``flyte`` to launch v2
    runs) and owns these protocols via s3fs/gcsfs/adlfs; clobbering the global registry breaks
    its I/O because obstore rejects s3fs/gcsfs/adlfs-only kwargs (e.g. ``cache_regions``).
    flyte's own I/O never resolves these protocols through the registry (see
    ``get_underlying_filesystem``), so this global registration only exists so that third-party
    libraries (pandas, pyarrow, ...) can still resolve e.g. ``s3://`` paths in images that
    don't ship s3fs/gcsfs/adlfs.
    """
    for protocol in _OBSTORE_SUPPORTED_PROTOCOLS:
        if protocol in fsspec.registry:
            # Someone (e.g. flytekit v1) already registered an implementation; leave it alone.
            continue
        try:
            # fsspec resolves known_implementations lazily, so an empty registry does not mean
            # no implementation exists. This probes whether one (e.g. s3fs) is importable, and
            # registers it in fsspec.registry as a side effect.
            fsspec.get_filesystem_class(protocol)
        except (ImportError, ValueError):
            register(protocol, asynchronous=True)


_register_obstore_for_missing_protocols()
fsspec.register_implementation("flyte", FlyteFS, clobber=True)