import os
import zipfile

from parallels.core.utils.windows_utils import add_windows_long_path_prefix


def zip_directory(directory, archive):
    """Pack directory contents to zip file

    For example if we have directory '/root/test' passed as the first argument, which contains files:
    /root/test/a.txt
    /root/test/b/c.txt

    Then we will get zip file with contents:
    a.txt
    b/c.txt

    :param str directory: name of a directory to pack contents of
    :param str archive: name of ZIP file to pack contents to
    :rtype: None
    """
    with zipfile.ZipFile(archive, 'w') as zfp:
        for root, _, files in os.walk(directory):
            for filename in files:
                zfp.write(
                    add_windows_long_path_prefix(os.path.join(root, filename)),
                    os.path.relpath(os.path.join(root, filename), directory)
                )


def unzip(filename, directory):
    """Unpack ZIP file to a directory

    :type filename: str | unicode
    :type directory: str | unicode
    :rtype: None
    """
    with ZipFileLongPaths(filename) as zfp:
        zfp.extractall(directory)


class ZipFileLongPaths(zipfile.ZipFile):
    """The same as zipfile.ZipFile but patched to work with long paths on Windows

    See for more details:
    https://stackoverflow.com/questions/40419395/python-zipfile-extractall-ioerror-on-windows-when-extracting-files-from-long-pat
    """

    def _extract_member(self, member, targetpath, pwd):
        targetpath = add_windows_long_path_prefix(targetpath)
        return zipfile.ZipFile._extract_member(self, member, targetpath, pwd)
