
"""Typing annotation utils.
Aadapted from https://github.com/tensorflow/datasets/blob/v4.4.0/tensorflow_datasets/core/utils/type_utils.py
"""


import abc
import os
import typing
from typing import Any, AnyStr, Dict, Iterator, List, Optional, Sequence, Tuple, Type, TypeVar, Union

from ..utils import lazy_import
tf = lazy_import('tensorflow')

# pylint: disable=g-import-not-at-top
try:
    from typing import Protocol
except ImportError:
    import typing_extensions
    Protocol = typing_extensions.Protocol
# pylint: enable=g-import-not-at-top

# Accept both `str` and `pathlib.Path`-like
PathLike = Union[str, os.PathLike]
PathLikeCls = (str, os.PathLike)    # Used in `isinstance`

T = TypeVar('T')

# Note: `TupleOrList` avoid abiguity from `Sequence` (`str` is `Sequence[str]`,
# `bytes` is `Sequence[int]`).
TupleOrList = Union[Tuple[T, ...], List[T]]

TreeDict = Union[T, Dict[str, 'TreeDict']]    # pytype: disable=not-supported-yet
Tree = Union[T, TupleOrList['Tree'], Dict[str, 'Tree']]    # pytype: disable=not-supported-yet

Tensor = Union[tf.Tensor, tf.SparseTensor, tf.RaggedTensor]

Dim = Optional[int]
Shape = TupleOrList[Dim]

JsonValue = Union[str, bool, int, float, None, List['JsonValue'],
                                    Dict[str, 'JsonValue'],    # pytype: disable=not-supported-yet
                                 ]
Json = Dict[str, JsonValue]

# Types for the tfrecord example construction.

Key = Union[int, str, bytes]
KeySerializedExample = Tuple[Key, bytes]    # `(key, serialized_proto)`

# pytype: disable=ignored-abstractmethod


class PurePath(Protocol):
    """Protocol for pathlib.PurePath-like API."""
    parts: Tuple[str, ...]
    drive: str
    root: str
    anchor: str
    name: str
    suffix: str
    suffixes: List[str]
    stem: str

    # pylint: disable=multiple-statements,line-too-long

    def __new__(cls: Type[T], *args: PathLike) -> T:
        raise NotImplementedError

    def __fspath__(self) -> str:
        raise NotImplementedError

    def __hash__(self) -> int:
        raise NotImplementedError

    def __lt__(self, other: 'PurePath') -> bool:
        raise NotImplementedError

    def __le__(self, other: 'PurePath') -> bool:
        raise NotImplementedError

    def __gt__(self, other: 'PurePath') -> bool:
        raise NotImplementedError

    def __ge__(self, other: 'PurePath') -> bool:
        raise NotImplementedError

    def __truediv__(self: T, key: PathLike) -> T:
        raise NotImplementedError

    def __rtruediv__(self: T, key: PathLike) -> T:
        raise NotImplementedError

    def __bytes__(self) -> bytes:
        raise NotImplementedError

    def as_posix(self) -> str:
        raise NotImplementedError

    def as_uri(self) -> str:
        raise NotImplementedError

    def is_absolute(self) -> bool:
        raise NotImplementedError

    def is_reserved(self) -> bool:
        raise NotImplementedError

    def match(self, path_pattern: str) -> bool:
        raise NotImplementedError

    def relative_to(self: T, *other: PathLike) -> T:
        raise NotImplementedError

    def with_name(self: T, name: str) -> T:
        raise NotImplementedError

    def with_suffix(self: T, suffix: str) -> T:
        raise NotImplementedError

    def joinpath(self: T, *other: PathLike) -> T:
        raise NotImplementedError

    @property
    def parents(self: T) -> Sequence[T]:
        raise NotImplementedError

    @property
    def parent(self: T) -> T:
        raise NotImplementedError
    
    @property
    def absolute_parent(self: T) -> T:
        raise NotImplementedError
    
    @property
    def is_cloud(self: T) -> bool:
        raise NotImplementedError

    @property
    def is_gcs(self: T) -> bool:
        raise NotImplementedError

    # py3.9 backport of PurePath.is_relative_to.
    def is_relative_to(self, *other: PathLike) -> bool:
        """Return True if the path is relative to another path or False."""
        try:
            self.relative_to(*other)
            return True
        except ValueError:
            return False

    # pylint: enable=multiple-statements,line-too-long


class ReadOnlyPath(PurePath, Protocol):
    """Protocol for read-only methods of pathlib.Path-like API.
    See [pathlib.Path](https://docs.python.org/3/library/pathlib.html)
    documentation.
    """

    def __new__(cls: Type[T], *args: PathLike) -> T:
        if cls in (ReadOnlyPath, ReadWritePath):
            from tensorflow_datasets.core.utils import generic_path    # pytype: disable=import-error    # pylint: disable=g-import-not-at-top
            return generic_path.as_path(*args)
        else:
            return super().__new__(cls, *args)

    @abc.abstractmethod
    def exists(self) -> bool:
        """Returns True if self exists."""

    @abc.abstractmethod
    def is_dir(self) -> bool:
        """Returns True if self is a dir."""

    def is_file(self) -> bool:
        """Returns True if self is a file."""
        return not self.is_dir()

    @abc.abstractmethod
    def iterdir(self: T) -> Iterator[T]:
        """Iterates over the directory."""

    @abc.abstractmethod
    def glob(self: T, pattern: str) -> Iterator[T]:
        """Yielding all matching files (of any kind)."""
        # Might be able to implement using `iterdir` (recursivelly for `rglob`).

    def rglob(self: T, pattern: str) -> Iterator[T]:
        """Yielding all matching files recursivelly (of any kind)."""
        return self.glob(f'**/{pattern}')

    def expanduser(self: T) -> T:
        """Returns a new path with expanded `~` and `~user` constructs."""
        if '~' not in self.parts:    # pytype: disable=attribute-error
            return self
        raise NotImplementedError

    @abc.abstractmethod
    def resolve(self: T, strict: bool = False) -> T:
        """Returns the absolute path."""

    @abc.abstractmethod
    def open(
            self,
            mode: str = 'r',
            encoding: Optional[str] = None,
            errors: Optional[str] = None,
            **kwargs: Any,
    ) -> typing.IO[AnyStr]:
        """Opens the file."""
    
    def read(self, mode='rb') -> bytes:
        with self.open(mode=mode) as f:
            return f.read()
    
    def readlines(self) -> List[str]:
        with self.open('r') as f:
            return f.readlines()

    def read_bytes(self) -> bytes:
        """Reads contents of self as bytes."""
        with self.open('rb') as f:
            return f.read()

    def read_text(self, encoding: Optional[str] = None) -> str:
        """Reads contents of self as bytes."""
        with self.open('r', encoding=encoding) as f:
            return f.read()

    def format(self: T, *args: Any, **kwargs: Any) -> T:
        """Apply `str.format()` to the path."""
        return type(self)(os.fspath(self).format(*args, **kwargs))    # pytype: disable=not-instantiable


class ReadWritePath(ReadOnlyPath, Protocol):
    """Protocol for pathlib.Path-like API.
    See [pathlib.Path](https://docs.python.org/3/library/pathlib.html)
    documentation.
    """

    @abc.abstractmethod
    def mkdir(
            self,
            mode: int = 0o777,
            parents: bool = False,
            exist_ok: bool = False,
    ) -> None:
        """Create a new directory at this given path."""

    @abc.abstractmethod
    def rmdir(self) -> None:
        """Remove the empty directory at this given path."""

    @abc.abstractmethod
    def rmtree(self) -> None:
        """Remove the directory, including all sub-files."""

    @abc.abstractmethod
    def unlink(self, missing_ok: bool = False) -> None:
        """Remove this file or symbolic link."""

    def write_bytes(self, data: bytes) -> None:
        """Writes content as bytes."""
        with self.open('wb') as f:
            return f.write(data)

    def write_text(
            self,
            data: str,
            encoding: Optional[str] = None,
            errors: Optional[str] = None,
    ) -> None:
        """Writes content as str."""
        with self.open('w') as f:
            return f.write(data)

    def touch(self, mode: int = 0o666, exist_ok: bool = True) -> None:
        """Create a file at this given path."""
        del mode    # Unused
        if self.exists():
            if exist_ok:
                return
            else:
                raise FileExistsError(f'{self} already exists.')
        self.write_text('')

    @abc.abstractmethod
    def rename(self: T, target: PathLike) -> T:
        """Renames the path."""

    @abc.abstractmethod
    def replace(self: T, target: PathLike) -> T:
        """Overwrites the destination path."""

    @abc.abstractmethod
    def copy(self: T, dst: PathLike, overwrite: bool = False) -> T:
        """Copy the current file to the given destination."""