from shutil import rmtree

from bzrlib.branch import Branch
from bzrlib.errors import NoWorkingTree, NotLocalUrl, NotBranchError
from bzrlib.workingtree import WorkingTree

from errors import (NotCheckout, UncommittedCheckout, ParentMissingRevisions, 
                    NoParent)


def zap(path, remove_branch=False):
    try:
        wt = WorkingTree.open(path)
    except (NoWorkingTree, NotBranchError):
        raise NotCheckout(path)
    tree_base = wt.bzrdir.transport.base
    branch = wt.branch
    branch_base = branch.bzrdir.transport.base
    if tree_base == branch_base:
        raise NotCheckout(path)
    delta = wt.changes_from(wt.basis_tree(), want_unchanged=False)
    if delta.has_changed():
        raise UncommittedCheckout()
    if remove_branch:
        parent_loc = branch.get_parent()
        if parent_loc is None:
            raise NoParent()
        parent = Branch.open(parent_loc)
        p_ancestry = parent.repository.get_ancestry(parent.last_revision())
        if branch.last_revision() not in p_ancestry:
            raise ParentMissingRevisions(branch.get_parent())
    rmtree(path)
    if remove_branch:
        t = branch.bzrdir.transport
        while t.base != branch_base:
            t = t.clone('..')
        t = t.clone('..')
        t.delete_tree('')


def test_suite():
    import os
    from unittest import makeSuite
    
    from bzrlib.bzrdir import BzrDir, BzrDirMetaFormat1
    from bzrlib.branch import BranchReferenceFormat
    from bzrlib.tests import TestCaseInTempDir

    class TestZap(TestCaseInTempDir):

        def make_checkout(self):
            wt = BzrDir.create_standalone_workingtree('source')
            os.mkdir('checkout')
            checkout = BzrDirMetaFormat1().initialize('checkout')
            BranchReferenceFormat().initialize(checkout, wt.branch)
            return checkout.create_workingtree()

        def make_checkout2(self):
            wt = self.make_checkout()
            wt2 = wt.branch.bzrdir.sprout('source2').open_workingtree()
            os.mkdir('checkout2')
            checkout = BzrDirMetaFormat1().initialize('checkout2')
            BranchReferenceFormat().initialize(checkout, wt2.branch)
            return checkout.create_workingtree()

        def test_is_checkout(self):
            self.assertRaises(NotCheckout, zap, '.')
            wt = BzrDir.create_standalone_workingtree('.')
            self.assertRaises(NotCheckout, zap, '.')

        def test_zap_works(self):
            self.make_checkout()
            self.assertIs(True, os.path.exists('checkout'))
            zap('checkout')
            self.assertIs(False, os.path.exists('checkout'))
            self.assertIs(True, os.path.exists('source'))

	def test_zap_branch(self):
            self.make_checkout2()
            base = WorkingTree.open('checkout').branch.base
            self.assertIs(True, os.path.exists('checkout'))
            self.assertRaises(NoParent, zap, 'checkout', remove_branch=True)
            self.assertIs(True, os.path.exists('checkout'))
            self.assertIs(True, os.path.exists('source'))
            zap('checkout2', remove_branch=True)
            self.assertIs(False, os.path.exists('checkout2'))
            self.assertIs(False, os.path.exists('source2'))

        def test_checks_modified(self):
            checkout = self.make_checkout()
            os.mkdir('checkout/foo')
            checkout.add('foo')
            self.assertRaises(UncommittedCheckout, zap, 'checkout')
            checkout.commit('commit changes to branch')
            zap('checkout')

    return makeSuite(TestZap)
