
import os
import sys
import stat
import types
import atexit
import time
# import atexit
import random
import hashlib
import cPickle
import threading
from gyropy import g_mailbox

import die
import gpkmisc

DEBUG = 0

UnpicklingError = cPickle.UnpicklingError

class CannotGetStat(OSError):
	def __init__(self, *s):
		OSError.__init__(self, *s)

def cachepath(f, tail='', root=''):
	"""OBSOLETE:
	Return a pathname suitable for cacheing some result.
	@return: C{(path_to_root,path_with_tail)}.   C{Path_to_root} is/will be a directory;
		C{path_with_tail} is a path to a data file within that directory.
		Normally, the actual cache is at the location
		C{os.path.join(path_to_root,path_with_tail)} on the disk;
		that is what you would pass as the C{fname} argument to L{load_cache} or
		L{dump_cache}.
	@rtype: C{tuple(str, str)}
	@param f: An arbitrary key, could be a pathname or a tuple of information about a file.
	@type f: often L{str}, but could be anything convertible to a L{str} via L{repr}.
	@param tail: something to add at the end of the constructed path.
	@type tail: str or None
	@raise ValueError: if C{suffix_to_del} is specified and C{f} doesn't end that way.
	"""
	if len(f)==0:
		raise ValueError, "Need to specify cache info"
	hf = hashlib.md5(repr(f)).hexdigest()
	d1 = hf[:2]
	d2 = hf[2:4]
	d3 = hf[4:6]
	fc = hf[6:]
	r = os.path.join(root, d1, d2, d3)
	# p = os.path.join(r, fc) + tail
	return (r, "%s%s" % (fc, tail))


def fileinfo(fname, *other):
	"""Collect enough information about a file to determine whether or
	not the cache can be used.
	"""
	if fname is None:
		return None
	try:
		s = os.stat(fname)
	except (IOError, OSError), ex:
		raise CannotGetStat("path=%s" % fname, *(ex.args))
	return (os.path.normpath(fname),
				# This is OK, as long as access patterns are stable, but it's clearly
				# possible to make the cacheing inefficient by accessing the same file
				# in different ways.   Using os.path.abspath() might be better.
			s[stat.ST_INO],	# This seems to be the same on different machines
					# that NFS mount the same file.
			# s[stat.ST_DEV],	# Not a good idea: this differs from machine to machine
						# when accessing the same file.
			s[stat.ST_SIZE], s[stat.ST_MTIME]) + other


def modFileInfo(fname, *other):
	"""Collect enough information about a file to determine whether or
	not the cache can be used.
	"""
	if fname is None:
		return None
	try:
		s = os.stat(fname)
	except (IOError, OSError), ex:
		raise CannotGetStat("path=%s" % fname, *(ex.args))
	# This is tuned to match .pyc files that may be identical, even if they
	# were installed at different times.   It's a bit weak, but the alternative
	# (i.e. including inode # or mtime) fails if you are trying to share computation
	# across a cluster.
	return (os.path.abspath(fname), s[stat.ST_SIZE]) + other


class BadFileFormat(Exception):
	def __init__(self, *s):
		Exception.__init__(self, *s)



def modinfo(m, seen):
	if not isinstance(m, tuple):
		m = (m,)
	for mm in m:
		try:
			modinfo_guts(mm, seen)
		except (IOError, OSError), ex:
			raise CannotGetStat('path: %s' % getattr(m, '__file__', '???'), *(ex.args))
	return tuple(sorted(seen.values()))


_modinfocache = {}	#: We cache module info because they are normally not reloaded during the run of a program.
def modinfo_guts(m, seen):
	assert isinstance(m, types.ModuleType)
	if m.__name__ in seen:
		return
	try:
		seen[m.__name__]  = _modinfocache[m.__name__]
	except KeyError:
		try:
			tmp = modFileInfo(m.__file__)
		except AttributeError:
			try:
				tmp = (m.__name__, len(m.__dict__))
			except AttributeError:
				tmp = m.__name__
		seen[m.__name__] = tmp
		_modinfocache[m.__name__] = tmp

	for o in m.__dict__.values():
		if isinstance(o, types.ModuleType):
			modinfo_guts(o, seen)


def namedModInfo(nm):
	if not isinstance(nm, tuple):
		nm = (nm,)
	s = {}
	return tuple( [ modinfo(sys.modules[q], s) for q in nm ] )


def _bg_dumper(mb):
	while True:
		try:
			ci, e = mb.get()
		except g_mailbox.EOF:
			break
		ci.dump(e)

def _bg_cleaner(mb):
	while True:
		try:
			ci, avoid = mb.get()
		except g_mailbox.EOF:
			break
		ci.makespace(avoid)


def _shutdown():
	cache_info._dump.putclose()
	cache_info._clean.putclose()
	cache_info._dumpthread.join()
	cache_info._cleanthread.join()

atexit.register(_shutdown)

# def _stop():
	# DumpBox.putclose()
	# CleanBox.putclose()

# atexit.register(_stop)


def _dodel(path, avoid, when, n):
	try:
		dx = os.listdir(path)
	except OSError:
		return 0
	ldx = len(dx)
	n *= ldx
	rv = 0
	empty = True
	for d in random.sample(dx, ldx):
		p = os.path.join(path, d)
		if p == avoid:
			empty = False
			continue
		try:
			s = os.stat(p)
		except OSError:
			continue
		if stat.S_ISDIR(s.st_mode):
			# Whoops.  There's a directory here.
			# Give up at this level and work downwards instead.
			return rv + _dodel(p, avoid, when, n)
		if s.st_mtime > when or n < 1:
			empty = False
			continue
		try:
			os.remove(p)
			rv += 1
		except OSError:
			pass
	# The directory may be empty now, so we'll delete it too.
	if empty:
		try:
			os.rmdir(path)
		except OSError:
			pass
	return rv


class cache_info(object):
	"""This class manages a disk cache of arbitrary objects.
	It first constructs a unique name, based on information that you give,
	then you can C{dump} data to that path, or C{load} data from that path.
	An attempt to C{load} data will either succeed or raise an exception;
	an attempt to C{dump} will either succeed or silently fail.

	Typical use::
	
		def cached_f(parameters):
			ci = cache_info(info=tuple(parameters))
			if ci is not None:
				try:
					return ci.load()
				except (BadFileFormat, IOError, OSError):
					pass
			o = f(parameters)
			if ci is not None:
				ci.dump(o)
			return o

	@note: This class assumes that the results it is cacheing are generated by a
		function C{f(parameters)}.  You need to be careful to give all of the
		relevant parameters to C{cache_info}, otherwise you can get the wrong
		results back.  For instance, if you have five parameters and you forget
		to give C{parameters[2]} to C{cache_info}, it will happily store
		values obtained with all different values of C{parameters[2]} in the same
		slot, and when you later call C{load}, you'll get whatever you asked for,
		even if it is not what you wanted.
	"""
	Age = 864000
	NumObj = 10000

	_dump = g_mailbox.mailbox()
	_dumpthread = threading.Thread(target=_bg_dumper, args=(_dump,), name='cache_dumper')
	_dumpthread.daemon = True
	_dumpthread.start()
	_clean = g_mailbox.mailbox()
	_cleanthread = threading.Thread(target=_bg_cleaner, args=(_clean,), name='cache_cleaner')
	_cleanthread.daemon = True
	_cleanthread.start()

	def __init__(self, root, info=(), fname=None, modname=None, mod=None):
		"""
		@param info: This is where you specify the parameters from which the cached value
			can be computed.  It is essentially a look-up key for the value.
		@type info: tuple(anything)
		@param fname: You can specify that the cached value depends on the contents of
			a file (in addition to other parameters).   See L{fileinfo} for details.
		@type fname: str
		@type modname: str or tuple(str)
		@param modname: You can specify that the value depends on a module (or a list of modules).
			(You give the names of the modules here.)  In which case, it tries
			to detect changes to the specified modules.   See L{namedModInfo} for details.
			You use this argument to protect youself against changes to the code
			used to compute the cached value.  Obviously, you don't want to L{load}
			a value from last weeks, buggy implementation.
		@param modname: You can specify that the value depends on a module (or a list of modules).
			(You give the module itself here.)  In which case, it tries
			to detect changes to the specified modules.   See L{modinfo} for details.
		@type mod: L{module} or L{tuple}(L{module}).
		@note: Certain compromises were made in the handling of C{modname} and C{mod}.
			Even if you use them, you are not 100% guaranteed to be protected from
			all changes to the code used to compute the cached values.   To be
			entirely safe, you should manually clear the cache when ever you
			change your code.   However, this will probably save your tail if
			you forget to clear the cache.  See L{modinfo} for details.
		@note: If there is an error when reading files
			(i.e. if C{fname}, C{modname}, or C{mod} is specified),
			then the object will be constructed with C{info=None}.
			This will lead to a L{OSError} if you then call L{load} on the object,
			which is what you'd get from a cache miss.
			Calling L{dump} or L{bg_dump} will silently do nothing.
		"""
		if info is None:
			raise ValueError, "info=None.  Is this the result of a previous constructor failure?"
		self.info = info
		try:
			if fname is not None:
				self.info += fileinfo(fname)
			if modname is not None:
				self.info += namedModInfo(modname)
			if mod is not None:
				self.info += modinfo(mod, {})
		except CannotGetStat, ex:
			die.warn("Error in getting cache info for %s" % str(ex))
			self.info = None
		self.root = root
		self.tail = '.pickle'
		self._dname = None	#: Pathname info, computed only when needed
		self._fpath = None	#: Pathname info, computed only when needed

	def __repr__(self):
		return "<cache_info %s %s %s>" % (self.root, self.info, self.tail)

	def copy(self):
		if self.info is None:
			return None
		return cache_info(self.root, info=self.info)


	def addinfo(self, *s, **kv):
		"""
		@note: This does I{not} modify self!  It creates a new object.
		"""
		if self.info is None:
			return None
		return cache_info(root=self.root, info=self.info+s,
					fname=kv.get('fname', None),
					modname=kv.get('modname', None),
					mod=kv.get('mod', None)
					)


	def makespace(self, avoid):
		n = 0
		p = 0
		when = time.time() - self.Age
		while n<2 and p<3:
			n += _dodel(self.root, avoid, when, 1.0/self.NumObj)
			p += 1
		if DEBUG and n > 0:
			die.info("Deleted %d files from the cache %s" % (n, self.root))


	def dump(self, e):
		"""Cache some data on the disk.
		@rtype: could be anything picklable.
		@return: whatever was passed as C{e}.
		@param e: the data to write.
		@type e: anything picklable.
		@note: This function might quietly fail.   Since this is a cache, failure to write
			is not considered a major problem.   In my experience, failure to write is
			often caused by intermittent network problems, and you don't want it to crash
			a long-running computation.
		"""
		if self.info is None:
			return
		li = len(self.info)
		assert li>0
		i1 = self.info[:(1+li)//2]
		i2 = self.info[(li-1)//2:]
		dname, fpath = self.cachepath()
		fd = None
		try:
			gpkmisc.makedirs(dname)
			fd = open(fpath, 'w')
			# We are assuming that the pickel protocol is sequential,
			# and that disk writes are sequential, so that if the two
			# ends of the dump are correct (i1 and i2) then the stuff
			# in between will be correct.
			cPickle.dump((i1, e, i2), fd, protocol=cPickle.HIGHEST_PROTOCOL)
			# os.rename(ftmp, fpath)
			# self.makespace(fpath)
		except (OSError, IOError, cPickle.PicklingError), ex:
			die.warn("cache:dump: %s on %s" % (ex, fpath))
			try:
				os.remove(fpath)
			except (IOError, OSError):
				pass
		finally:
			if fd is not None:
				try:
					fd.close()
				except IOError, ex:
					die.warn("Dump to cache failed on fd.close(): %s" % ex)
		self._clean.put((self, fpath))
		return e
	

	def bg_dump(self, e):
		if self.info is None:
			return
		if len(self.info)==0:
			raise ValueError, "Need to specify cache info"
		# t = threading.Thread(target=self.dump, args=(e,), name='cache_dumper%s' % id(e))
		# t.start()
		self._dump.put((self, e))


	def load(self):
		"""Pull in some data from the disk.
		@rtype: could be anything picklable.
		@return: whatever was cached on disk.
		@raise BadFileFormat: when the data isn't valid.
		@raise OSError: on cache miss.
		@raise IOError: e.g. network problems.
		"""
		if self.info is None:
			raise IOError, "cache_path.load: does not know what file to open, as info=None"
		dname, fpath = self.cachepath()
		fd = open(fpath, 'r')
		x = cPickle.load(fd)
		del fd
		li = len(self.info)
		assert li>0
		if(isinstance(x, tuple) and len(x)==3
			and len(x[0])>0 and len(x[2])>0 and len(x[0])+len(x[2])==li+1
			and x[0][-1]==x[2][0]
			and x[0]==self.info[:(1+li)//2] and x[2]==self.info[(li-1)//2:]
			):
			if DEBUG:
				die.info("Cache hit: %s" % fpath)
			return x[1]
		die.info("Cache fail: %s" % fpath)
		# Cache fail should catch cases of file truncation and data corruption,
		# even if the cPickle.load call doesn't.
		raise BadFileFormat, "Cache fail %s" % fpath


	def cachepath(self):
		"""Return a pathname suitable for cacheing some result.
		@return: C{(path_to_root,path_with_tail)}.   C{Path_to_root} is/will be a directory;
			C{path_with_tail} is a path to a data file within that directory.
			Normally, the actual cache is at the location
			C{os.path.join(path_to_root,path_with_tail)} on the disk;
			that is what you would pass as the C{fname} argument to L{load_cache} or
			L{dump_cache}.
		@rtype: C{tuple(str, str)}
		@raise ValueError: if you haven't specified any C{info} yet.
		"""
		if self._dname is not None:
			return (self._dname, self._fpath)
		if len(self.info)==0:
			raise ValueError, "Need to specify cache info"
		hf = hashlib.md5(repr(self.info)).hexdigest()
		d1 = hf[:2]
		d2 = hf[2:4]
		d3 = hf[4:6]
		fc = hf[6:]
		self._dname = os.path.join(self.root, d1, d2, d3)
		self._fpath = os.path.join(self._dname, fc) + self.tail
		return (self._dname, self._fpath)

	
	Errors = (IOError, EOFError, UnpicklingError, BadFileFormat)



def walkcache(top):
	"""This is to help humans read the cache.
	"""
	for (dirpath, dirnames, filenames) in os.walk(top):
		for fn in filenames:
			p = os.path.join(dirpath, fn)
			try:
				x = cPickle.load(open(p, 'r'))
			except cPickle.UnpicklingError, ex:
				print '#', p, str(ex)
				continue
			if(isinstance(x, tuple) and len(x)==3
				and len(x[0])>0 and len(x[2])>0
				and x[0][-1]==x[2][0]
				):
				print p, x[0][:-1] + x[2]
			else:
				print '#', p, '???'


def test_errs():
	cache_info('/tmp', fname="/u/ajdn/wlkenql/q!!")

def test_normal():
	import shutil
	import g_exec
	N = 500
	tmp = '/tmp/ci_cache_py_test'
	cache_info.NumObj = N
	cache_info.Age = 1.0
	for i in range(2*N):
		cache_info(tmp, info=(i,0), mod=die).bg_dump(i%97)
		# if i%100 == 99:
			# die.info("filling: i=%d" % i)
	nmiss = 0
	nhit = 0
	for i in range(2*N-N//10,2*N):
		try:
			x = cache_info(tmp, info=(i,0), mod=die).load()
			assert x == (i%97)
			nhit += 1
		except IOError:
			nmiss += 1
	print 'BG dump done'
	for i in range(2*N):
		cache_info(tmp, info=(i,1), mod=die).dump(i%197)
		# if i%100 == 99:
			# die.info("filling: i=%d" % i)
	for i in range(2*N-N//10,2*N):
		try:
			x = cache_info(tmp, info=(i,1), mod=die).load()
			assert x == (i%197)
			nhit += 1
		except IOError:
			nmiss += 1
	die.info("Cache: misses=%d hits=%d" % (nmiss, nhit))
	assert nmiss < nhit//10
	n = len(list(g_exec.getiter_raw(None, ["find", tmp, "-type", "f", "-name", "*.pickle"])))
	die.info("Number of files=%d" % n)
	assert 0.3*N < n < 3*N
	shutil.rmtree(tmp);

def test():
	test_normal()
	test_errs()

if __name__ == '__main__':
	test()
	# import sys
	# walkcache(sys.argv[1])
