#!/usr/bin/env python2.4
# revctl.py - bzr-based toolage
# Copyright Catalis, Inc.; written by Charles Duffy.

import logging
import os
import shutil
import stat
import string

import bzrlib, bzrlib.add, bzrlib.builtins, bzrlib.commit

logger = logging.getLogger('revctl')

class InvalidOperationException(Exception):
	"""The operation requested does not make sense"""
class BadStateException(Exception):
	"""The tree is in an unexpected state"""

class TransactionalFolder(object):
	"""Allow a filesystem folder to be treated transactionally"""
	def __init__(self, dirname):
		self.transaction_count = 0
		self.commit_queued = None # is a commit queued for after this transaction?
		self.dirname = dirname
		if not os.path.exists(os.path.join(self.dirname, '.bzr')):
			self.initialize()

	def _get_wt(self):
		"""Return a working tree object"""
		return bzrlib.workingtree.WorkingTree.open(self.dirname)

	def begin_transaction(self):
		logger.debug('Beginning a transaction')
		self.transaction_count += 1

	def end_transaction(self, *args, **kwargs):
		logger.debug('Ending a transaction')
		if self.transaction_count > 0:
			self.transaction_count -= 1
		else:
			logger.warning('Attempting to end a nonexistant transaction')
		if self.transaction_count == 0 and self.commit_queued:
			logger.debug('Running a queued commit')
			if len(args) == 0 and len(kwargs) == 0:
				args, kwargs = self.commit_queued
			self.commit(*args, **kwargs)
			self.commit_queued = None

	def initialize(self, initial_commit = True):
		"""Create a new working tree"""
		if os.path.exists(os.path.join(self.dirname, '.bzr')):
			raise InvalidOperationException('Attempting to reinitialize an already-revision-controlled folder')
		if not os.path.exists(self.dirname):
			os.mkdir(self.dirname)
		bzrlib.builtins.cmd_init().run(location=self.dirname)
		if initial_commit and self.has_changes():
			self.commit('Initial commit')

	def has_changes(self):
		"""Return true if the tree has been changed since creation; False otherwise; or throw an exception if it has changed in ways which will not be properly represented within revision control"""
		wt = self._get_wt()
		## See if there are unknown files in the tree
		try: wt.unknowns().next()
		except StopIteration: pass
		else: return True
		## Look for more conventional changes
		delta = wt.changes_from(wt.basis_tree(), want_unchanged=False)
		return delta.has_changed()

	def commit(self, message, allow_pointless=False, strict=True):
		"""Commit all changes. Includes adding new files to the repository."""
		if self.transaction_count == 0:
			wt = self._get_wt()
			bzrlib.add.smart_add_tree(wt, [ os.path.join(wt.basedir, '.') ], recurse=True, save=True)
			commit = bzrlib.commit.Commit()
			commit.commit(working_tree = wt, message=message, allow_pointless=allow_pointless, strict=strict)
		else:
			logger.debug('Deferring a commit due to an ongoing transaction')
			self.commit_queued = ([message], {'allow_pointless':allow_pointless, 'strict':strict})

	def revert(self, remove_unknowns = True):
		"""Completely revert the repository to its former state."""
		self.commit_queued = None
		wt = self._get_wt()
		wt.revert([], backups=False)
		if remove_unknowns:
			for unknown_file in wt.unknowns():
				fullpath = os.path.join(self.dirname, unknown_file)
				if stat.S_ISDIR(os.lstat(fullpath)[stat.ST_MODE]):
					shutil.rmtree(fullpath)
				else:
					os.remove(fullpath)

def build_message_for_call(func, self, args, kwargs):
	if self != None:
		classprefix = self.__class__.__name__ + '.'
	else:	classprefix = ''
	kwargs_items = [ (item_name, repr(item_val)) for (item_name, item_val) in kwargs.items() ]
	kwargs_items.sort()
	return "%s%s(%s)" % (classprefix, func.__name__, string.join([repr(arg) for arg in args] + ['%s=%s' % (key, value) for (key, value) in kwargs_items], ', '))

global instances
instances = {}
def get_instance_for(folder_name):
	global instances
	stat_data = os.stat(folder_name)
	folder_id = (stat_data[stat.ST_DEV], stat_data[stat.ST_INO]) ## FIXME: Need reasonable behavior on Win32
	if not instances.has_key(folder_id):
		instances[folder_id] = TransactionalFolder(folder_name)
	return instances[folder_id]

def atomic_operation(func):
	def wrapped_method(self, *args, **kwargs):
		if kwargs.has_key('transactional_folder_instance'):
			transactional_folder_instance = kwargs['transactional_folder_instance']
			del kwargs['transactional_folder_instance']
		else:
			transactional_folder_instance = None
		if transactional_folder_instance == None:
			if hasattr(self, 'transactional_folder_instance') and self.transactional_folder_instance != None:
				logger.debug('Using a preexisting TFI tied to this object')
				transactional_folder_instance = self.transactional_folder_instance
			elif hasattr(self, 'dirname'):
				logger.debug('Building or retrieving a TFI for this object')
				transactional_folder_instance = get_instance_for(self.dirname)
			else:
				logger.warn('Using current working directory for revision-controlled root')
				transactional_folder_instance = get_instance_for(os.getcwd())
		else:
			logger.debug('Initialized with transactional_folder_instance %s' % repr(transactional_folder_instance))
		assert transactional_folder_instance != None
		assert isinstance(transactional_folder_instance, TransactionalFolder)
		if transactional_folder_instance.transaction_count == 0 and transactional_folder_instance.has_changes():
			raise BadStateException('Tree has changes prior to operation even though transactions are not deferred')
		try:
			try:
				transactional_folder_instance.begin_transaction()
				retval = func(self, *args, **kwargs)
			except:
				transactional_folder_instance.revert()
				raise
			else:
				transactional_folder_instance.commit(build_message_for_call(func, self, args, kwargs))
		finally:
			transactional_folder_instance.end_transaction(build_message_for_call(func, self, args, kwargs))
		return retval
	return wrapped_method

### test suite!

import unittest
from tempfile import mkdtemp

class ArtificialException(Exception):
	"""An exception thrown to exercise test code"""

class RevisionedFolderTest(unittest.TestCase):
	def setUp(self):
		self.dirname = mkdtemp(prefix='test_py_revctl', suffix='.d')
	def tearDown(self):
		shutil.rmtree(self.dirname)
		del self.dirname

	def nonAtomicWriteFile(self, filename, contents, fail=False):
		f = open(os.path.join(self.dirname, str(filename)), 'w')
		f.write(str(contents))
		f.close()
		if fail: raise ArtificialException()

	def nonAtomicMkdir(self, dirname):
		return os.mkdir(os.path.join(self.dirname, dirname))

	@atomic_operation
	def writeFile(self, *args, **kwargs):
		self.nonAtomicWriteFile(*args, **kwargs)

	@atomic_operation
	def mkdir(self, dirname):
		return self.nonAtomicMkdir(dirname)

class CommitTest(RevisionedFolderTest):
	"""Validate that we can do individual commits"""
	def runTest(self):
		tfi = get_instance_for(self.dirname)
		self.writeFile(1,1)
		self.assert_(not tfi.has_changes())
		self.mkdir('subdir')
		self.assert_(not tfi.has_changes())
		self.writeFile('subdir/1', 2)
		self.assert_(not tfi.has_changes())

class TransactionSuccessTest(RevisionedFolderTest):
	"""Validate that we can store commits for batching"""
	def runTest(self):
		tfi = get_instance_for(self.dirname)
		tfi.begin_transaction()
		self.writeFile(1,1)
		self.assert_(tfi.has_changes())
		self.mkdir('subdir')
		self.assert_(tfi.has_changes())
		self.writeFile('subdir/1', 2)
		self.assert_(tfi.has_changes())
		tfi.end_transaction(message='Perform a queued batch')
		self.assert_(not tfi.has_changes())

class RevertTest(RevisionedFolderTest):
	"""Validate that manual reversion works correctly"""
	def runTest(self):
		tfi = get_instance_for(self.dirname)
		self.writeFile(1,1)
		self.assert_(not tfi.has_changes())
		self.nonAtomicWriteFile(1,2)
		self.nonAtomicWriteFile(2,2)
		self.nonAtomicMkdir('test.d')
		self.nonAtomicWriteFile('test.d/1', 2)
		self.assert_(tfi.has_changes())
		tfi.revert()
		self.assert_(not tfi.has_changes())

class TransactionFailTest(RevisionedFolderTest):
	"""Validate that failures will revert"""
	def runTest(self):
		tfi = get_instance_for(self.dirname)
		tfi.begin_transaction()
		self.writeFile(1,1)
		self.assert_(tfi.has_changes())
		self.assertRaises(ArtificialException,
			self.writeFile, 1, 2, fail=True)
		self.assert_(not tfi.has_changes())

class PriorInterferenceTest(RevisionedFolderTest):
	def runTest(self):
		tfi = get_instance_for(self.dirname)
		self.writeFile(1,1)
		self.nonAtomicWriteFile(1,2)
		self.assert_(tfi.has_changes())
		self.assertRaises(BadStateException,
			self.writeFile, 2, 1)
		tfi.revert()
		self.assert_(not tfi.has_changes())

class InitialCommitTest(RevisionedFolderTest):
	"""Initial tree state should be automatically committed on initialization, rather than causing a BadStateException"""
	def runTest(self):
		self.nonAtomicWriteFile(1,1)
		self.writeFile(1,2)
		## TODO: assert that there are two revisions in the repository's history

if __name__ == '__main__':
	logging.basicConfig(level=logging.WARNING)
	unittest.main()