#!/usr/bin/env python
"""
Simple filesystem tester.  Writes and reads back a set of large 1GB files.
Mark Moraes
"""

import md5 as hash
import sys, os, time
from cStringIO import StringIO

def main():
	if len(sys.argv) != 3:
		print >> sys.stderr, "Usage: filetest filename size-in-MiB"
		return 1
	progname, fname, size = sys.argv
	try:
		nmax = long(size)
		if nmax < 0:
			print >> sys.stderr, "%s: Error: size %s must be > 0" % (progname, size)
			return 3
	except Exception, e:
		print >> sys.stderr, "%s: Error: failed testing size %s: %s" % (progname, `size`, str(e))
		return 4
	maxsize = 1024
	k1 = long(nmax/maxsize)
	k2 = long(nmax%maxsize)
	for i in xrange(k1):
		verifyfile(progname, "%s-%d" % (fname, i), maxsize)
	if k2:
		verifyfile(progname, "%s-%d" % (fname, k1), k2)

def verifyfile(progname, fname, nmax):
	try:
		f = open(fname, 'wb')
	except Exception, e:
		print >> sys.stderr, "%s: Error: failed opening file %s for writing: %s" % (progname, fname, str(e))
		return 5
	gen = hash.new("FILE TEST")
	n = 0
	blksz = 1024*1024
	blk = genblock(gen, blksz-10)
	reportinterval = max(min(int(nmax / 10),100),1)
	p = Progress(nmax, "Writing " + fname, "MiB", reportinterval)
	while n < nmax:
 		s = ("N%9s" % n) + blk
		try:
			f.write(s)
		except Exception, e:
			print >> sys.stderr, "%s: Error: writing block %d to %s: %s" % (progname, n, fname, str(e))
			return 5
		n += 1
		p.click()
 	try:
		f.flush()
		f.close()
		sys.stdout.flush()
	except Exception, e:
		print >> sys.stderr, "%s: Error: flushing file %s: %s" % (progname, fname, str(e))
		return 6
	p.done()

	# Now read back and verify
	try:
		f = open(fname, 'rb')
	except Exception, e:
		print >> sys.stderr, "%s: Error: opening file %s for reading: %s" % (progname, fname, str(e))
		return 7

	n = 0
	p = Progress(nmax, "Verifying "+fname, "MiB", reportinterval)
	while n < nmax:
 		s = ("N%9s" % n) + blk
		try:
			d = f.read(blksz)
		except Exception, e:
			print >> sys.stderr, "%s: Error: reading from file %s block %d: %s" % (progname, fname, n, str(e))
			return 8
		n += 1
		p.click()
		assert len(d) == blksz
		if s != d:
			print >> sys.stderr, "%s: Error: mismatch in block %d" % (progname, n)
			return 9
	p.done()
	
class Progress:
	def __init__(self, maxclicks, desc, units, reportinterval):
		self.maxclicks = maxclicks
		self.newclicks = self.clicks = 0
		self.start = self.tlast = time.time()
		self.desc = desc
		self.units = units
		self.reportinterval = reportinterval
		sys.stdout.write('>>> ' + desc + '\n')
	def click(self, nclicks = 1):
		self.newclicks += nclicks
		if self.newclicks >= self.reportinterval:
			self.clicks += self.newclicks
			t = self.report(self.newclicks, self.tlast)
			self.newclicks = 0
			self.tlast = t
	def report(self, nclicks, tlast):
		t = time.time()
		dt = t - self.start
		rate = nclicks / (t - tlast)
		left = (self.maxclicks - self.clicks)/rate
		sys.stdout.write("%9.1f secs %9d %s %9.1f %s/sec %9.1f secs left\n" % (dt, self.clicks, self.units, rate, self.units, left))
		return t
	def done(self):
		sys.stdout.write("Done "+self.desc+'\n')
		self.report(self.clicks, self.start)

def main2():
	gen = hash.new('FILE TEST')
	blksz = 1024*1024
	nmax = 10
	p = Progress(nmax, "Generating", "MiB", 1)
	for i in xrange(nmax):
		s = genblock(gen, blksz)
		p.click()
	p.done()

def main3():
	gen = hash.new('FILE TEST')
	blksz = 1024*1024
	blk = genblock(gen, blksz-10)
	nmax = 100
	p = Progress(nmax*nmax, "Generating", "MiB", 100)
	for n in xrange(nmax):
		for i in xrange(nmax):
	 		s = ("N%9s" % n) + blk
		p.click(nmax)
	p.done()

def genblock(gen, blksz):
	k = int((blksz + gen.digest_size - 1)/gen.digest_size)
	s = StringIO()
	for i in xrange(k):
		gen.update("I%09d" % i)
		s.write(gen.digest())
	s.seek(0)
	ret = s.read()[:blksz]
	assert len(ret) >= blksz
	return ret

if __name__ == '__main__':
	sys.exit(main())