#!/usr/bin/env python

# This version is much better abstracted than the first attempt

import os
import pty
import sys
import tty
import time
import fcntl
import random
import select
import signal
import struct
import termios

version = '1.01'

sys.path.insert(0, '/usr/local/lib')
sys.path.insert(0, os.path.expanduser('~/lib'))

def usage(retval):
	sys.stderr.write("Usage: %s [file]\n" % sys.argv[0])
	sys.stderr.write('\t-a           Append the output to file or typescript, retaining the prior contents\n')
	sys.stderr.write('\t-c command\n')
	sys.stderr.write('\t-f           Flush output after each write.\n')
	sys.stderr.write('\t-q           Be quiet.\n')
	sys.stderr.write('\t-d           Use dated files.\n')
	sys.stderr.write('\t-t           Output timing data\n')
	sys.stderr.write('\t             In dated mode, using a .timing file.\n')
	sys.stderr.write('\t             In non-dated mode, use standard error.\n')
	sys.stderr.write('\t--verbosity levelno\n')
	sys.stderr.write('\t-v\n')
	sys.stderr.write('\t--version    report the version of the program\n')
	sys.stderr.write('\t--help\n')
	sys.exit(retval)

def rnd():
	# We let python worry about seeding.  For this application, that should be sufficient
	characters = '0123456789abcdefghijklmnopqrstuvwxyzABCDEFGHIJKLMNOPQRSTUVWXYZ'
	len_characters = len(characters)
	result_list = []
	for i in xrange(12):
		result_list.append(characters[int(random.random() * len_characters)])
	return ''.join(result_list)

class Timed_dated_file:
	def __init__(self, program_name, base_output_filename, dated_files=False, append=False, quiet=False, flush=False, \
		timing_data=False):
		self.program_name = program_name.title()
		self.base_output_filename = base_output_filename
		self.dated_files = dated_files
		self.append = append
		self.quiet = quiet
		self.flush = flush
		self.timing_data = timing_data
		self.timing_file = None
		self.prior_filename = ''
		self.prior_session_file_part = ''
		self.file = None
		self.filename = ''
		self.session_file_part = ''
		self.first_time = True
		if dated_files:
			self.session_id = rnd()
		else:
			self.session_id = ''
		self.__open()

	def close(self, continued=False):
		if not self.quiet:
			s = '%s finished, file was %s\r\n' % (self.program_name, self.prior_filename)
			if continued:
				s = s + '...To be continued\r\n'
			print s
			sys.stdout.flush()
			self.write(s)
		if self.file != None:
			self.file.close()
		if self.timing_file != sys.stderr and self.timing_file != None:
			self.timing_file.close()

	def flush(self):
		self.file.flush()
		if self.timing_data:
			self.timing_file.flush()

	def __open(self):
		if self.dated_files:
			year = time.strftime('%Y')
			month = time.strftime('%m')
			day = time.strftime('%d')
			year_and_month = os.path.join(year, month)
			year_month_and_day = os.path.join(year_and_month, day)
			self.session_file_part = \
				os.path.join(year_month_and_day, '%s-%s' % (self.base_output_filename, self.session_id))
#				+ time.strftime('-%H-%M')
		else:
			self.session_file_part = self.base_output_filename
		if self.session_file_part != self.prior_session_file_part:
			self.prior_session_file_part = self.session_file_part
			if self.dated_files:
				self.filename = '%s-%s' % (self.session_file_part, rnd())
			else:
				self.filename = self.session_file_part
			if self.file != None:
				self.close(continued=True)
			if self.append:
				open_flag='a'
			else:
				open_flag='w'
			if self.dated_files:
				# Try to open the file.  Only if it that fails do we try to create any leading subdirectories
				# This should be quite a bit faster
				try:
					self.file = open(self.filename, open_flag)
				except IOError:
					# The open failed with an OSError, possibly a No such file or directory
					# Create the leading directories and try again.  If that fails too,
					# go ahead and error out.
					for dir in [ year, year_and_month, year_month_and_day ]:
						if not os.path.isdir(dir):
							os.mkdir(dir)
					self.file = open(self.filename, open_flag)
			else:
				self.file = open(self.filename, open_flag)
			if not self.quiet:
				s = '%s started, file is %s\r\n' % (self.program_name, self.filename)
				sys.stdout.write(s)
				#sys.stdout.flush()
				self.file.write(s)
			self.prior_session_file_part = self.session_file_part
			self.prior_filename = self.filename
			if self.timing_data:
				self.prior_time = time.time()
				if self.dated_files:
					self.timing_filename = '%s.timing' % self.filename
					self.timing_file = open(self.timing_filename, open_flag)
					if self.first_time:
						self.first_time = False
				else:
					self.timing_file = sys.stderr
				self.timing_file.write('0.0 1\n')

	def write(self, data):
		self.__open()
		if self.timing_data:
			current_time = time.time()
			self.timing_file.write('%f %d\n' % (current_time - self.prior_time, len(data)))
			self.prior_time = current_time
		self.file.write(data)
		if self.flush:
			self.file.flush()
			if self.timing_data:
				self.timing_file.flush()

def resize(dummy1, dummy2):
	if hasattr(termios, 'TIOCGWINSZ'):
		ws = struct.pack("HHHH", 0, 0, 0, 0)
		ws = fcntl.ioctl(sys.stdin.fileno(), termios.TIOCGWINSZ, ws)
		#lines, columns, x, y = struct.unpack("HHHH", ws)
		fcntl.ioctl(globals.fd, termios.TIOCSWINSZ, ws)
	os.kill(globals.pid, signal.SIGWINCH)

class Globals:
	def __init__(self):
		pass

globals = Globals()

def main():
	verbosity=0
	base_output_filename = 'typescript'
	append=False
	command=''
	flush=False
	timing_data=False
	dated_files=False
	use_psyco=True
	quiet=False

	while sys.argv[1:]:
		if sys.argv[1] in [ '--help', '-h' ]:
			usage(0)
		elif sys.argv[1] == '--verbosity':
			verbosity=int(sys.argv[2])
			del sys.argv[1]
		elif sys.argv[1] == '--version':
			print version
			sys.exit(0)
		elif sys.argv[1] == '-v':
			verbosity+=1
		elif sys.argv[1] == '-a':
			append=True
		elif sys.argv[1] == '-c':
			command=sys.argv[2]
			del sys.argv[1]
		elif sys.argv[1] == '-f':
			flush=True
		elif sys.argv[1] == '-q':
			quiet=True
		elif sys.argv[1] == '-t':
			timing_data=True
		elif sys.argv[1] == '-d':
			dated_files=True
		elif sys.argv[1] == '--no-psyco':
			use_psyco=False
		elif sys.argv[1].startswith('-'):
			sys.stderr.write("%s: Illegal option: %s\n" % (sys.argv[0], sys.argv[1]))
			usage(1)
		else:
			base_output_filename = sys.argv[1]
			del sys.argv[1]
			if sys.argv[1:]:
				sys.stderr.write('%s: Too many arguments starting with %s\n' % (sys.argv[0], sys.argv[1]))
				sys.exit(1)
			break
		del sys.argv[1]

	program_name = os.path.basename(sys.argv[0])

	if append and dated_files:
		sys.stderr.write('%s: You may not combine -a and -d\n' % sys.argv[0])
		usage(1)

	if use_psyco:
		try:
			import psyco
			psyco.full()
		except:
			if verbosity >= 1:
				sys.stderr.write('%s: psyco initialization failed\n' % sys.argv[0])
		else:
			if verbosity >= 1:
				sys.stderr.write('%s: psyco initialization succeeded\n' % sys.argv[0])

	pid, fd = pty.fork()

	globals.pid = pid
	globals.fd = fd

	if pid == 0:
		# child process
		if os.environ.has_key('SHELL'):
			shell=os.environ['SHELL']
		else:
			shell='/bin/sh'
		if command == '':
			#os.execvp('bash', [ 'bash', '-i' ])
			os.execvp(shell, [ shell, '-i' ])
		else:
			os.execvp(shell, [ shell, '-c', command ])
	else:
		# parent process
		signal.signal(signal.SIGWINCH, resize)
		resize(None, None)
		stdin = sys.stdin.fileno()
		stdout = sys.stdout.fileno()
		stderr = sys.stderr.fileno()
		read_length=2^16
		timed_dated_file = Timed_dated_file(program_name, base_output_filename, dated_files, append, quiet, flush, timing_data)
		old_termios = termios.tcgetattr(stdin)
		tty.setraw(stdin)
		while 1:
			while 1:
				try:
					(input_fds, output_fds, exception_fds) = select.select([stdin, fd], [], [])
				except select.error:
					pass
				else:
					break
			if exception_fds:
				sys.stderr.write('%s: should never get an exception fd from select\n' % sys.argv[0])
				sys.exit(1)
			if output_fds:
				sys.stderr.write('%s: should never get an output fd from select\n' % sys.argv[0])
				sys.exit(1)
			if fd in input_fds:
				try:
					ch = os.read(fd, read_length)
				except OSError:
					# Unfortunately, with pty's a read past EOF (IOW, when the
					# child process exits) raises an OSError exception
					termios.tcsetattr(stdin, termios.TCSADRAIN, old_termios)
					timed_dated_file.close()
					sys.exit(0)
				os.write(stdout, ch)
				timed_dated_file.write(ch)
			if stdin in input_fds:
				ch = os.read(stdin, read_length)
				timed_dated_file.write('')
				os.write(fd, ch)

main()