"""Import upstream source into a branch"""

from bz2 import BZ2File
import errno
import os
from shutil import rmtree
from StringIO import StringIO
import tarfile
from unittest import makeSuite

from bzrlib.bzrdir import BzrDir
from bzrlib.errors import NoSuchFile, BzrCommandError, NotBranchError
from bzrlib.osutils import pathjoin, isdir, file_iterator
from bzrlib.tests import TestCaseInTempDir
from bzrlib.trace import warning
from bzrlib.transform import TreeTransform, resolve_conflicts, cook_conflicts
from bzrlib.workingtree import WorkingTree


def top_directory(path):
    """Return the top directory given in a path."""
    dirname = os.path.dirname(path)
    last_dirname = dirname
    while True:
        dirname = os.path.dirname(dirname)
        if dirname == '' or dirname == last_dirname:
            return last_dirname
        last_dirname = dirname


def common_directory(names):
    """Determine a single directory prefix from a list of names"""
    possible_prefix = None
    for name in names:
        name_top = top_directory(name)
        if possible_prefix is None:
            possible_prefix = name_top
        else:
            if name_top != possible_prefix:
                return None
    return possible_prefix


def do_directory(tt, trans_id, tree, relative_path, path):
    if isdir(path) and tree.path2id(relative_path) is not None:
        tt.cancel_deletion(trans_id)
    else:
        tt.create_directory(trans_id)


def add_implied_parents(implied_parents, path):
    """Update the set of implied parents from a path"""
    parent = os.path.dirname(path)
    if parent in implied_parents:
        return
    implied_parents.add(parent)
    add_implied_parents(implied_parents, parent)


def names_of_files(tar_file):
    for member in tar_file.getmembers():
        if member.type != "g":
            yield member.name


def import_tar(tree, tar_input):
    """Replace the contents of a working directory with tarfile contents.
    The tarfile may be a gzipped stream.  File ids will be updated.
    """
    tar_file = tarfile.open('lala', 'r', tar_input)
    prefix = common_directory(names_of_files(tar_file))
    tt = TreeTransform(tree)

    removed = set()
    for path, entry in tree.inventory.iter_entries():
        if entry.parent_id is None:
            continue
        trans_id = tt.trans_id_tree_path(path)
        tt.delete_contents(trans_id)
        removed.add(path)

    added = set() 
    implied_parents = set()
    seen = set()
    for member in tar_file.getmembers():
        if member.type == 'g':
            # type 'g' is a header
            continue
        relative_path = member.name 
        if prefix is not None:
            relative_path = relative_path[len(prefix)+1:]
        if relative_path == '':
            continue
        add_implied_parents(implied_parents, relative_path)
        trans_id = tt.trans_id_tree_path(relative_path)
        added.add(relative_path.rstrip('/'))
        path = tree.abspath(relative_path)
        if member.name in seen:
            tt.cancel_creation(trans_id)
        seen.add(member.name)
        if member.isreg():
            tt.create_file(file_iterator(tar_file.extractfile(member)), 
                           trans_id)
        elif member.isdir():
            do_directory(tt, trans_id, tree, relative_path, path)
        elif member.issym():
            tt.create_symlink(member.linkname, trans_id)

    for relative_path in implied_parents.difference(added):
        if relative_path == "":
            continue
        trans_id = tt.trans_id_tree_path(relative_path)
        path = tree.abspath(relative_path)
        do_directory(tt, trans_id, tree, relative_path, path)
        added.add(relative_path)

    for conflict in cook_conflicts(resolve_conflicts(tt), tt):
        warning(conflict)
    tt.apply()
    update_ids(tree, added, removed)


def update_ids(tree, added, removed):
    """Make sure that all present files files have file_ids.
    """
    # XXX detect renames
    new = added.difference(removed)
    deleted = removed.difference(added)
    tree.add(sorted(new))
    tree.remove(sorted(deleted, reverse=True))


def do_import(source, tree_directory=None):
    """Implementation of import command.  Intended for UI only"""
    if tree_directory is not None:
        try:
            tree = WorkingTree.open(tree_directory)
        except NotBranchError:
            if not os.path.exists(tree_directory):
                os.mkdir(tree_directory)
            branch = BzrDir.create_branch_convenience(tree_directory)
            tree = branch.bzrdir.open_workingtree()
    else:
        tree = WorkingTree.open_containing('.')[0]
    tree.lock_write()
    try:
        if tree.changes_from(tree.basis_tree()).has_changed():
            raise BzrCommandError("Working tree has uncommitted changes.")

        if (source.endswith('.tar') or source.endswith('.tar.gz') or 
            source.endswith('.tar.bz2')) or source.endswith('.tgz'):
            try:
                if source.endswith('.bz2'):
                    tar_input = BZ2File(source, 'r')
                    tar_input = StringIO(tar_input.read())
                else:
                    tar_input = file(source, 'rb')
            except IOError, e:
                if e.errno == errno.ENOENT:
                    raise NoSuchFile(source)
            try:
                import_tar(tree, tar_input)
            finally:
                tar_input.close()
    finally:
        tree.unlock()

class TestImport(TestCaseInTempDir):

    def make_tar(self, mode='w'):
        result = StringIO()
        tar_file = tarfile.open('project-0.1.tar', mode, result)
        os.mkdir('project-0.1')
        tar_file.add('project-0.1')
        os.mkdir('project-0.1/junk')
        tar_file.add('project-0.1/junk')
        
        f = file('project-0.1/README', 'wb')
        f.write('What?')
        f.close()
        tar_file.add('project-0.1/README')

        f = file('project-0.1/FEEDME', 'wb')
        f.write('Hungry!!')
        f.close()
        tar_file.add('project-0.1/FEEDME')

        tar_file.close()
        rmtree('project-0.1')
        result.seek(0)
        return result

    def make_tar2(self):
        result = StringIO()
        tar_file = tarfile.open('project-0.2.tar', 'w', result)
        os.mkdir('project-0.2')
        tar_file.add('project-0.2')
        
        os.mkdir('project-0.2/junk')
        tar_file.add('project-0.2/junk')

        f = file('project-0.2/README', 'wb')
        f.write('Now?')
        f.close()
        tar_file.add('project-0.2/README')
        tar_file.close()

        tar_file = tarfile.open('project-0.2.tar', 'a', result)
        tar_file.add('project-0.2/README')

        rmtree('project-0.2')
        return result

    def make_messed_tar(self):
        result = StringIO()
        tar_file = tarfile.open('project-0.1.tar', 'w', result)
        os.mkdir('project-0.1')
        tar_file.add('project-0.1')

        os.mkdir('project-0.2')
        tar_file.add('project-0.2')
        
        f = file('project-0.1/README', 'wb')
        f.write('What?')
        f.close()
        tar_file.add('project-0.1/README')
        tar_file.close()
        rmtree('project-0.1')
        result.seek(0)
        return result

    def test_top_directory(self):
        self.assertEqual(top_directory('ab/b/c'), 'ab')
        self.assertEqual(top_directory('/etc'), '/')

    def test_common_directory(self):
        self.assertEqual(common_directory(['ab/c/d', 'ab/c/e']), 'ab')
        self.assertIs(common_directory(['ab/c/d', 'ac/c/e']), None)

    def test_untar(self):
        tar_file = self.make_tar()
        tree = BzrDir.create_standalone_workingtree('tree')
        import_tar(tree, tar_file)
        self.assertTrue(tree.path2id('README') is not None) 
        self.assertTrue(tree.path2id('FEEDME') is not None)
        self.assertTrue(os.path.isfile(tree.abspath('README')))
        self.assertEqual(tree.inventory[tree.path2id('README')].kind, 'file')
        self.assertEqual(tree.inventory[tree.path2id('FEEDME')].kind, 'file')
        
        f = file(tree.abspath('junk/food'), 'wb')
        f.write('I like food\n')
        f.close()

        tar_file = self.make_tar2()
        import_tar(tree, tar_file)
        self.assertTrue(tree.path2id('README') is not None) 
        self.assertTrue(not os.path.exists(tree.abspath('FEEDME')))


    def test_untar2(self):
        tar_file = self.make_messed_tar()
        tree = BzrDir.create_standalone_workingtree('tree')
        import_tar(tree, tar_file)
        self.assertTrue(tree.path2id('project-0.1/README') is not None) 

    def test_untar_gzip(self):
        tar_file = self.make_tar(mode='w:gz')
        tree = BzrDir.create_standalone_workingtree('tree')
        import_tar(tree, tar_file)
        self.assertTrue(tree.path2id('README') is not None) 


def test_suite():
    return makeSuite(TestImport)
