"""
	A python module to control a linux style
	terminal

	Copyright (c) 2009, Phillip Berndt

	Available under the terms of the BSD (Berkeley) license. If you fail to
	find the text yourself, you're prohibited to use this module.

	This program can also be used from the command line. Try invoking
	 ./consolectrl -h
	for help.

	For more information, see console_codes(4)
"""
import sys
import termios
import tty
import os
import warnings
import fcntl
import array

class Consolectrl(object):
	"""
		Control terminals, i.e. color, cursor position, etc.

		To be honest: Control Linux terminals. I didn't test this
		for VT100-only devices.
	"""

	def _send_esc_seq(self, sequence): # {{{
		"""
			Send <ESC>sequence to the terminal
		"""
		sys.stdout.write("\033" + sequence)
		sys.stdout.flush()
	# }}}
	def _send_and_receive_esc_seq(self, sequence, answer_ends_with = None, num_chars = None, autosplit = False): # {{{
		"""
			Send an escape sequence and read the response
			from the terminal.

			The response is assumed to be ending with
			answer_ends_with. You may specify a number
			of chars alternatively.
			
			If autosplit is enabled,
			 <ESC>[xx;yyE stuff will be returned as
			(xx,yy).

		"""
		if answer_ends_with and num_chars:
			raise Exception("Specify either answer_ends_with or read_chars")
		if num_chars and autosplit:
			raise Exception("read_chars is incompatible with autosplit")

		stored_attrs = termios.tcgetattr(sys.stdin)
		tty.setraw(sys.stdin)
		try:
			self._send_esc_seq(sequence)

			resp = ""
			is_escaped = sys.stdin.read(2)
			if is_escaped != "\033[":
				raise Exception("Invalid response")

			if answer_ends_with:
				while resp[-len(answer_ends_with):] != answer_ends_with:
					resp += sys.stdin.read(1)
			else:
				while len(resp) != num_chars - 1:
					resp += sys.stdin.read(1)

		finally:
			termios.tcsetattr(sys.stdin, termios.TCSADRAIN, stored_attrs)

		if autosplit:
			resp = resp[:-len(answer_ends_with)]
			return resp.split(";")
		return resp
	# }}}
	def _is_compliant(self): # {{{
		"""
			Check if the terminal is compliant
		"""
		if not sys.stdin.isatty():
			return False
		if not ("term" in os.environ["TERM"] or "vt" in os.environ["TERM"]):
			return False
		try:
			response = (self._send_and_receive_esc_seq("[5n", answer_ends_with="n",
				autosplit=True) != ['0']) # 5n means status report, response 0 means everything ok
		except:
			return False
		return True
	# }}}
	def __init__(self): # {{{
		assert(self._is_compliant())
	# }}}

	def reset_settings(self):
		"""
			Reset terminal settings to default
		"""
		self._send_esc_seq("c")
	
	# Basic settings {{{
	def set_line_wrap(self, enable = True):
		"""
			Enable or disable line wrapping
		"""
		self._send_esc_seq("[7" + ("h" if enable else "l"))
	
	def set_charset(self, utf8 = True):
		"""
			Switch between UTF-8 and ISO-8859-1 mode
		"""
		self._send_esc_seq("%G" if utf8 else "%@")
	def toggle_insert_mode(self):
		"""
			Put terminal into insert mode
		"""
		self._send_esc_seq("[4h")
	def set_led(self, status):
		"""
			Set LED status.
				0 - disable all LEDs
				n - enable corresponding LED
		"""
		assert(type(status) == int)
		self._send_esc_seq("[%dq" % status)
		
	# }}}
	# Information {{{
	def query_terminal_size(self):
		"""
			Return terminal size as (rows, cols)
		"""
		size = array.array("B", [ 0, 0, 0, 0 ])
		assert(fcntl.ioctl(0, termios.TIOCGWINSZ, size, True) == 0)
		return size[0], size[2]
	def query_mouse(self):
		"""
			Wait for a mouse click and report it's position
			as (row, column)
		"""
		data = self._send_and_receive_esc_seq("[?9h", num_chars=5)
		self._send_esc_seq("[?1000l")
		place = array.array('B')
		place.fromstring(data[1:])
		return place[2] - 040, place[1] - 040
	# }}}
	# Cursor stuff {{{
	def cursor_get_position(self):
		"""
			Return cursor position as a (row, col) tuple
		"""
		return map(int, self._send_and_receive_esc_seq("[6n", answer_ends_with="R", autosplit=True))
	def cursor_set_position(self, row, col):
		"""
			Set cursor position to row, column as given
		"""
		assert(type(row) == int and type(col) == int)
		self._send_esc_seq("[%d;%dH" % (row, col))
	def cursor_set_col(self, col):
		"""
			Move cursor to the specified column
		"""
		assert(type(col) == int)
		self._send_esc_seq("[%d`" % col)
	def cursor_move_up(self, num_lines = 1):
		"""
			Move the cursor up by num_lines lines
		"""
		assert(type(num_lines) == int)
		self._send_esc_seq("[%dA" % num_lines)
	def cursor_move_down(self, num_lines = 1):
		"""
			Move the cursor down by num_lines lines
		"""
		assert(type(num_lines) == int)
		self._send_esc_seq("[%dB" % num_lines)
	def cursor_move_left(self, num_columns = 1):
		"""
			Move the cursor left by num_columns columns
		"""
		assert(type(num_columns) == int)
		self._send_esc_seq("[%dD" % num_columns)
	def cursor_move_right(self, num_columns = 1):
		"""
			Move the cursor right by num_columns columns
		"""
		assert(type(num_columns) == int)
		self._send_esc_seq("[%dC" % num_columns)
	def cursor_store_position(self):
		"""
			Store the cursor's position into an internal
			variable
		"""
		self._send_esc_seq("[s")
	def cursor_store_position_and_attributes(self):
		"""
			Store the cursor's position and attributes into an
			internal variable
		"""
		self._send_esc_seq("[7")
	def cursor_restore_position(self):
		"""
			Restore the cursor's position from an internal
			variable
		"""
		self._send_esc_seq("[u")
	def cursor_restore_position_and_attributes(self):
		"""
			Restore the cursor's position and attributes from an
			internal variable
		"""
		self._send_esc_seq("[8")
	# }}}
	# Scroll area {{{
	def scroll_area(self, top_row=0, bottom_row=0):
		"""
			Force the scroll area into the given rows. The difference
			between the lines must be at least 1. Skip the parameters
			to restore default behaviour
		"""
		assert(type(top_row) == int and type(bottom_row) == int)
		assert(bottom_row - top_row > 0 or (bottom_row == top_row and bottom_row == 0))
		self._send_esc_seq("[%d;%dr" % (top_row, bottom_row))
	def scroll_down(self, num_lines):
		"""
			Scroll the screen num_lines down
		"""
		assert(type(num_lines) == int)
		for i in range(num_lines):
			self._send_esc_seq("D")
	def scroll_up(self, num_lines):
		"""
			Scroll the screen num_lines up
		"""
		assert(type(num_lines) == int)
		for i in range(num_lines):
			self._send_esc_seq("M")
	# }}}
	# Tabstop {{{
	def tabstop_set(self, column):
		"""
			Set a tabstop in a specific column
		"""
		saved_position = self.cursor_get_position()
		sys.stdout.write("\r")
		self.cursor_move_right(column)
		self._send_esc_seq("H")
		apply(self.cursor_set_position, saved_position)

	def tabstop_clear(self, column):
		"""
			Clear a tabstop from a specific column
		"""
		saved_position = self.cursor_get_position()
		sys.stdout.write("\r" + " " * column)
		self._send_esc_seq("[g")
		apply(self.cursor_set_position, saved_position)
	def tabstop_clear_all(self):
		"""
			Clear all tabstops
		"""
		self._send_esc_seq("[3g")
	# }}}
	# Erasing the screen {{{
	def erase_to_eol(self):
		"""
			Erase to end of line
		"""
		self._send_esc_seq("[K")
	def erase_to_sol(self):
		"""
			Erase to start of line
		"""
		self._send_esc_seq("[1K")
	def erase_line(self):
		"""
			Erase the current line
		"""
		self._send_esc_seq("[2K")
	def erase_down(self):
		"""
			Erase the screen from the current line to end of screen
		"""
		self._send_esc_seq("[J")
	def erase_up(self):
		"""
			Erase the screen from the current line to top of screen
		"""
		self._send_esc_seq("[1J")
	def erase(self):
		"""
			Erase the screen
		"""
		self._send_esc_seq("[2J")
	# }}}
	# Color and fun stuff {{{
	ATTR_RESET	= 0
	ATTR_BRIGHT	= 1
	ATTR_BOLD	= 1
	ATTR_DIM	= 2
	ATTR_UNDERSCORE = 4
	ATTR_BLINK	= 5
	ATTR_REVERSE	= 7
	ATTR_INVERSE	= 7
	ATTR_HIDDEN	= 8
	ATTR_FG_BLACK	= 30
	ATTR_FG_RED	= 31
	ATTR_FG_GREEN	= 32
	ATTR_FG_YELLOW	= 33
	ATTR_FG_BLUE	= 34
	ATTR_FG_MAGENTA	= 35
	ATTR_FG_CYAN	= 36
	ATTR_FG_WHITE	= 37
	ATTR_BG_BLACK	= 40
	ATTR_BG_RED	= 41
	ATTR_BG_GREEN	= 42
	ATTR_BG_YELLOW	= 43
	ATTR_BG_BLUE	= 44
	ATTR_BG_MAGENTA	= 45
	ATTR_BG_CYAN	= 46
	ATTR_BG_WHITE	= 47

	def attrib_set(self, *attributes):
		"""
			Set terminal attributes, including color. See
			ATTR_* constants in this class.
		"""
		for attrib in attributes:
			assert(type(attrib) == int)
			assert(attrib in (range(9) + range(30, 38) + range(40, 48)))
		self._send_esc_seq("[%sm" % ";".join(map(str, attributes)))
	
	def attrib_reset(self):
		"""
			Reset attributes to default settings
		"""
		self._send_esc_seq(self.ATTR_RESET)
	
	def color_write(self, text):
		"""
			Write text to screen, replacing %{FOO} with ATTR_FOO.
			i.E. %{FG_BLUE} will write blue text
		"""
		for constant in filter(lambda s: s[0:5] == "ATTR_", dir(self)):
			text = text.replace("%%{%s}" % constant[5:], "\033[%sm" % str(getattr(self, constant)))
		self.attrib_set(self.ATTR_RESET)
		print text
		self.attrib_set(self.ATTR_RESET)
	# }}}
	# Xterm specific {{{
	def xterm_title(self, title):
		"""
			Set an X-terminal's title
		"""
		self._send_esc_seq("]2;%s\033\\" % title.replace("\033\\", ""))
	def xterm_icon_name(self, icon_name):
		"""
			Set an X-terminal's icon
		"""
		self._send_esc_seq("]1;%s\033\\" % icon_name.replace("\033\\", ""))
	def xterm_font(self, font):
		"""
			Set an X-terminal's font
		"""
		self._send_esc_seq("]50;%s\033\\" % font.replace("\033\\", ""))
	# }}}


if __name__ == "__main__": 
	# Code to use this from the command line {{{
	from optparse import OptionParser
	parser = OptionParser(usage="%prog <action> [arguments]", description="Change terminal behaviour")
	parser.add_option("-c", "--color", action="store", help="Display a colored string, parsing stuff like %{FG_RED}")
	parser.add_option("-s", "--store-cursor", action="store_true", help="Store cursor position")
	parser.add_option("-r", "--restore-cursor", action="store_true", help="Restore cursor position")
	parser.add_option("-p", "--place-cursor", action="store", help="Set cursor position to x, y")
	parser.add_option("-g", "--get-cursor", action="store_true", help="Output cursor position")
	parser.add_option("-e", "--erase-screen", action="store_true", help="Erase the screen")
	parser.add_option("-l", "--limit-scroll", action="store", help="Limit the scroll area from line1 to line2")
	parser.add_option("-t", "--terminal-size", action="store_true", help="Query terminal size")
	parser.add_option("-m", "--get-mouse", action="store_true", help="Wait for a mouse click and return position")
	parser.add_option("-x", "--xtitle", action="store", help="Set xterm title")
	(args, params) = parser.parse_args(sys.argv)

	def parse_pair(pair):
		try:
			values = pair.split(",")
			values = map(int, values)
			assert(len(values) == 2)
			assert(values[0] > 0 and values[1] > 0)
		except:
			parser.error("The parameter needs to have the format x,y")
		return values

	obj = Consolectrl()
	if args.xtitle:
		obj.xterm_title(args.xtitle)
	if args.erase_screen:
		obj.erase()
	if args.limit_scroll:
		l1, l2 = parse_pair(args.limit_scroll)
		obj.scroll_area(l1, l2)
	if args.get_cursor:
		print " ".join(map(str, obj.cursor_get_position()))
	if args.get_mouse:
		print " ".join(map(str, obj.query_mouse()))
	if args.terminal_size:
		print " ".join(map(str, obj.query_terminal_size()))
	if args.store_cursor:
		obj.cursor_store_position()
	if args.restore_cursor:
		obj.cursor_restore_position()
	if args.place_cursor:
		l1, l2 = parse_pair(args.place_cursor)
		obj.cursor_set_position(l1, l2)
	if args.color:
		obj.color_write(args.color)
	# }}}