#!/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 errno
import fcntl
import random
import select
import signal
import struct
import termios

version = '1.30'

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-s           Output window size data\n')
	sys.stderr.write('\t             Use .sizing file whether in dated mode or not\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, sizing_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.sizing_data = sizing_data
		self.timing_file = None
		self.sizing_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()
		if self.sizing_file != None:
			self.sizing_file.close()

	def flush(self):
		self.file.flush()
		if self.timing_data:
			self.timing_file.flush()
		# we don't need to flush the sizing file here, because it's done in the SIGWINCH signal handler

	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 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')
			if self.sizing_data:
				self.sizing_filename = '%s.sizing' % self.filename
				self.sizing_file = open(self.sizing_filename, open_flag)
				globals.sizing_file = self.sizing_file
				globals.flush = self.flush
				respond_to_resize_event(None, None)
			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)
				globals.byte_no += len(s)
			self.prior_session_file_part = self.session_file_part
			self.prior_filename = self.filename

	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)
		globals.byte_no += len(data)
		if self.flush:
			self.file.flush()
			if self.timing_data:
				self.timing_file.flush()

def respond_to_resize_event(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)
		if globals.sizing_data:
			globals.sizing_file.write('%d %d %d\n' % (globals.byte_no, lines, columns))
			if globals.flush:
				# we could almost just flush this irrespective of -f, because this file should change pretty
				# infrequently
				globals.sizing_file.flush()
	os.kill(globals.pid, signal.SIGWINCH)

class Globals:
	def __init__(self):
		pass

globals = Globals()

def maybe_restarted_syscall(fn, loop_exc, loop_errnos):
	while 1:
		try:
			result = fn()
		except loop_exc, (errno, error_string):
			if errno in loop_errnos:
				continue
			else:
				sys.stderr.write('%s: %s\n' % (sys.argv[0], error_string))
				raise
		else:
			break
	return result

def main():
	verbosity=0
	base_output_filename = 'typescript'
	append=False
	command=''
	flush=False
	timing_data=False
	sizing_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] == '-s':
			sizing_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
	globals.sizing_data = sizing_data
	globals.sizing_file = None
	globals.byte_no = 0
	globals.flush = flush

	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
		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, sizing_data)
		signal.signal(signal.SIGWINCH, respond_to_resize_event)
		old_termios = termios.tcgetattr(stdin)
		#tty.setraw(stdin)
		maybe_restarted_syscall(lambda : tty.setraw(stdin), termios.error, [ errno.EINTR ])
		while 1:
#			Traceback (most recent call last):
#				File "./pypty", line 332, in <module>
#					main()
#				File "./pypty", line 309, in main
#					(input_fds, output_fds, exception_fds) = select.select([stdin, fd], [], [])
#					select.error: (4, 'Interrupted system call')
			#(input_fds, output_fds, exception_fds) = select.select([stdin, fd], [], [])
			(input_fds, output_fds, exception_fds) = \
				maybe_restarted_syscall(lambda : select.select([stdin, fd], [], []), select.error, [ errno.EINTR ])
			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)
				maybe_restarted_syscall(lambda : os.write(stdout, ch), OSError, [ errno.EINTR ])
				timed_dated_file.write(ch)
			if stdin in input_fds:
				ch = os.read(stdin, read_length)
				timed_dated_file.write('')
				#os.write(fd, ch)
				maybe_restarted_syscall(lambda : os.write(fd, ch), OSError, [ errno.EINTR ])

#File "/home/dstromberg/bin/pypty", line 294    , in <module>
	#main()
#File "/home/dstromberg/bin/pypty", line 287, in main
	#os.write(stdout, ch)
	#OSError: [Errno 4] Interrupted system call

main()