#!/usr/bin/python
#
# This program is free software; you can redistribute it and/or modify
# it under the terms of the GNU Library General Public License as published by
# the Free Software Foundation; version 2 only
#
# 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 Library General Public License for more details.
#
# You should have received a copy of the GNU Library 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.
# Copyright 2004 Dag Wieers <dag@wieers.com>

import os, signal, socket, sys, threading, time

class Reader(threading.Thread):
	def __init__(self, path, filesize, bitrate):
		threading.Thread.__init__(self)
		self.name = self.getName()
		self.filename = os.path.join(path, hostname, 'dload-reader-%s-%1.1fmbps.bin' % (self.name, bitrate))
		self.filesize = filesize
		self.bitrate = bitrate * 1024 * 1024
		self.block = threading.Event()

		if self.filesize / self.bitrate < 16:
			self.filesize = self.bitrate * 8 * 2

		if not os.path.exists(os.path.join(path, hostname)):
			os.mkdir(os.path.join(path, hostname))

		if not os.path.exists(self.filename):
			info(1, 'File %s does not exist, creating.' % self.filename)
			self.fd = open(self.filename, 'w+')
			self.fd.truncate(self.filesize)
			self.fd.write(charlist * int(self.filesize/256))
			self.fd.flush(); os.fsync(self.fd)
			self.fd.close()

	def handler(self):
		self.block.set()

	def run(self):
		mbitrate = self.bitrate / 1024 / 1024
		print '%s: Start reading file %s at %1.1f Mbps (%1.2f MB/s)' % (self.name, self.filename, mbitrate, mbitrate / 8.0)
		while True:
			t1 = []; t2 = []; avg = 0; slip = 0

			fd = os.open(self.filename, os.O_RDONLY|os.O_NONBLOCK)
			fo = os.fdopen(fd, 'r', int(self.bitrate/slice))

			for i in range(int(self.filesize / self.bitrate / slice)):
				threading.Timer(slice, self.handler).start()
				t1.append(float(time.time()))
				fo.read(int(self.bitrate/slice/256/8))
				t2.append(float(time.time()))
#				print '%s: Read in %1.4f secs at %1.4f Mbps' % (self.getName(), t2[-1]-t1[-1], mbitrate/slice/(t2[-1]-t1[-1]))
				if t2[-1]-t1[-1] > slice:
					print '%s: WARNING: Reading %d KB at %1.1f Mbps took %1.2f secs at only %1.1f Mbps' % (self.name, int(mbitrate*1024/slice/8), mbitrate, t2[-1]-t1[-1], mbitrate/slice/(t2[-1]-t1[-1]))
					slip = slip + 1
				self.block.wait(slice)
				self.block.clear()
			for j in range(i): avg = avg + t2[j]-t1[j]; avg = avg / len(t1)
			if avg == 0: avg = 0.00000001
			print '%s: Restart reading at %1.1f Mbps. Avg bitrate: %1.1f Mbps, Slips: %d' % (self.name, mbitrate, mbitrate/slice/avg, slip)
			fo.close()

class Writer(threading.Thread):
	def __init__(self, path, filesize, bitrate):
		threading.Thread.__init__(self)
		self.name = self.getName()
		self.filename = os.path.join(path, hostname, 'dload-writer-%s-%1.1fmbps.bin' % (self.name, bitrate))
		self.filesize = filesize
		self.bitrate = bitrate * 1024 * 1024
		self.block = threading.Event()

		if self.filesize / self.bitrate < 16:
			self.filesize = self.bitrate * 8 * 2

		if not os.path.exists(os.path.join(path, hostname)):
			os.mkdir(os.path.join(path, hostname))

	def handler(self):
		self.block.set()

	def run(self):
		mbitrate = self.bitrate / 1024 / 1024
		print '%s: Start writing file %s at %1.1f Mbps (%1.1f MB/s)' % (self.name, self.filename, mbitrate, mbitrate / 8.0)
		while True:
			t1 = []; t2 = []; avg = 0; slip = 0

			fd = os.open(self.filename, os.O_CREAT|os.O_WRONLY|os.O_NONBLOCK|os.O_SYNC)
			fo = os.fdopen(fd, 'w', int(self.bitrate/slice))
			fo.truncate(self.filesize)
			os.fdatasync(fd)

			for i in range(int(self.filesize / self.bitrate / slice)):
				threading.Timer(slice, self.handler).start()
				t1.append(float(time.time()))
				fo.write(charlist * int(self.bitrate/slice/256/8))
				fo.flush(); os.fsync(fd)
				t2.append(float(time.time()))
#				print '%s: Written in %1.4f secs at %1.4f Mbps' % (self.getName(), t2[-1]-t1[-1], mbitrate/slice/(t2[-1]-t1[-1]))
				if t2[-1]-t1[-1] > slice:
					print '%s: WARNING: Writing %d KB at %1.1f Mbps took %1.2f secs at only %1.1f Mbps' % (self.name, int(mbitrate*1024/slice/8), mbitrate, t2[-1]-t1[-1], mbitrate/slice/(t2[-1]-t1[-1]))
					slip = slip + 1
				self.block.wait(slice)
				self.block.clear()
			fo.close()
			for j in range(i): avg = avg + t2[j]-t1[j]; avg = avg * 1.0 / len(t1)
			if avg == 0: avg = 0.00000001
			print '%s: Restart writing at %1.1f Mbps. Avg bitrate: %1.1f Mbps, Slips: %d' % (self.name, mbitrate, mbitrate/slice/avg, slip)

def info(level, str):
	"Output info message"
#	if level <= op.verbose:
	print str

def die(ret, str):
	"Print error and exit with errorcode"
	info(0, str)
	sys.exit(ret)

def charlist():
	ret = ''
	for i in range(256):
		ret = ret + chr(i)
	return ret

def main():
	global charlist, hostname, slice

	charlist = charlist()
#	hostname = os.uname()[1].split('.')[0]
	hostname = socket.gethostbyaddr(socket.gethostname())[0]
	slice = 1

	filesize = 1
	path = '/mnt/extra'
#	path = '/dev/shm'
	writer, reader = (0, 1)

	threadlist = (
		( writer, path, filesize, 10 ),
		( writer, path, filesize, 15 ),
		( writer, path, filesize, 20 ),
		( writer, path, filesize, 50 ),
		( writer, path, filesize, 8 ),
		( writer, path, filesize, 5 ),
		( writer, path, filesize, 2.5 ),
		( reader, path, filesize, 5 ),
		( reader, path, filesize, 2.5 ),
#		( reader, path, filesize, 20 ),
	)

	threads = []; wtr = wbr = rtr = rbr = 0
	for th in threadlist:
		role, path, filesize, bitrate = th
		if role == writer:
			wtr = wtr + 1
			wbr = wbr + bitrate
			threads.append(Writer(path, filesize, bitrate))
		elif role == reader:
			rtr = rtr + 1
			rbr = rbr + bitrate
			threads.append(Reader(path, filesize, bitrate))
		else:
			info(1, 'Role %d not implemented yet!' % role)

	info(2, 'Starting %d write (%1.1f MB/s) and %d read (%1.1f MB/s) threads.' % (wtr, wbr / 8, rtr, rbr / 8))
	for th in threads:
		th.start()
		time.sleep(slice / len(threadlist))

### Main entrance
if __name__ == '__main__':
	try:
#		import profile
#		profile.run('main()','dloadprof')
#		import pstats
#		p = pstats.Stats('dloadprof')
#		p.strip_dirs().sort_stats(-1).print_stats()
#		p.sort_stats('name')
#		p.print_stats()
#		p.sort_stats('cumulative').print_stats(10)
		main()
	except KeyboardInterrupt, e:
		die(6, 'Exiting on user request')
	except OSError, e:
#		print e.errno
		die(7, 'OSError: %s' % e)

# vim:ts=4:sw=4