# 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."""

import os
import sys
import tarfile
from contextlib import closing
from io import BytesIO

from .. import errors, osutils
from ..export import _export_iter_entries


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
    """
    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"
                              % (final_path, entry.kind))
    return (item, fileobj)


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

    Args:
      tree: Tree to export
      subdir: Sub directory to export
      force_mtime: Option mtime to force, instead of using tree
        timestamps.
    Returns: A generator that will produce file content chunks.
    """
    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, recurse_nested=recurse_nested):
            (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, recurse_nested=False):
    """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(''):
            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,
                recurse_nested=recurse_nested):
            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, recurse_nested=False):
    """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',
        recurse_nested=recurse_nested)


def plain_tar_generator(tree, dest, root, subdir,
                        force_mtime=None, recurse_nested=False):
    """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='',
        recurse_nested=recurse_nested)


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


def tar_lzma_generator(tree, dest, root, subdir, force_mtime=None,
                       compression_format="alone", recurse_nested=False):
    """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 ModuleNotFoundError as exc:
        raise errors.DependencyNotPresent('lzma', e) from exc

    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,
            recurse_nested=recurse_nested):
        yield compressor.compress(chunk)

    yield compressor.flush()
