# ByteDance Volcengine EMR, Copyright 2024.
#
# 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.

"""The compatible module about AbstractFileSystem in fsspec."""
import os
import re
from concurrent.futures import ThreadPoolExecutor, as_completed
from typing import Any, Optional, Union

from fsspec import AbstractFileSystem
from fsspec.implementations.local import LocalFileSystem
from fsspec.utils import other_paths

magic_check_bytes = re.compile(b"([*?[])")
magic_check = re.compile("([*?[])")


def has_magic(s: str) -> bool:
    """Check if a string has glob characters."""
    if isinstance(s, bytes):
        match = magic_check_bytes.search(s)
    else:
        match = magic_check.search(s)
    return match is not None


class EnhancedLocalFileSystem(LocalFileSystem):
    """Enhanced LocalFileSystem than fsspec's LocalFileSystem."""

    def expand_path(
        self,
        path: Union[str, list[str]],
        recursive: bool = False,
        maxdepth: Optional[int] = None,
        disable_glob: bool = False,
        **kwargs: Any,
    ) -> Union[str, list[str]]:
        """Turn one or more globs or dirs into a list of files or dirs.

        kwargs are passed to ``glob`` or ``find``, which may in turn call ``ls``
        """
        # type: ignore
        if maxdepth is not None and maxdepth < 1:
            raise ValueError("maxdepth must be at least 1")

        if isinstance(path, (str, os.PathLike)):
            out = self.expand_path(
                [path], recursive, maxdepth, disable_glob=disable_glob
            )
        else:
            out = set()  # type: ignore
            path = [self._strip_protocol(p) for p in path]
            for p in path:
                if has_magic(p) and not disable_glob:
                    bit = set(self.glob(p, maxdepth=maxdepth, **kwargs))
                    out |= bit  # type: ignore
                    if recursive:
                        # glob call above expanded one depth so if maxdepth is defined
                        # then decrement it in expand_path call below. If it is zero
                        # after decrementing then avoid expand_path call.
                        if maxdepth is not None and maxdepth <= 1:
                            continue
                        out |= set(  # type: ignore
                            self.expand_path(
                                list(bit),
                                recursive=recursive,
                                maxdepth=maxdepth - 1 if maxdepth is not None else None,
                                **kwargs,
                            )
                        )
                    continue
                elif recursive:
                    rec = set(
                        self.find(
                            p, maxdepth=maxdepth, withdirs=True, detail=False, **kwargs
                        )
                    )
                    out |= rec  # type: ignore
                if p not in out and (recursive is False or self.exists(p)):
                    # should only check once, for the root
                    out.add(p)  # type: ignore
        if not out:
            raise FileNotFoundError(path)
        return sorted(out)


class FsspecCompatibleFS(AbstractFileSystem):
    """A fsspec compatible file system.

    Used to be compatible with old version in some special methods.
    """

    def walk(  # noqa
        self,
        path: str,
        maxdepth: Optional[int] = None,
        topdown: bool = True,
        on_error: str = "omit",
        **kwargs: Any,
    ) -> Any:
        """Return all files belows path.

        Copied from fsspec(2024.9.0) to fix fsspec(2023.5.0.)

        List all files, recursing into subdirectories; output is iterator-style,
        like ``os.walk()``. For a simple list of files, ``find()`` is available.

        When topdown is True, the caller can modify the dirnames list in-place (perhaps
        using del or slice assignment), and walk() will
        only recurse into the subdirectories whose names remain in dirnames;
        this can be used to prune the search, impose a specific order of visiting,
        or even to inform walk() about directories the caller creates or renames before
        it resumes walk() again.
        Modifying dirnames when topdown is False has no effect. (see os.walk)

        Note that the "files" outputted will include anything that is not
        a directory, such as links.

        Parameters
        ----------
        path: str
            Root to recurse into
        maxdepth: int
            Maximum recursion depth. None means limitless, but not recommended
            on link-based file-systems.
        topdown: bool (True)
            Whether to walk the directory tree from the top downwards or from
            the bottom upwards.
        on_error: "omit", "raise", a collable
            if omit (default), path with exception will simply be empty;
            If raise, an underlying exception will be raised;
            if callable, it will be called with a single OSError instance as argument
        kwargs: passed to ``ls``

        """
        # type: ignore
        if maxdepth is not None and maxdepth < 1:
            raise ValueError("maxdepth must be at least 1")

        path = self._strip_protocol(path)
        full_dirs = {}
        dirs = {}
        files = {}

        detail = kwargs.pop("detail", False)
        try:
            listing = self.ls(path, detail=True, **kwargs)
        except (FileNotFoundError, OSError) as e:
            if on_error == "raise":
                raise
            elif callable(on_error):
                on_error(e)
            if detail:
                return path, {}, {}  # type: ignore
            return path, [], []  # type: ignore

        for info in listing:
            # each info name must be at least [path]/part , but here
            # we check also for names like [path]/part/
            pathname = info["name"].rstrip("/")  # type: ignore
            name = pathname.rsplit("/", 1)[-1]
            if info["type"] == "directory" and pathname != path:  # type: ignore
                # do not include "self" path
                full_dirs[name] = pathname
                dirs[name] = info
            elif pathname == path:
                # file-like with same name as give path
                files[""] = info
            else:
                files[name] = info

        if not detail:
            dirs = list(dirs)  # type: ignore
            files = list(files)  # type: ignore

        if topdown:
            # Yield before recursion if walking top down
            yield path, dirs, files

        if maxdepth is not None:
            maxdepth -= 1
            if maxdepth < 1:
                if not topdown:
                    yield path, dirs, files
                return

        with ThreadPoolExecutor() as executor:
            futures = {
                executor.submit(
                    self.walk,
                    full_dirs[d],
                    maxdepth=maxdepth,
                    detail=detail,
                    topdown=topdown,
                    **kwargs,
                ): d
                for d in dirs
            }
            for future in as_completed(futures):
                for result in future.result():
                    yield result

        if not topdown:
            # Yield after recursion if walking bottom up
            yield path, dirs, files

    def find(  # noqa #
        self,
        path: str,
        maxdepth: Optional[int] = None,
        withdirs: bool = False,
        detail: bool = False,
        **kwargs: Any,  # type: ignore
    ) -> Any:
        """List all files below path.

        Copied from fsspec(2024.9.0) to fix fsspec(2023.5.0.)

        Like posix ``find`` command without conditions

        Parameters
        ----------
        path : str
        maxdepth: int or None
            If not None, the maximum number of levels to descend
        withdirs: bool
            Whether to include directory paths in the output. This is True
            when used by glob, but users usually only want files.
        kwargs are passed to ``ls``.

        """
        # TODO: allow equivalent of -name parameter
        path = self._strip_protocol(path)
        out = {}

        # Add the root directory if withdirs is requested
        # This is needed for posix glob compliance
        if withdirs and path != "" and self.isdir(path):
            out[path] = self.info(path)

        for _, dirs, files in super().walk(path, maxdepth, detail=True, **kwargs):
            if withdirs:
                files.update(dirs)
            out.update({info["name"]: info for name, info in files.items()})
        if not out and self.isfile(path):
            # walk works on directories, but find should also return [path]
            # when path happens to be a file
            out[path] = {}
        names = sorted(out)
        if not detail:
            return names
        else:
            return {name: out[name] for name in names}

    def put(
        self,
        lpath: str,
        rpath: str,
        recursive: bool = False,
        callback: Any = None,
        maxdepth: Optional[int] = None,
        disable_glob: bool = False,
        **kwargs: Any,
    ) -> None:
        """Copy file(s) from local.

        Copies a specific file or tree of files (if recursive=True). If rpath
        ends with a "/", it will be assumed to be a directory, and target files
        will go within.

        Calls put_file for each source.
        """
        if isinstance(lpath, list) and isinstance(rpath, list):
            # No need to expand paths when both source and destination
            # are provided as lists
            rpaths = rpath
            lpaths = lpath
        else:
            from fsspec.implementations.local import make_path_posix, trailing_sep

            source_is_str = isinstance(lpath, str)
            if source_is_str:
                lpath = make_path_posix(lpath)
            fs = EnhancedLocalFileSystem()
            lpaths = fs.expand_path(
                lpath, recursive=recursive, maxdepth=maxdepth, disable_glob=disable_glob
            )
            if source_is_str and (not recursive or maxdepth is not None):
                # Non-recursive glob does not copy directories
                lpaths = [p for p in lpaths if not (trailing_sep(p) or fs.isdir(p))]
                if not lpaths:
                    return

            source_is_file = len(lpaths) == 1
            dest_is_dir = isinstance(rpath, str) and (
                trailing_sep(rpath) or self.isdir(rpath)
            )

            rpath = (
                self._strip_protocol(rpath)
                if isinstance(rpath, str)
                else [self._strip_protocol(p) for p in rpath]
            )
            exists = source_is_str and (
                (has_magic(lpath) and source_is_file)
                or (
                    (not has_magic(lpath) or disable_glob)
                    and dest_is_dir
                    and not trailing_sep(lpath)
                )
            )
            rpaths = other_paths(
                lpaths,
                rpath,
                exists=exists,
                flatten=not source_is_str,
            )

        for lpath, rpath in zip(lpaths, rpaths):
            self.put_file(lpath, rpath, **kwargs)
