# Copyright (C) 2005, 2006, 2008-2011 Canonical Ltd
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation; either version 2 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program; if not, write to the Free Software
# Foundation, Inc., 51 Franklin Street, Fifth Floor, Boston, MA 02110-1301 USA

"""Export a tree to a tarball."""

from __future__ import absolute_import

from contextlib import closing
import os
import sys
import tarfile

from .. import (
    errors,
    osutils,
    )
from ..export import _export_iter_entries
from ..sixish import (
    BytesIO,
    )


def prepare_tarball_item(tree, root, final_path, tree_path, entry, force_mtime=None):
    """Prepare a tarball item for exporting

    :param tree: Tree to export
    :param final_path: Final path to place item
    :param tree_path: Path for the entry in the tree
    :param entry: Entry to export
    :param force_mtime: Option mtime to force, instead of using tree
        timestamps.

    Returns a (tarinfo, fileobj) tuple
    """
    file_id = getattr(entry, 'file_id', None)
    filename = osutils.pathjoin(root, final_path)
    item = tarfile.TarInfo(filename)
    if force_mtime is not None:
        item.mtime = force_mtime
    else:
        item.mtime = tree.get_file_mtime(tree_path)
    if entry.kind == "file":
        item.type = tarfile.REGTYPE
        if tree.is_executable(tree_path):
            item.mode = 0o755
        else:
            item.mode = 0o644
        # This brings the whole file into memory, but that's almost needed for
        # the tarfile contract, which wants the size of the file up front.  We
        # want to make sure it doesn't change, and we need to read it in one
        # go for content filtering.
        content = tree.get_file_text(tree_path)
        item.size = len(content)
        fileobj = BytesIO(content)
    elif entry.kind in ("directory", "tree-reference"):
        item.type = tarfile.DIRTYPE
        item.name += '/'
        item.size = 0
        item.mode = 0o755
        fileobj = None
    elif entry.kind == "symlink":
        item.type = tarfile.SYMTYPE
        item.size = 0
        item.mode = 0o755
        item.linkname = tree.get_symlink_target(tree_path)
        fileobj = None
    else:
        raise errors.BzrError("don't know how to export {%s} of kind %r"
                              % (file_id, entry.kind))
    return (item, fileobj)


def tarball_generator(tree, root, subdir=None, force_mtime=None, format=''):
    """Export tree contents to a tarball.

    :returns: A generator that will produce file content chunks.

    :param tree: Tree to export

    :param subdir: Sub directory to export

    :param force_mtime: Option mtime to force, instead of using tree
        timestamps.
    """
    buf = BytesIO()
    with closing(tarfile.open(None, "w:%s" % format, buf)) as ball, tree.lock_read():
        for final_path, tree_path, entry in _export_iter_entries(tree, subdir):
            (item, fileobj) = prepare_tarball_item(
                tree, root, final_path, tree_path, entry, force_mtime)
            ball.addfile(item, fileobj)
            # Yield the data that was written so far, rinse, repeat.
            yield buf.getvalue()
            buf.truncate(0)
            buf.seek(0)
    yield buf.getvalue()


def tgz_generator(tree, dest, root, subdir, force_mtime=None):
    """Export this tree to a new tar file.

    `dest` will be created holding the contents of this tree; if it
    already exists, it will be clobbered, like with "tar -c".
    """
    with tree.lock_read():
        import gzip
        if force_mtime is not None:
            root_mtime = force_mtime
        elif (getattr(tree, "repository", None) and
              getattr(tree, "get_revision_id", None)):
            # If this is a revision tree, use the revisions' timestamp
            rev = tree.repository.get_revision(tree.get_revision_id())
            root_mtime = rev.timestamp
        elif tree.is_versioned(u''):
            root_mtime = tree.get_file_mtime('')
        else:
            root_mtime = None

        is_stdout = False
        basename = None
        # gzip file is used with an explicit fileobj so that
        # the basename can be stored in the gzip file rather than
        # dest. (bug 102234)
        basename = os.path.basename(dest)
        buf = BytesIO()
        zipstream = gzip.GzipFile(basename, 'w', fileobj=buf,
                                  mtime=root_mtime)
        for chunk in tarball_generator(tree, root, subdir, force_mtime):
            zipstream.write(chunk)
            # Yield the data that was written so far, rinse, repeat.
            yield buf.getvalue()
            buf.truncate(0)
            buf.seek(0)
        # Closing zipstream may trigger writes to stream
        zipstream.close()
        yield buf.getvalue()


def tbz_generator(tree, dest, root, subdir, force_mtime=None):
    """Export this tree to a new tar file.

    `dest` will be created holding the contents of this tree; if it
    already exists, it will be clobbered, like with "tar -c".
    """
    return tarball_generator(
        tree, root, subdir, force_mtime, format='bz2')


def plain_tar_generator(tree, dest, root, subdir,
                        force_mtime=None):
    """Export this tree to a new tar file.

    `dest` will be created holding the contents of this tree; if it
    already exists, it will be clobbered, like with "tar -c".
    """
    return tarball_generator(
        tree, root, subdir, force_mtime, format='')


def tar_xz_generator(tree, dest, root, subdir, force_mtime=None):
    return tar_lzma_generator(tree, dest, root, subdir, force_mtime, "xz")


def tar_lzma_generator(tree, dest, root, subdir, force_mtime=None,
                       compression_format="alone"):
    """Export this tree to a new .tar.lzma file.

    `dest` will be created holding the contents of this tree; if it
    already exists, it will be clobbered, like with "tar -c".
    """
    try:
        import lzma
    except ImportError as e:
        raise errors.DependencyNotPresent('lzma', e)

    if sys.version_info[0] == 2:
        compressor = lzma.LZMACompressor(
            options={"format": compression_format})
    else:
        compressor = lzma.LZMACompressor(
            format={
                'xz': lzma.FORMAT_XZ,
                'raw': lzma.FORMAT_RAW,
                'alone': lzma.FORMAT_ALONE,
                }[compression_format])

    for chunk in tarball_generator(
            tree, root, subdir, force_mtime=force_mtime):
        yield compressor.compress(chunk)

    yield compressor.flush()