import os,sys,string,getopt

import fchksum

#from time import time as timer
import time
timer = time.time

def pathexists(p, whenfound, notfound=None):
	path=string.split(os.environ.get('PATH',os.defpath), os.pathsep)
	for d in path:
		if os.path.exists(os.path.join(d,p)):
			return whenfound
	return notfound

import md5,zlib

try:
	import mmap
	if hasattr(mmap, 'PROT_READ'):
		def dommap(fileno, len):#unix mmap.  python default is PROT_READ|PROT_WRITE, but we open readonly.
			if len==0: return '' #mmap doesn't like length=0
			return mmap.mmap(fileno, len, mmap.MAP_SHARED, mmap.PROT_READ)
	else:
		def dommap(fileno, len):#windows mmap.
			if len==0: return ''
			return mmap.mmap(fileno, len)
	def getfilemd5_mmap(file):
		f=open(file,'rb')
		s = os.path.getsize(file)
		m = md5.new(dommap(f.fileno(), s))
		c=""
		for z in m.digest():
			c = c+'%02x'%ord(z)
		return c,s
	def getfilecrc_mmap(file):
		f=open(file,'rb')
		s = os.path.getsize(file)
		c = zlib.crc32(dommap(f.fileno(), s))
		return "%08X"%c,s

except ImportError:
	getfilecrc_mmap = None
	getfilemd5_mmap = None

def getfilemd5(file):
	f=open(file,'rb')
	m=md5.new()
	s=0L
	while 1:
		x=f.read(65536)
		if not len(x):
			c=""
			for z in m.digest():
				c=c+'%02x'%ord(z)
			return c,s
		s=s+len(x)
		m.update(x)
		
def getfilecrc(file):
	f=open(file,'rb')
	c=zlib.crc32('')
	s=0L
	while 1:
		x=f.read(65536)
		if not len(x):
			return "%08X"%c,s
		s=s+len(x)
		c=zlib.crc32(x,c)

chksums = {
	'md5':(
		('fchksum', fchksum.fmd5t,),
		('pymd5', getfilemd5,),
		('pymd5+mmap', getfilemd5_mmap,),
		('md5sum', pathexists('md5sum', lambda fn: os.system('md5sum -b %s'%fn)),),
		('md5', pathexists('md5', lambda fn: os.system('md5 %s'%fn)),),
	),
	'crc32':(
		('fchksum', fchksum.fcrc32t,),
		('pycrc', getfilecrc,),
		('pycrc+mmap', getfilecrc_mmap,),
#		('crc32', lambda fn: os.system('crc32 -b %s'%fn),),
		('cksfv', pathexists('cksfv', lambda fn: os.system('cksfv %s'%fn)),),
	),
	'cksum':(
		('fchksum', fchksum.fcksumt,),
		('cksum', pathexists('cksum', lambda fn: os.system('cksum %s'%fn)),),
	),
	'bsdsum':(
		('fchksum', fchksum.fbsdsumt,),
		('sum', pathexists('sum', lambda fn: os.system('sum %s'%fn)),),
	),
	'sysvsum':(
		('fchksum', fchksum.fsysvsumt,),
		('sum', pathexists('sum', lambda fn: os.system('sum -s %s'%fn)),),
	),
}

def runtest(func, arg, n=1):
	start=timer()
	for i in range(0,n):
		r=func(arg)
		if r:
			print r
	return timer()-start

if __name__=='__main__':
	def printusage(err):
		print 'speedtest.py [-n number of runs] [-t tests to run] [--fchksum-only] filename'
		sys.exit(err)

	try:
		optlist, args = getopt.getopt(sys.argv[1:], 'n:t:h', ['help', 'fchksum-only'])
	except getopt.error, a:
		print "speedtest: %s"%a
		printusage(1)

	testruns = 1
	tests = chksums.keys()
	test_fchksum_only = 0
	for o,a in optlist:
		if o=='-n':
			testruns = int(a)
		elif o=='-t':
			tests = string.split(a,',')
		elif o=='--fchksum-only':
			test_fchksum_only = 1
		elif o=='-h' or o=='--help':
			printusage(0)

	testfile = args[0] #for now, only allow one filename
	fchksum.fcrc32(testfile) #read once so its cached already

	print 'tests to run:',tests
	for t in tests:
		print 'testing',t
		for k, func in chksums[t]:
			if test_fchksum_only and k != 'fchksum':
				continue
			if func:
				elapsed = runtest(func, testfile, testruns)
				print k, 'processed', testfile, 'in', elapsed, '(', os.path.getsize(testfile)*testruns / elapsed / 1024, 'K/s)'