'Convenience decorators for reading and writing to compressed archives'
import os
import shutil
import zipfile
import tarfile
import tempfile
from collections import defaultdict
from decorator import decorator
from six import string_types


class Archive(object):

    def __init__(self, path, extension=None):
        """
        The path can be a string or file-like object.
        The extension determines the archive format.
        """
        # If the extension is specified, use it
        if extension:
            try:
                make_consumer, make_generator = dict(
                    EXTENSION_PACKS)[extension.lower()]
            except KeyError:
                raise ArchiveError(
                    'Could not recognize file extension: %s' % extension)
        # If path is a file-like object, raise Exception
        elif hasattr(path, 'read'):
            raise ArchiveError(
                'Must specify file extension when using a file-like object')
        # If path is a string, try to recognize the extension
        else:
            pathLower = path.lower()
            for extension, (make_consumer, make_generator) in EXTENSION_PACKS:
                if pathLower.endswith(extension):
                    break
            else:
                raise ArchiveError(
                    'Could not recognize archive format from file extension')
        # Set
        self.__path = path
        self.__extension = extension
        self.__make_consumer = make_consumer
        self.__make_generator = make_generator

    def save(self, filePaths, basePath=''):
        """
        Compress filePaths using the extension specified in archivePath,
        truncating each filePath into a relativePath using basePath.
        """
        # Convert filePaths into a list if it isn't one already
        if isinstance(filePaths, string_types):
            filePaths = [filePaths]
        if basePath:
            baseIndex = len(os.path.abspath(basePath)) + 1
            truncate_basePath = lambda x: os.path.abspath(x)[baseIndex:]
        else:
            truncate_basePath = lambda x: x
        consumer = self.__make_consumer(self.__path)
        next(consumer)
        for filePath in expand_paths(filePaths):
            relativePath = truncate_basePath(filePath)
            consumer.send((filePath, relativePath))
        consumer.close()
        # If path is a file-like object, prepare it
        if hasattr(self.__path, 'read'):
            self.__path.seek(0)
        return self.__path

    def load(self, targetFolder):
        'Uncompress archivePath to a targetFolder.'
        # If path is a file-like object, prepare it
        if hasattr(self.__path, 'read'):
            self.__path.seek(0)
        return self.__make_generator(self.__path, targetFolder)

    def get_extension(self):
        return self.__extension


class ArchiveError(Exception):
    'Custom exception for archiveIO'
    pass


class TemporaryFolder(object):
    'Create a temporary folder on entry and remove it on exit'

    def __init__(self, suffix='', prefix='tmp', dir=None):
        self.suffix = suffix
        self.prefix = prefix
        if dir is None:
            dir = os.path.expanduser('~/.tmp')
            try:
                os.makedirs(dir)
            except OSError:
                pass
        self.dir = dir

    def __enter__(self):
        self.temporaryFolder = tempfile.mkdtemp(
            self.suffix, self.prefix, self.dir)
        return self.temporaryFolder

    def __exit__(self, type, value, traceback):
        shutil.rmtree(self.temporaryFolder)


@decorator
def save(function, *args, **kw):
    """
    Decorator to support saving to compressed files for functions
    whose first argument is the targetPath

    If the first argument ends with a recognized extension,
    the decorator runs the function in a temporary folder and
    compresses the resulting output to targetPath.

    Archive format is determined by file extension:
    .zip .tar.gz .tar.bz2 .tar
    """
    targetPath = kw.get('targetPath', args[0])
    try:
        archive = Archive(targetPath, extension=kw.get('targetExtension'))
    # If we did not recognize the extension, run function as usual
    except ArchiveError:
        return function(*args, **kw)
    targetName = kw.get('targetName')
    if not targetName:
        # If path is a file-like object, raise exception
        if hasattr(targetPath, 'read'):
            raise ArchiveError(
                'Must specify targetName when using a file-like object')
        # If path is a string, remove matching extension from filename
        else:
            targetName = os.path.basename(
                targetPath[:targetPath.lower().rfind(archive.get_extension())])
    # Make temporaryFolder
    with TemporaryFolder() as temporaryFolder:
        # Run function in temporaryFolder
        function(os.path.join(temporaryFolder, targetName), *args[1:], **kw)
        # Save
        return archive.save(temporaryFolder, temporaryFolder)


def load(function=None, extensions=None, CustomException=IOError):
    """
    Decorator to support loading from compressed files for functions
    whose first argument is the sourcePath

    If the first argument ends with a recognized extension,
    the decorator uncompresses sourcePath to a temporary folder
    and runs the function on each resulting file until it succeeds.

    Archive format is determined by file extension:
    .zip .tar.gz .tar.bz2 .tar

    # Try load(sourcePath) on each file in the archive until one works
    @archiveIO.load
    def load(sourcePath):
        return open(sourcePath).read()
    # Try load(sourcePath) on TXT files and if all fail, try CSV files
    @archiveIO.load(extensions=['.txt', '.csv'])
    def load(sourcePath):
        return open(sourcePath).read()
    # Raise AppError instead of IOError
    @archiveIO.load(CustomException=AppError)
    def load(sourcePath):
        return open(sourcePath).read()
    """
    def load(function, *args, **kw):
        sourcePath = kw.get('sourcePath', args[0])
        try:
            archive = Archive(sourcePath)
        # If we did not recognize the extension, run function as usual
        except ArchiveError:
            return function(*args, **kw)
        # Make temporaryFolder
        with TemporaryFolder() as temporaryFolder:
            try:
                # Convert a disposable generator into a reusable list
                paths = list(archive.load(temporaryFolder))
                # For each path, run function and exit if successful
                for path in select_extensions(paths, extensions or []):
                    try:
                        return function(path, *args[1:], **kw)
                    except:
                        pass
            except:
                raise CustomException('Could not open archive')
            raise CustomException('Could not load archive')
    if function:
        return decorator(load, function)
    else:
        return decorator(load)


def select_extensions(paths, extensions):
    'Order paths by matching file extensions or return everything'
    pathsByExtension = defaultdict(list)
    for path in paths:
        extension = os.path.splitext(path)[1].lower()
        pathsByExtension[extension].append(path)
    selectedPaths = []
    for extension in extensions:
        selectedPaths.extend(sorted(pathsByExtension[extension.lower()]))
    return selectedPaths or paths


def expand_paths(paths):
    'Expand folderPaths'
    filePaths = []
    for path in paths:
        if os.path.isdir(path):
            filePaths.extend(walk_paths(path))
        elif os.path.exists(path):
            filePaths.append(path)
    return set(filePaths)


def walk_paths(path):
    'Yield filePaths one by one from the specified path'
    path = os.path.abspath(path)
    for rootPath, folderNames, fileNames in os.walk(path):
        for fileName in fileNames:
            yield os.path.join(rootPath, fileName)
        if not folderNames and not fileNames and os.path.abspath(
                rootPath) != path:
            yield rootPath


def open_tarfile(targetPath, mode):
    valueByKey = dict(mode=mode)
    if hasattr(targetPath, 'read'):
        valueByKey['fileobj'] = targetPath
    else:
        valueByKey['name'] = targetPath
    return tarfile.open(**valueByKey)


def make_consumer_tar(targetPath, mode='w'):
    'Save .tar file'
    def filter_(tarInfo):
        'Anonymize file data'
        tarInfo.uid = tarInfo.gid = 0
        tarInfo.uname = tarInfo.gname = 'root'
        return tarInfo
    with open_tarfile(targetPath, mode) as targetFile:
        while True:
            filePath, relativePath = yield
            targetFile.add(
                filePath, relativePath, recursive=False, filter=filter_)


def make_consumer_tar_gz(targetPath):
    'Save .tar.gz file'
    return make_consumer_tar(targetPath, mode='w:gz')


def make_consumer_tar_bz2(targetPath):
    'Save .tar.bz2 file'
    return make_consumer_tar(targetPath, mode='w:bz2')


def make_consumer_zip(targetPath):
    'Save .zip file'
    with zipfile.ZipFile(targetPath, 'w', zipfile.ZIP_DEFLATED) as targetFile:
        while True:
            filePath, relativePath = yield
            targetFile.write(filePath, relativePath)


def make_generator_tar(sourcePath, temporaryFolder):
    'Load .tar.gz or .tar.bz2 or .tar file'
    with open_tarfile(sourcePath, 'r') as sourceFile:
        sourceFile.extractall(temporaryFolder)
        for relativePath in sourceFile.getnames():
            yield os.path.join(temporaryFolder, relativePath)


def make_generator_zip(sourcePath, temporaryFolder):
    'Load .zip file'
    with zipfile.ZipFile(sourcePath) as sourceFile:
        sourceFile.extractall(temporaryFolder)
        for relativePath in sourceFile.namelist():
            yield os.path.join(temporaryFolder, relativePath)


EXTENSION_PACKS = [
    ('.zip', (make_consumer_zip, make_generator_zip)),
    ('.tar.gz', (make_consumer_tar_gz, make_generator_tar)),
    ('.tar.bz2', (make_consumer_tar_bz2, make_generator_tar)),
    ('.tar', (make_consumer_tar, make_generator_tar)),
]
