# Copyright (C) 2006-2007, 2010-2011 Aaron Bentley <aaron@aaronbentley.com>
# Copyright (C) 2013 Aaron Bentley <aaron@aaronbentley.com>
# Copyright (C) 2007 Charlie Shepherd <masterdriverz@gentoo.org>
# Copyright (C) 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., 59 Temple Place, Suite 330, Boston, MA  02111-1307  USA

from shutil import rmtree

from bzrlib import (
    bzrdir,
    revision as _mod_revision,
    )
from bzrlib.branch import Branch
from bzrlib.errors import BzrCommandError, NoWorkingTree, NotBranchError
from bzrlib import registry
from bzrlib.workingtree import WorkingTree

from errors import (NotCheckout, UncommittedCheckout, ParentMissingRevisions,
                    NoParent)
from bzrlib.plugins.bzrtools.bzrtools import read_locked


class AllowChanged(object):

    @classmethod
    def check_changed(klass, wt, remove_branch):
        pass


class CheckChanged(object):

    @classmethod
    def check_changed(klass, wt, remove_branch):
        delta = wt.changes_from(wt.basis_tree(), want_unchanged=False)
        if delta.has_changed():
            klass.handle_changed(wt, remove_branch)


class HaltOnChange(CheckChanged):

    @staticmethod
    def handle_changed(wt, remove_branch):
        raise UncommittedCheckout()


class StoreChanges(CheckChanged):

    @staticmethod
    def handle_changed(wt, remove_branch):
        from bzrlib.plugins.pipeline.pipeline import PipeManager
        if remove_branch:
            raise BzrCommandError('Cannot store changes in deleted branch.')
        PipeManager.from_checkout(wt).store_uncommitted()


change_policy_registry = registry.Registry()


change_policy_registry.register('force', AllowChanged,
                                'Delete tree even if contents are modified.')


change_policy_registry.register('store', StoreChanges,
                                'Store changes in branch.  (Requires'
                                ' bzr-pipeline.)')


change_policy_registry.register('check', HaltOnChange,
                                'Stop if tree contents are modified.')


change_policy_registry.default_key = 'check'



def zap(path, remove_branch=False, policy=HaltOnChange):
    try:
        wt = bzrdir.BzrDir.open(path).open_workingtree(path,
                                                       recommend_upgrade=False)
    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)
    policy.check_changed(wt, remove_branch)
    if remove_branch:
        parent_loc = branch.get_parent()
        if parent_loc is None:
            raise NoParent()
        with read_locked(Branch.open(parent_loc)) as parent:
            last_revision = _mod_revision.ensure_null(parent.last_revision())
            if last_revision != _mod_revision.NULL_REVISION:
                graph = parent.repository.get_graph()
                heads = graph.heads([last_revision, branch.last_revision()])
                if branch.last_revision() in heads:
                    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.tests import TestCaseWithTransport


    class PipelinePluginFeature:

        @staticmethod
        def available():
            try:
                import bzrlib.plugins.pipeline
            except ImportError:
                return False
            else:
                return True


    class TestZap(TestCaseWithTransport):

        def make_checkout(self):
            wt = bzrdir.BzrDir.create_standalone_workingtree('source')
            return wt.branch.create_checkout('checkout', lightweight=True)

        def make_checkout2(self):
            wt = self.make_checkout()
            wt2 = wt.branch.bzrdir.sprout('source2').open_workingtree()
            return wt2.branch.create_checkout('checkout2', lightweight=True)

        def test_is_checkout(self):
            self.assertRaises(NotCheckout, zap, '.')
            wt = bzrdir.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')

        def make_modified_checkout(self):
            checkout = self.make_checkout()
            os.mkdir('checkout/foo')
            checkout.add('foo')
            return checkout

        def test_allow_modified(self):
            self.make_modified_checkout()
            self.assertRaises(UncommittedCheckout, zap, 'checkout')
            zap('checkout', policy=AllowChanged)

        def test_store(self):
            self.requireFeature(PipelinePluginFeature)
            checkout = self.make_modified_checkout()
            zap('checkout', policy=StoreChanges)
            self.assertIn('stored-transform',
                checkout.branch.bzrdir.transport.list_dir('branch'))

        def test_store_remove_branch(self):
            self.requireFeature(PipelinePluginFeature)
            checkout = self.make_modified_checkout()
            branch = self.make_branch('branch')
            checkout.branch.set_parent(branch.base)
            e = self.assertRaises(BzrCommandError, zap, 'checkout',
                                  policy=StoreChanges, remove_branch=True)
            self.assertEqual('Cannot store changes in deleted branch.', str(e))

        def test_store_remove_branch_unmodified(self):
            self.requireFeature(PipelinePluginFeature)
            checkout = self.make_checkout()
            branch = self.make_branch('branch')
            checkout.branch.set_parent(branch.base)
            zap('checkout', policy=StoreChanges, remove_branch=True)

        def test_zap_branch_with_unique_revision(self):
            parent = self.make_branch_and_tree('parent')
            parent.commit('foo')
            new_branch = self.make_branch('new')
            new_branch.set_parent(parent.branch.base)
            checkout = new_branch.create_checkout('checkout', lightweight=True)
            checkout.commit('unique')
            self.assertRaises(ParentMissingRevisions, zap, 'checkout',
                              remove_branch=True)

    return makeSuite(TestZap)
