# Source code for ubelt.util_zip

    """
    Abstractions for working with zipfiles and archives

    This may be renamed to util_archive in the future.

    The :func:`ubelt.split_archive` works with paths that reference a file inside
    of an archive (e.g. a zipfile). It splits it into two parts, the full path to
    the archive and then the path to the file inside of the archive. By convention
    these are separated with either a pathsep or a colon.

    The :func:`ubelt.zopen` works to open a file that lives inside of an archive
    without the user needing to worry about extracting it first. When possible it
    will read it directly from the archive, but in some cases it may extract it to
    a temporary directory first.
    """
    import io
    import os
    from os.path import exists, join
    from ubelt.util_mixins import NiceRepr

    __all__ = ['zopen', 'split_archive']

    def split_archive(fpath, ext='.zip'):
        """
        If fpath specifies a file inside a zipfile, it breaks it into two parts the
        path to the zipfile and the internal path in the zipfile.

        Args:
            fpath (str | PathLike): path that specifies a path inside of an archive
            ext (str): archive extension

        Returns:
            Tuple[str, str | None]

        Example:
            >>> split_archive('/a/b/foo.txt')
            >>> split_archive('/a/b/foo.zip/bar.txt')
            >>> split_archive('/a/b/foo.zip/baz/biz.zip/bar.py')
            >>> split_archive('archive.zip')
            >>> import ubelt as ub
            >>> split_archive(ub.Path('/a/b/foo.zip/baz/biz.zip/bar.py'))
            >>> split_archive('/a/b/foo.zip/baz.pt/bar.zip/bar.zip', '.pt')

        TODO:
            Fix got/want for win32

            (None, None)
            ('/a/b/foo.zip', 'bar.txt')
            ('/a/b/foo.zip/baz/biz.zip', 'bar.py')
            ('archive.zip', None)
            ('/a/b/foo.zip/baz/biz.zip', 'bar.py')
            ('/a/b/foo.zip/baz.pt', 'bar.zip/bar.zip')
        """
        import re
        fpath = os.fspath(fpath)
        # fpath = os.fspath(fpath)
        pat = '({}[{}/:])'.format(re.escape(ext), re.escape(os.path.sep))
        # pat = r'(\'' + ext + '[' + re.escape(os.path.sep) + '/:])'
        parts = re.split(pat, fpath, flags=re.IGNORECASE)
        if len(parts) > 2:
            archivepath = ''.join(parts[:-1])[:-1]
            internal = parts[-1]
        elif len(parts) == 1:
            archivepath = parts[0]
            if not archivepath.endswith(ext):
                archivepath = None
            internal = None
        else:  # nocover
            raise AssertionError('impossible state')
        return archivepath, internal

    class zopen(NiceRepr):
        """
        An abstraction of the normal :func:`open` function that can also handle
        reading data directly inside of zipfiles.

        This is a file-object like interface [FileObj] --- i.e. it supports the
        read and write methods to an underlying resource.

        Can open a file normally or open a file within a zip file (readonly).
        Tries to read from memory only, but will extract to a tempfile if necessary.

        Just treat the zipfile like a directory,
        e.g. /path/to/myzip.zip/compressed/path.txt OR?
        e.g. /path/to/myzip.zip:compressed/path.txt

        References:
            .. [FileObj] https://docs.python.org/3/glossary.html#term-file-object

        TODO:
            - [ ] Fast way to open a base zipfile, query what is inside, and
                  then choose a file to further zopen (and passing along the same
                  open zipfile reference maybe?).
            - [ ] Write mode in some restricted setting?

        Attributes:
            name (str | PathLike):
                path to a file or reference to an item in a zipfile.

        Example:
            >>> from ubelt.util_zip import *  # NOQA
            >>> import pickle
            >>> import ubelt as ub
            >>> dpath = ub.Path.appdir('ubelt/tests/util_zip').ensuredir()
            >>> dpath = ub.Path(dpath)
            >>> data_fpath = dpath / 'test.pkl'
            >>> data = {'demo': 'data'}
            >>> with open(str(data_fpath), 'wb') as file:
            >>>     pickle.dump(data, file)
            >>> # Write data
            >>> import zipfile
            >>> zip_fpath = dpath / 'test_zip.archive'
            >>> stl_w_zfile = zipfile.ZipFile(os.fspath(zip_fpath), mode='w')
            >>> stl_w_zfile.write(os.fspath(data_fpath), os.fspath(data_fpath.relative_to(dpath)))
            >>> stl_w_zfile.close()
            >>> stl_r_zfile = zipfile.ZipFile(os.fspath(zip_fpath), mode='r')
            >>> stl_r_zfile.namelist()
            >>> stl_r_zfile.close()
            >>> # Test zopen
            >>> self = zopen(zip_fpath / 'test.pkl', mode='rb', ext='.archive')
            >>> print(self._split_archive())
            >>> print(self.namelist())
            >>> self.close()
            >>> self = zopen(zip_fpath / 'test.pkl', mode='rb', ext='.archive')
            >>> recon1 = pickle.loads(self.read())
            >>> self.close()
            >>> self = zopen(zip_fpath / 'test.pkl', mode='rb', ext='.archive')
            >>> recon2 = pickle.load(self)
            >>> self.close()
            >>> assert recon1 == recon2
            >>> assert recon1 is not recon2

        Example:
            >>> # Test we can load json data from a zipfile
            >>> from ubelt.util_zip import *  # NOQA
            >>> import ubelt as ub
            >>> import json
            >>> import zipfile
            >>> dpath = ub.Path.appdir('ubelt/tests/util_zip').ensuredir()
            >>> infopath = join(dpath, 'info.json')
            >>> ub.writeto(infopath, '{"x": "1"}')
            >>> zippath = join(dpath, 'infozip.zip')
            >>> internal = 'folder/info.json'
            >>> with zipfile.ZipFile(zippath, 'w') as myzip:
            >>>     myzip.write(infopath, internal)
            >>> fpath = zippath + '/' + internal
            >>> # Test context manager
            >>> with zopen(fpath, 'r') as self:
            >>>     info2 = json.load(self)
            >>>     assert info2['x'] == '1'
            >>> # Test outside of context manager
            >>> self = zopen(fpath, 'r')
            >>> print(self._split_archive())
            >>> info2 = json.load(self)
            >>> assert info2['x'] == '1'
            >>> # Test nice repr (with zfile)
            >>> print('self = {!r}'.format(self))
            >>> self.close()

        Example:
            >>> # Coverage tests --- move to unit-test
            >>> from ubelt.util_zip import *  # NOQA
            >>> import ubelt as ub
            >>> import json
            >>> import zipfile
            >>> dpath = ub.Path.appdir('ubelt/tests/util_zip').ensuredir()
            >>> textpath = join(dpath, 'seekable_test.txt')
            >>> text = chr(10).join(['line{}'.format(i) for i in range(10)])
            >>> ub.writeto(textpath, text)
            >>> zippath = join(dpath, 'seekable_test.zip')
            >>> internal = 'folder/seekable_test.txt'
            >>> with zipfile.ZipFile(zippath, 'w') as myzip:
            >>>     myzip.write(textpath, internal)
            >>> ub.delete(textpath)
            >>> fpath = zippath + '/' + internal
            >>> # Test seekable
            >>> self_seekable = zopen(fpath, 'r', seekable=True)
            >>> assert self_seekable.seekable()
            >>> self_seekable.seek(8)
            >>> assert self_seekable.readline() == 'ne1' + chr(10)
            >>> assert self_seekable.readline() == 'line2' + chr(10)
            >>> self_seekable.seek(8)
            >>> assert self_seekable.readline() == 'ne1' + chr(10)
            >>> assert self_seekable.readline() == 'line2' + chr(10)
            >>> # Test non-seekable?
            >>> # Sometimes non-seekable files are still seekable
            >>> maybe_seekable = zopen(fpath, 'r', seekable=False)
            >>> if maybe_seekable.seekable():
            >>>     maybe_seekable.seek(8)
            >>>     assert maybe_seekable.readline() == 'ne1' + chr(10)
            >>>     assert maybe_seekable.readline() == 'line2' + chr(10)
            >>>     maybe_seekable.seek(8)
            >>>     assert maybe_seekable.readline() == 'ne1' + chr(10)
            >>>     assert maybe_seekable.readline() == 'line2' + chr(10)

        Example:
            >>> # More coverage tests --- move to unit-test
            >>> from ubelt.util_zip import *  # NOQA
            >>> import ubelt as ub
            >>> import pytest
            >>> dpath = ub.Path.appdir('ubelt/tests/util_zip').ensuredir()
            >>> with pytest.raises(OSError):
            >>>     self = zopen('', 'r')
            >>> # Test open non-zip existing file
            >>> existing_fpath = join(dpath, 'exists.json')
            >>> ub.writeto(existing_fpath, '{"x": "1"}')
            >>> self = zopen(existing_fpath, 'r')
            >>> assert self.read() == '{"x": "1"}'
            >>> # Test dir
            >>> dir(self)
            >>> # Test nice
            >>> print(self)
            >>> print('self = {!r}'.format(self))
            >>> self.close()
            >>> # Test open non-zip non-existing file
            >>> nonexisting_fpath = join(dpath, 'does-not-exist.txt')
            >>> ub.delete(nonexisting_fpath)
            >>> with pytest.raises(OSError):
            >>>     self = zopen(nonexisting_fpath, 'r')
            >>> with pytest.raises(NotImplementedError):
            >>>     self = zopen(nonexisting_fpath, 'w')
            >>> # Test nice-repr
            >>> self = zopen(existing_fpath, 'r')
            >>> print('self = {!r}'.format(self))
            >>> # pathological
            >>> self = zopen(existing_fpath, 'r')
            >>> self._handle = None
            >>> dir(self)
        """
        def __init__(self, fpath, mode='r', seekable=False, ext='.zip'):
            """
            Args:
                fpath (str | PathLike):
                    path to a file, or a special path that denotes both a
                    path to a zipfile and a path to a archived file inside of
                    the zipfile.

                mode (str):
                    Currently only "r" - readonly mode is supported

                seekable (bool):
                    If True, attempts to force "seekability" of the underlying
                    file-object, for compressed files this will first extract
                    the file to a temporary location on disk.  If False, any underlying
                    compressed file will be opened directly which may result in the
                    object being non-seekable.

                ext (str):
                    The extension of the zipfile. Modify this is a non-standard
                    extension is used (e.g. for torch packages).
            """
            self.fpath = fpath
            self.ext = ext
            self.name = fpath
            self.mode = mode
            self._seekable = seekable
            self._zfpath = None  # points to the base zipfile (if appropriate)
            self._temp_dpath = None  # for temporary extraction
            self._zfile_read = None  # underlying opened zipfile object
            # The _handle pointer should be a file-like object that this zopen
            # object impersonate, by forwarding most every getattr call to it.
            self._handle = None
            self._open()

        @property
        def zfile(self):
            """
            Access the underlying archive file
            """
            if self._zfile_read is None:
                import zipfile
                archivefile, internal = self._split_archive()
                myzip = zipfile.ZipFile(archivefile, 'r')
                self._zfile_read = myzip
            return self._zfile_read

        def namelist(self):
            """
            Lists the contents of this zipfile
            """
            myzip = self.zfile
            namelist = myzip.namelist()
            return namelist

        def __nice__(self):
            if self._zfpath is None:
                return 'handle={}, mode={}'.format(str(self._handle), self.mode)
            else:
                return 'handle={} in zipfpath={}, mode={}'.format(self._handle, self._zfpath, self.mode)

        def __getattr__(self, key):
            # Expose attributes of wrapped handle
            if hasattr(self._handle, key):
                assert self._handle is not self
                return getattr(self._handle, key)
            raise AttributeError(key)

        def __dir__(self):
            # Expose attributes of wrapped handle
            zopen_attributes = {
                'namelist',
                'zfile',
            }
            keyset = set(dir(super(zopen, self)))
            keyset.update(set(self.__dict__.keys()))
            if self._handle is not None:
                keyset.update(set(dir(self._handle)))
            return sorted(keyset | zopen_attributes)

        def _cleanup(self):
            # print('self._cleanup = {!r}'.format(self._cleanup))
            if self._handle is not None:
                if not getattr(self, 'closed', True):
                    closemethod = getattr(self, 'close', None)
                    if closemethod is not None:  # nocover
                        closemethod()
                    closemethod = None
            self._handle = None
            if self._temp_dpath and exists(self._temp_dpath):
                # os.unlink(self._temp_dpath)
                from ubelt.util_io import delete
                delete(self._temp_dpath)

        def __del__(self):
            self._cleanup()

        def _split_archive(self):
            archivefile, internal = split_archive(self.fpath, self.ext)
            return archivefile, internal

        def _open(self):
            """
            This logic sets the "_handle" to the appropriate backend object
            such that zopen can behave like a standard IO object.

            In read-only mode:
                * If fpath is a normal file, _handle is the standard `open` object
                * If fpath is a seekable zipfile, _handle is an IOWrapper pointing
                    to the internal data
                * If fpath is a non-seekable zipfile, the data is extracted behind
                    the scenes and a standard `open` object to the extracted file
                    is given.

            In write mode:
                * NotImpelemented
            """
            if 'r' not in self.mode:
                raise NotImplementedError('Only read mode is supported for now')
            _handle = None
            fpath = os.fspath(self.fpath)
            if exists(fpath):
                _handle = open(fpath, self.mode)
            elif self.ext + '/' in fpath or self.ext + os.path.sep in fpath:
                archivefile, internal = self._split_archive()
                myzip = self.zfile
                if self._seekable:
                    import tempfile
                    # If we need data to be seekable, then we must extract it to a
                    # temporary file first.
                    self._temp_dpath = tempfile.mkdtemp(prefix='zopen_')
                    temp_fpath = join(self._temp_dpath, internal)
                    myzip.extract(internal, self._temp_dpath)
                    _handle = open(temp_fpath, self.mode)
                else:
                    # Try to load data directly from the zipfile
                    _handle = myzip.open(internal, 'r')
                    if self.mode == 'rb':
                        data = _handle.read()
                        _handle = io.BytesIO(data)
                    elif self.mode == 'r':
                        # FIXME: does not always work. handle seems to be closed
                        # too soon in the case util.zopen(module.__file__).read()
                        _handle = io.TextIOWrapper(_handle)
                    else:
                        raise KeyError(self.mode)
                    self._zfpath = archivefile
            if _handle is None:
                raise IOError('file {!r} does not exist'.format(fpath))
            self._handle = _handle

        def __enter__(self):
            return self

        def __exit__(self, ex_type, ex_value, ex_traceback):
            """
            Args:
                ex_type (Type[BaseException] | None):
                ex_value (BaseException | None):
                ex_traceback (TracebackType | None):

            Returns:
                bool | None
            """
            self.close()

        # TODO: Allow for navigating inside of the zipfile

        # TODO: opening a member should not force disk decompression unless we
        # really need to do real seeks. If we are just streaming the first few
        # bytes, then a standard handle will work fine.