import logging
import pathlib
import sys
import warnings
from typing import TYPE_CHECKING, List, Optional, Tuple, Union
from urllib.parse import quote, unquote, urlparse

from ray.data._internal.util import (
    RetryingPyFileSystem,
    _normalize_paths_to_strings,
    _resolve_custom_scheme,
)
from ray.util.annotations import RayDeprecationWarning

logger = logging.getLogger(__name__)

if TYPE_CHECKING:
    import fsspec.spec
    import pyarrow


def _get_fsspec_http_filesystem() -> "pyarrow.fs.PyFileSystem":
    """Get fsspec HTTPFileSystem wrapped in PyArrow PyFileSystem.

    Returns:
        PyFileSystem wrapping fsspec HTTPFileSystem.

    Raises:
        ImportError: If fsspec is not installed.
    """
    try:
        import fsspec  # noqa: F401
        from fsspec.implementations.http import HTTPFileSystem
    except ModuleNotFoundError:
        raise ImportError("Please install fsspec to read files from HTTP.") from None

    from pyarrow.fs import FSSpecHandler, PyFileSystem

    return PyFileSystem(FSSpecHandler(HTTPFileSystem()))


def _validate_and_wrap_filesystem(
    filesystem: Optional[
        Union["pyarrow.fs.FileSystem", "fsspec.spec.AbstractFileSystem"]
    ],
) -> Optional["pyarrow.fs.FileSystem"]:
    """Validate filesystem and wrap fsspec filesystems in PyArrow.

    Args:
        filesystem: Filesystem to validate and potentially wrap. Can be None,
            a pyarrow.fs.FileSystem, or an fsspec.spec.AbstractFileSystem.

    Returns:
        None if filesystem is None, otherwise a pyarrow.fs.FileSystem
        (either the original if already PyArrow, or wrapped if fsspec).

    Raises:
        TypeError: If filesystem is not None and not a valid pyarrow or fsspec filesystem.
    """
    if filesystem is None:
        return None

    from pyarrow.fs import FileSystem

    if isinstance(filesystem, FileSystem):
        return filesystem

    try:
        import fsspec  # noqa: F401
    except ModuleNotFoundError:
        raise TypeError("fsspec is not installed") from None

    if not isinstance(filesystem, fsspec.spec.AbstractFileSystem):
        raise TypeError(
            f"Filesystem must conform to pyarrow.fs.FileSystem or "
            f"fsspec.spec.AbstractFileSystem, got: {type(filesystem).__name__}"
        )

    from pyarrow.fs import FSSpecHandler, PyFileSystem

    return PyFileSystem(FSSpecHandler(filesystem))


def _try_resolve_with_encoding(
    path: str,
    filesystem: Optional["pyarrow.fs.FileSystem"],
) -> Tuple["pyarrow.fs.FileSystem", str]:
    """Try resolving a path with URL encoding for special characters.

    This handles paths with special characters like ';', '?', '#' that
    may cause URI parsing errors.

    Args:
        path: The path to resolve.
        filesystem: Optional filesystem to validate against.

    Returns:
        Tuple of (resolved_filesystem, resolved_path).
    """
    from pyarrow.fs import _resolve_filesystem_and_path

    encoded_path = quote(path, safe="/:", errors="ignore")
    resolved_filesystem, resolved_path = _resolve_filesystem_and_path(
        encoded_path, filesystem
    )
    return resolved_filesystem, unquote(resolved_path, errors="ignore")


def _has_file_extension(path: str, extensions: Optional[List[str]]) -> bool:
    """Check if a path has a file extension in the provided list.

    Examples:
        >>> _has_file_extension("foo.csv", ["csv"])
        True
        >>> _has_file_extension("foo.CSV", ["csv"])
        True
        >>> _has_file_extension("foo.CSV", [".csv"])
        True
        >>> _has_file_extension("foo.csv", ["json", "jsonl"])
        False
        >>> _has_file_extension("foo.csv", None)
        True

    Args:
        path: The path to check.
        extensions: A list of extensions to check against. If `None`, any extension is
            considered valid.

    Returns:
        ``True`` if ``path`` ends with one of the provided extensions (or
        ``extensions`` is ``None``), otherwise ``False``.
    """
    assert extensions is None or isinstance(extensions, list), type(extensions)

    if extensions is None:
        return True

    # If the user-specified extensions don't contain a leading dot, we add it here
    extensions = [
        f".{ext.lower()}" if not ext.startswith(".") else ext.lower()
        for ext in extensions
    ]
    # Ignore query components when checking extensions (for example,
    # versioned object-store paths like `...parquet?versionId=...`).
    # Keep `#` untouched because it can be part of object keys.
    parsed_path = path.split("?", 1)[0]
    return any(parsed_path.lower().endswith(ext) for ext in extensions)


# Mapping from URI schemes to compatible filesystem type_name values.
# Used to validate that a cached filesystem is compatible with a given URI scheme
# before attempting to use it, avoiding silent failures from PyArrow when the
# wrong filesystem type is passed to _resolve_filesystem_and_path.
_SCHEME_TO_FS_TYPE_NAMES = {
    "": ("local",),  # No scheme = local filesystem
    "file": ("local",),  # file:// = local filesystem
    "s3": ("s3",),  # s3:// = S3 filesystem
    "s3a": ("s3",),  # s3a:// = S3 filesystem (Hadoop compat)
    "gs": ("gcs",),  # gs:// = GCS filesystem
    "gcs": ("gcs",),  # gcs:// = GCS filesystem
    "hdfs": ("hdfs",),  # hdfs:// = Hadoop filesystem
    "viewfs": ("hdfs",),  # viewfs:// = Hadoop filesystem
    "abfs": ("abfs",),  # abfs:// = Azure Blob FileSystem
    "abfss": ("abfs",),  # abfss:// = Azure Blob FileSystem (TLS)
    "http": ("py",),  # http:// = fsspec HTTP (wrapped in PyFileSystem)
    "https": ("py",),  # https:// = fsspec HTTP (wrapped in PyFileSystem)
}


def _is_filesystem_compatible_with_scheme(
    filesystem: "pyarrow.fs.FileSystem",
    scheme: str,
) -> bool:
    """Check if a filesystem is compatible with a URI scheme.

    Uses PyArrow's `type_name` property for reliable filesystem type detection.
    This prevents silently using the wrong filesystem for a URI, which can result
    in malformed paths or incorrect behavior.

    Args:
        filesystem: The PyArrow filesystem to check.
        scheme: The URI scheme (e.g., 's3', 'gs', 'http', 'file', '').

    Returns:
        True if the filesystem can handle the scheme, False otherwise.
    """
    # Get expected type names for this scheme
    expected_types = _SCHEME_TO_FS_TYPE_NAMES.get(scheme.lower())
    if expected_types is None:
        # Unknown scheme (e.g., abfs://, az://, custom protocols) - trust user's filesystem
        # This preserves backward compatibility for custom filesystems
        return True

    # Unwrap RetryingPyFileSystem to get the underlying filesystem's type
    from ray.data._internal.util import RetryingPyFileSystem

    unwrapped = (
        filesystem.unwrap()
        if isinstance(filesystem, RetryingPyFileSystem)
        else filesystem
    )

    # Get the actual filesystem type
    fs_type = unwrapped.type_name

    # For PyFileSystem (fsspec wrappers), check the inner fsspec protocol
    # rather than relying on type_name alone, since all fsspec wrappers
    # share type_name "py" regardless of the underlying protocol.
    if fs_type in ("py", "RetryingPyFileSystem") or fs_type.startswith("py::"):
        from pyarrow.fs import FSSpecHandler, PyFileSystem

        actual_fs = filesystem
        if isinstance(actual_fs, RetryingPyFileSystem):
            actual_fs = actual_fs.unwrap()

        # After unwrapping, the inner filesystem may be a native PyArrow
        # filesystem (e.g., S3FileSystem) rather than a PyFileSystem wrapper.
        # Fall back to direct type_name matching in that case.
        if not isinstance(actual_fs, PyFileSystem):
            return actual_fs.type_name in expected_types

        if isinstance(actual_fs.handler, FSSpecHandler):
            inner_fs = actual_fs.handler.fs
            protocol = getattr(inner_fs, "protocol", None)
            if protocol is not None:
                if isinstance(protocol, str):
                    protocol = (protocol,)
                # Match scheme against fsspec protocol(s)
                if scheme in protocol:
                    return True
                # For bare paths (empty scheme), trust user-provided filesystem
                if scheme == "":
                    return True

        # Fallback: check HTTP
        if scheme in ("http", "https"):
            return _is_http_filesystem(filesystem)

        return False

    # Direct match for native PyArrow filesystems (s3, gcs, local, hdfs, etc.)
    return fs_type in expected_types


_AZURE_BLOB_HTTPS_SUFFIXES = (".blob.core.windows.net", ".dfs.core.windows.net")


def _rewrite_azure_blob_https_url(path: str) -> Tuple[str, Optional[str]]:
    """Rewrite Azure Blob/DFS HTTPS URLs to ``abfs://<container>/<blob>`` form.

    PyArrow's scheme auto-detection only recognises ``abfs://`` / ``abfss://`` for
    Azure; an ``https://<account>.blob.core.windows.net/<container>/<blob>`` URL
    falls through to the fsspec HTTP filesystem and fetches anonymously, which
    fails for private blobs.

    We deliberately avoid emitting the Hadoop-style
    ``abfs://<container>@<account>.dfs.core.windows.net/<path>`` form. PyArrow's
    C++ ``AzureOptions::FromUri`` parses it correctly when the resolver builds a
    filesystem from the URI alone — but ``_resolve_filesystem_and_path`` *skips*
    that parser whenever a filesystem is already supplied (the common case for
    ``download(..., filesystem=AzureFileSystem(...))``). In that branch PyArrow
    strips the scheme and treats the authority as part of the path, so requests
    land on ``https://<fs.account>.blob.core.windows.net/<cont>@<acct>.dfs.core.windows.net/<blob>``
    and 404. The bare ``abfs://<container>/<blob>`` form sidesteps this entirely
    because there is no authority to misinterpret.

    Returns ``(rewritten_path, account_name)``. ``account_name`` is the host
    prefix extracted from the original URL so the caller can construct an
    ``AzureFileSystem(account_name=...)`` when the user did not pass an explicit
    filesystem. ``account_name`` is ``None`` when the path was not rewritten.

    URLs carrying a query string (e.g. a SAS token) are returned unchanged: the
    SAS already carries auth, and the fsspec HTTP path will fetch them with the
    signature intact.
    """
    if not path.startswith("https://"):
        return path, None

    parsed = urlparse(path, allow_fragments=False)
    if parsed.query:
        return path, None

    # `parsed.hostname` strips explicit port and userinfo (and lower-cases the
    # host per RFC 3986), so the suffix check works for URLs that include
    # `:443` or `user:pass@`.
    host = parsed.hostname
    if not host:
        return path, None
    suffix = next(
        (s for s in _AZURE_BLOB_HTTPS_SUFFIXES if host.endswith(s)),
        None,
    )
    if suffix is None or len(host) == len(suffix):
        return path, None

    account = host[: -len(suffix)]
    object_path = parsed.path.lstrip("/")
    if "/" not in object_path:
        # No <container>/<blob> split available — leave it alone.
        return path, None

    container, blob_path = object_path.split("/", 1)
    if not container or not blob_path:
        return path, None

    return f"abfs://{container}/{blob_path}", account


def _resolve_single_path_with_fallback(
    path: str,
    filesystem: Optional["pyarrow.fs.FileSystem"] = None,
) -> Tuple["pyarrow.fs.FileSystem", str]:
    """Resolve a single path with filesystem, with fallback to re-resolution on error.

    This is a helper for lazy filesystem resolution. If a filesystem is provided,
    it first validates that the filesystem type is compatible with the URI scheme,
    then attempts to resolve the path. If the filesystem is incompatible or
    resolution fails, it re-resolves without the cached filesystem.

    Args:
        path: A single file/directory path.
        filesystem: Optional cached filesystem from previous resolution.

    Returns:
        Tuple of (resolved_filesystem, resolved_path).

    Raises:
        ValueError: If path resolution fails.
        ImportError: If required dependencies are missing.
    """
    import pyarrow as pa
    import pyarrow.fs
    from pyarrow.fs import _resolve_filesystem_and_path

    path = _resolve_custom_scheme(path)

    # Only rewrite Azure Blob HTTPS URLs when this PyArrow build actually has
    # AzureFileSystem support. Otherwise the abfs:// form we would produce
    # cannot be resolved by PyArrow, and we would block the original https://
    # URL from falling through to the fsspec HTTP filesystem below.
    azure_fs_cls = getattr(pyarrow.fs, "AzureFileSystem", None)
    if azure_fs_cls is not None:
        path, azure_account = _rewrite_azure_blob_https_url(path)
    else:
        azure_account = None

    # Validate/wrap filesystem if needed
    try:
        filesystem = _validate_and_wrap_filesystem(filesystem)
    except TypeError as e:
        raise ValueError(f"Invalid filesystem provided: {e}") from e

    # If we rewrote an Azure HTTPS URL and the caller did not supply a
    # filesystem, build one with the account name we extracted. AzureFileSystem
    # still resolves credentials (account_key, SAS, default credential chain)
    # from the environment; supplying a filesystem explicitly to download() with
    # custom credentials remains the override path for non-default auth.
    # _resolve_paths_and_filesystem caches the filesystem returned by the first
    # path and reuses it, so the auto-injected AzureFileSystem is shared across
    # the rest of the column — multi-account columns are not supported on the
    # PyArrow path; use obstore for that case.
    if azure_account is not None and filesystem is None:
        filesystem = azure_fs_cls(account_name=azure_account)

    # Parse scheme to validate filesystem compatibility
    parsed = urlparse(path, allow_fragments=False)
    scheme = parsed.scheme.lower() if parsed.scheme else ""

    # Check HTTP scheme FIRST - PyArrow doesn't support HTTP/HTTPS natively
    if scheme in ("http", "https"):
        # If we have a compatible cached HTTP filesystem, use it
        if filesystem is not None and _is_filesystem_compatible_with_scheme(
            filesystem, scheme
        ):
            return filesystem, path
        # Otherwise create a new HTTP filesystem
        try:
            resolved_filesystem = _get_fsspec_http_filesystem()
            resolved_path = path
            return resolved_filesystem, resolved_path
        except ImportError as import_error:
            raise ImportError(
                f"Cannot resolve HTTP path '{path}': {import_error}"
            ) from import_error

    # Try with provided filesystem only if scheme is compatible (fast path for cached FS)
    if filesystem is not None and _is_filesystem_compatible_with_scheme(
        filesystem, scheme
    ):
        try:
            _, resolved_path = _resolve_filesystem_and_path(path, filesystem)
            # Return the wrapped filesystem we passed in.
            return filesystem, resolved_path
        except Exception:
            # Fall through to full resolution without cached filesystem
            pass

    # Full resolution without cached filesystem
    try:
        resolved_filesystem, resolved_path = _resolve_filesystem_and_path(path, None)
    except (pa.lib.ArrowInvalid, ValueError) as original_error:
        # Try URL encoding for paths with special characters that may cause parsing issues
        try:
            resolved_filesystem, resolved_path = _try_resolve_with_encoding(path, None)
        except (pa.lib.ArrowInvalid, ValueError, TypeError) as encoding_error:
            # If encoding doesn't help, raise with both errors for full context
            raise ValueError(
                f"Failed to resolve path '{path}'. Initial error: {original_error}. "
                f"URL encoding fallback also failed: {encoding_error}"
            ) from original_error
    except TypeError as e:
        raise ValueError(f"The path: '{path}' has an invalid type {e}") from e

    return resolved_filesystem, resolved_path


def _resolve_paths_and_filesystem(
    paths: Union[str, List[str]],
    filesystem: Optional["pyarrow.fs.FileSystem"] = None,
) -> Tuple[List[str], "pyarrow.fs.FileSystem"]:
    """
    Resolves and normalizes all provided paths, infers a filesystem from the
    paths and assumes that all paths use the same filesystem.

    Args:
        paths: A single file/directory path or a list of file/directory paths.
            A list of paths can contain both files and directories.
        filesystem: The filesystem implementation that should be used for
            reading these files. If None, a filesystem will be inferred. If not
            None, the provided filesystem will still be validated against all
            filesystems inferred from the provided paths to ensure
            compatibility.

    Returns:
        A pair ``(resolved_paths, filesystem)``. *resolved_paths* lists the
        normalized paths for each input path that resolved successfully, in
        order.

        If *filesystem* was ``None``, the returned *filesystem* is set from
        ``resolved_filesystem`` on the first successful path and is left
        unchanged on later iterations whenever it is already non-``None``.

        If *filesystem* was not ``None``, the returned value is always that
        same validated instance, even when ``_resolve_single_path_with_fallback``
        inferred a different filesystem for a given path. Callers should pass
        ``None`` or a filesystem compatible with the path URIs so returned paths
        and filesystem stay consistent.

        All paths are assumed to use one storage backend; mixing unrelated URI
        schemes in a single call is unsupported and may fail when reading.
    """
    paths = _normalize_paths_to_strings(paths)
    if any(path.lower().startswith("local:") for path in paths):
        warnings.warn(
            "`local://` paths in Ray Data are deprecated and will be removed after "
            "January 2027. Use shared or cloud storage for distributed execution.",
            RayDeprecationWarning,
            stacklevel=3,
        )

    # Validate/wrap filesystem upfront so we return a proper PyArrow filesystem
    filesystem = _validate_and_wrap_filesystem(filesystem)

    resolved_paths = []
    for path in paths:
        try:
            resolved_filesystem, resolved_path = _resolve_single_path_with_fallback(
                path, filesystem
            )
        except (ValueError, ImportError) as e:
            logger.warning(f"Failed to resolve path '{path}': {e}, skipping")
            continue

        if filesystem is None:
            filesystem = resolved_filesystem

        # If the PyArrow filesystem is handled by a fsspec HTTPFileSystem, the protocol/
        # scheme of paths should not be unwrapped/removed, because HTTPFileSystem
        # expects full file paths including protocol/scheme. This is different behavior
        # compared to other file system implementation in pyarrow.fs.FileSystem.
        if not _is_http_filesystem(resolved_filesystem):
            resolved_path = _unwrap_protocol(resolved_path)

        resolved_path = resolved_filesystem.normalize_path(resolved_path)
        resolved_paths.append(resolved_path)

    return resolved_paths, filesystem


def _split_uri(uri: str):
    """Split a URI into (store_url, path) for use with obstore.

    e.g. "s3://my-bucket/a/b/c.jpg"               -> ("s3://my-bucket", "a/b/c.jpg")
         "https://host.com/a/b?X-Amz-Signature=x" -> ("https://host.com", "a/b?X-Amz-Signature=x")

    The query string is preserved so signed URLs (e.g. pre-signed S3 HTTPS)
    reach obstore intact. Semicolons in object keys normally appear in
    ``parsed.path`` (not ``parsed.params``) for typical ``urlparse`` output.

    Only the first leading ``/`` after the authority (as reported in
    ``parsed.path``) is removed. Extra leading slashes belong to the object
    key (e.g. ``s3://bucket//abs/key`` -> key ``/abs/key``), so
    ``str.lstrip("/")`` is not used.
    """
    parsed = urlparse(uri, allow_fragments=False)
    store_url = f"{parsed.scheme}://{parsed.netloc}"
    raw_path = parsed.path
    path = raw_path[1:] if raw_path.startswith("/") else raw_path
    if parsed.query:
        path = f"{path}?{parsed.query}"
    return store_url, path


def _is_http_filesystem(fs: "pyarrow.fs.FileSystem") -> bool:
    """Return whether ``fs`` is a PyFileSystem handled by a fsspec HTTPFileSystem."""
    from pyarrow.fs import FSSpecHandler, PyFileSystem

    # Try to import HTTPFileSystem
    try:
        from fsspec.implementations.http import HTTPFileSystem
    except ModuleNotFoundError:
        return False

    if isinstance(fs, RetryingPyFileSystem):
        fs = fs.unwrap()

    if not isinstance(fs, PyFileSystem):
        return False

    return isinstance(fs.handler, FSSpecHandler) and isinstance(
        fs.handler.fs, HTTPFileSystem
    )


def _unwrap_protocol(path):
    """
    Slice off any protocol prefixes on path.
    """
    if sys.platform == "win32" and _is_local_windows_path(path):
        # Represent as posix path such that downstream functions properly handle it.
        # This is executed when 'file://' is NOT included in the path.
        return pathlib.Path(path).as_posix()

    parsed = urlparse(path, allow_fragments=False)  # support '#' in path
    params = ";" + parsed.params if parsed.params else ""  # support ';' in path
    query = "?" + parsed.query if parsed.query else ""  # support '?' in path
    netloc = parsed.netloc
    if parsed.scheme == "s3" and "@" in parsed.netloc:
        # If the path contains an @, it is assumed to be an anonymous
        # credentialed path, and we need to strip off the credentials.
        netloc = parsed.netloc.split("@")[-1]

    parsed_path = parsed.path
    # urlparse prepends the path with a '/'. This does not work on Windows
    # so if this is the case strip the leading slash.
    if (
        sys.platform == "win32"
        and not netloc
        and len(parsed_path) >= 3
        and parsed_path[0] == "/"  # The problematic leading slash
        and parsed_path[1].isalpha()  # Ensure it is a drive letter.
        and parsed_path[2:4] in (":", ":/")
    ):
        parsed_path = parsed_path[1:]

    return netloc + parsed_path + params + query


def _filesystem_root_from_uri(uri: str) -> str:
    """The path form the filesystem backing ``uri`` expects to be given.

    Returns what ``pyarrow.fs.FileSystem.from_uri`` would hand back as its
    path, without also *building* a filesystem the way ``from_uri`` does. That
    matters because for an ``s3://`` URI ``from_uri`` resolves the bucket's
    region against real AWS: it fails outright against any S3-compatible
    endpoint (MinIO, moto, Ceph) even with ``AWS_ENDPOINT_URL`` set, and costs a
    network round-trip every time it's called.

    Scheme conventions, matching ``from_uri``:

    * ``s3``/``s3a``/``gs``/``gcs``: ``bucket/key``.
    * ``abfs``/``abfss``: ``container/key``. The storage account belongs to the
      ``AzureFileSystem`` object rather than the path, so the
      ``@account.dfs.core.windows.net`` host is dropped -- leaving it in points
      writes at a container literally named ``container@account...``.
    * local paths: unchanged.
    """
    parsed = urlparse(uri, allow_fragments=False)
    if parsed.scheme in ("abfs", "abfss") and "@" in parsed.netloc:
        container = parsed.netloc.split("@")[0]
        return f"{container}{parsed.path}"
    return _unwrap_protocol(uri)


def _is_http_url(path) -> bool:
    parsed = urlparse(path)
    return parsed.scheme in ("http", "https")


def _is_local_windows_path(path: str) -> bool:
    """Determines if path is a Windows file-system location."""
    if sys.platform != "win32":
        return False

    if len(path) >= 1 and path[0] == "\\":
        return True
    if (
        len(path) >= 3
        and path[1] == ":"
        and (path[2] == "/" or path[2] == "\\")
        and path[0].isalpha()
    ):
        return True
    return False