#!/usr/bin/python

"""
pg_activity utility
version: 0.2.0 DEV
author: Julien Tachoires <julmon@gmail.com>
license: New BSD License

Copyright (c) 2012, Julien Tachoires
All rights reserved.

Redistribution and use in source and binary forms, with or without
modification, are permitted provided that the following conditions are met:
    * Redistributions of source code must retain the above copyright
      notice, this list of conditions and the following disclaimer.
    * Redistributions in binary form must reproduce the above copyright
      notice, this list of conditions and the following disclaimer in the
      documentation and/or other materials provided with the distribution.
    * Neither the name of pg_activity nor the
      names of its contributors may be used to endorse or promote products
      derived from this software without specific prior written permission.

THIS SOFTWARE IS PROVIDED BY THE COPYRIGHT HOLDERS AND CONTRIBUTORS "AS IS" AND
ANY EXPRESS OR IMPLIED WARRANTIES, INCLUDING, BUT NOT LIMITED TO, THE IMPLIED
WARRANTIES OF MERCHANTABILITY AND FITNESS FOR A PARTICULAR PURPOSE ARE
DISCLAIMED. IN NO EVENT SHALL JULIEN TACHOIRES BE LIABLE FOR ANY
DIRECT, INDIRECT, INCIDENTAL, SPECIAL, EXEMPLARY, OR CONSEQUENTIAL DAMAGES
(INCLUDING, BUT NOT LIMITED TO, PROCUREMENT OF SUBSTITUTE GOODS OR SERVICES;
LOSS OF USE, DATA, OR PROFITS; OR BUSINESS INTERRUPTION) HOWEVER CAUSED AND
ON ANY THEORY OF LIABILITY, WHETHER IN CONTRACT, STRICT LIABILITY, OR TORT
(INCLUDING NEGLIGENCE OR OTHERWISE) ARISING IN ANY WAY OUT OF THE USE OF THIS
SOFTWARE, EVEN IF ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.
"""

PGTOP_VERSION="0.2.0"


import os
import sys
if os.name != 'posix':
	sys.exit('FATAL: Platform not supported.')
import getpass
if getpass.getuser() != 'postgres':
	sys.exit('FATAL: Must be run as postgres.')
import signal
import time
import curses
import atexit
from datetime import datetime, timedelta
import psutil
import psycopg2
import psycopg2.extras
import re
from optparse import OptionParser, OptionGroup
import optparse
import socket

"""
Called at exit time.
Rollback to default values.
"""
def at_exit_curses():
	win.keypad(0)
	win.move(0,0)
	win.erase()
	curses.nocbreak()
	curses.echo()
	curses.curs_set(1)
	curses.endwin()

"""
Function called on a kill
"""
def signal_handler(signal, frame):
	at_exit_curses()
	print "FATAL: Killed."
	sys.exit(0)


# Curses initialization
curses.setupterm()
win = curses.initscr()
win.keypad(1)
curses.noecho()
curses.cbreak()
# deactivate cursor
curses.curs_set(0)
# use colors
curses.start_color()
curses.use_default_colors()
lineno = 0
curses.endwin()
win.scrollok(0)
lines = []

# define some color pairs
C_BLACK_GREEN = 1
C_CYAN = 2
C_RED = 3
C_GREEN = 4
C_YELLOW = 5
C_MAGENTA = 6
C_WHITE = 7
C_BLACK_CYAN = 8
C_RED_BLACK = 9
# Columns
PGTOP_FLAG_DATABASE = 1
PGTOP_FLAG_CLIENT = 2
PGTOP_FLAG_CPU = 4
PGTOP_FLAG_MEM = 8
PGTOP_FLAG_READ = 16
PGTOP_FLAG_WRITE = 32
PGTOP_FLAG_TIME = 64
PGTOP_FLAG_WAIT = 128
PGTOP_FLAG_NONE = None
PGTOP_FLAGS_COL = PGTOP_FLAG_DATABASE | PGTOP_FLAG_CPU | PGTOP_FLAG_MEM | PGTOP_FLAG_READ | PGTOP_FLAG_WRITE | PGTOP_FLAG_TIME | PGTOP_FLAG_WAIT

PGTOP_COLS = {
	'pid'		: {'n': 1, 'name': 'PID', 'template_h': '%-6s ', 'flag': PGTOP_FLAG_NONE, 'mandatory': True},
	'database'	: {'n': 2, 'name': 'DATABASE', 'template_h': '%-16s ', 'flag': PGTOP_FLAG_DATABASE, 'mandatory': False},
	'client'	: {'n': 3, 'name': 'CLIENT', 'template_h': '%16s ', 'flag': PGTOP_FLAG_CLIENT, 'mandatory': False},
	'cpu'		: {'n': 4, 'name': 'CPU%', 'template_h': '%5s ', 'flag': PGTOP_FLAG_CPU, 'mandatory': False},
	'mem'		: {'n': 5, 'name': 'MEM%', 'template_h': '%4s ', 'flag': PGTOP_FLAG_MEM, 'mandatory': False},
	'read'		: {'n': 6, 'name': 'READ/s', 'template_h': '%8s ', 'flag': PGTOP_FLAG_READ, 'mandatory': False},
	'write'		: {'n': 7, 'name': 'WRITE/s', 'template_h': '%8s ', 'flag': PGTOP_FLAG_WRITE, 'mandatory': False},
	'time'		: {'n': 8, 'name': 'TIME+', 'template_h': '%9s ', 'flag': PGTOP_FLAG_TIME, 'mandatory': False},
	'wait'		: {'n': 9, 'name': 'W', 'template_h': '%2s ', 'flag': PGTOP_FLAG_WAIT, 'mandatory': False},
	'query'		: {'n': 10, 'name': 'Query', 'template_h': ' %2s', 'flag': PGTOP_FLAG_NONE, 'mandatory': True},
}

PGTOP_LINE_COLORS = {
	'pid'			: {'default': curses.color_pair(C_CYAN), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'database'		: {'default': curses.A_UNDERLINE | curses.color_pair(0), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'client'		: {'default': curses.color_pair(C_CYAN), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'cpu'			: {'default': curses.color_pair(0), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'mem'			: {'default': curses.color_pair(0), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'read'			: {'default': curses.color_pair(0), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'write'			: {'default': curses.color_pair(0), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'time_red'		: {'default': curses.color_pair(C_RED), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'time_yellow'	: {'default': curses.color_pair(C_YELLOW), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'time_green'	: {'default': curses.color_pair(C_GREEN), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'wait_green'	: {'default': curses.color_pair(C_GREEN) | curses.A_BOLD, 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'wait_red'		: {'default': curses.color_pair(C_RED) | curses.A_BOLD, 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)},
	'query'			: {'default': curses.color_pair(0), 'cursor': curses.color_pair(C_CYAN)|curses.A_REVERSE, 'lock': curses.color_pair(C_RED)}
}

# display query mode
PGTOP_VERBOSE_MODE = 2
PGTOP_TRUNCATE = 1
PGTOP_WRAP_NOINDENT = 2
PGTOP_WRAP = 3
# max IOPs
PGTOP_MAX_IOPS = 0
# color ?
PGTOP_COLOR = True
# sort
PGTOP_SORT = 't'

"""
Replace colors by white.
"""
def set_nocolor():
	global PGTOP_COLOR
	PGTOP_COLOR = False
	curses.init_pair(C_BLACK_GREEN, curses.COLOR_BLACK, curses.COLOR_WHITE)
	curses.init_pair(C_CYAN, curses.COLOR_WHITE, -1)
	curses.init_pair(C_RED, curses.COLOR_WHITE, -1)
	curses.init_pair(C_RED_BLACK, curses.COLOR_WHITE, curses.COLOR_BLACK)
	curses.init_pair(C_GREEN, curses.COLOR_WHITE, -1)
	curses.init_pair(C_YELLOW, curses.COLOR_WHITE, -1)
	curses.init_pair(C_MAGENTA, curses.COLOR_WHITE, -1)
	curses.init_pair(C_WHITE, curses.COLOR_WHITE, -1)
	curses.init_pair(C_BLACK_CYAN, curses.COLOR_WHITE, -1)

"""
Set colors.
"""
def set_color():
	global PGTOP_COLOR
	PGTOP_COLOR = True
	curses.init_pair(C_BLACK_GREEN, curses.COLOR_BLACK, curses.COLOR_GREEN)
	curses.init_pair(C_CYAN, curses.COLOR_CYAN, -1)
	curses.init_pair(C_RED, curses.COLOR_RED, -1)
	curses.init_pair(C_RED_BLACK, curses.COLOR_RED, curses.COLOR_BLACK)
	curses.init_pair(C_GREEN, curses.COLOR_GREEN, -1)
	curses.init_pair(C_YELLOW, curses.COLOR_YELLOW, -1)
	curses.init_pair(C_MAGENTA, curses.COLOR_MAGENTA, -1)
	curses.init_pair(C_WHITE, curses.COLOR_WHITE, -1)
	curses.init_pair(C_BLACK_CYAN, curses.COLOR_BLACK, curses.COLOR_CYAN)

"""
Strip and replace some special characters.
"""
def clean_str(string):
	msg = str(string)
	msg = msg.replace("\n", " ")
	msg = re.sub(r'\s+', r' ', msg)
	msg = msg.replace("FATAL:", "")
	msg = re.sub(r'^\s', r'', msg)
	msg = re.sub(r'\s$', r'', msg)
	return msg

"""
Connect to a PostgreSQL server and returns
cursor & connector.
"""
def pg_connect(host = 'localhost', port = 5432, user = 'postgres', password = None):
	conn = psycopg2.connect(
			database = 'postgres',
			host = host,
			port = port,
			user = user,
			password = str(password),
			connection_factory=psycopg2.extras.DictConnection	
		)
	conn.set_isolation_level(0)
	return conn

"""
Get PostgreSQL server version
"""
def pg_get_version(conn):
	query = """
	SELECT version() AS pg_version
	"""
	cur = conn.cursor()
	res = cur.execute(query)
	ret = cur.fetchone()
	return ret['pg_version']

"""
Cancel a backend
"""
def pg_cancel_backend(conn, pid):
	query = """
	SELECT pg_cancel_backend(%s) AS canceled
	"""
	cur = conn.cursor()
	res = cur.execute(query, (pid,))
	ret = cur.fetchone()
	return ret['canceled']

"""
Get PostgreSQL short & numeric version from
a string (SELECT version())
"""
def pg_get_num_version(text_version):
	res = re.match(r'^PostgreSQL ([0-9]+)\.([0-9]+)\.([0-9]+)', text_version)
	if res is not None:
		r = res.group(1)
		if int(res.group(2)) < 10:
			r += '0'
		r += res.group(2)
		if int(res.group(3)) < 10:
			r += '0'
		r += res.group(3)
		return (res.group(0), int(r))
	return None

"""
Get activity from pg_stat_activity view
"""
def pg_get_activities(conn, pg_num_version):
	if pg_num_version >= 90200:
		# PostgreSQL 9.2.0 and more
		query = """
		SELECT
			pg_stat_activity.pid AS pid,
			CASE WHEN LENGTH(pg_stat_activity.datname) > 16 THEN SUBSTRING(pg_stat_activity.datname FROM 0 FOR 6)||'...'||SUBSTRING(pg_stat_activity.datname FROM '........$') ELSE pg_stat_activity.datname END AS database,
			pg_stat_activity.client_addr AS client,
			EXTRACT(epoch FROM (NOW() - pg_stat_activity.query_start)) AS duration,
			pg_stat_activity.waiting AS wait,
			pg_stat_activity.usename AS user,
			pg_stat_activity.query AS query
		FROM
			pg_stat_activity
		WHERE
			state <> 'idle'
		ORDER BY
			EXTRACT(epoch FROM (NOW() - pg_stat_activity.query_start)) DESC
		"""
	elif pg_num_version < 90200:
		# PostgreSQL prior to 9.2.0
		query = """
		SELECT
			pg_stat_activity.procpid AS pid,
			CASE WHEN LENGTH(pg_stat_activity.datname) > 16 THEN SUBSTRING(pg_stat_activity.datname FROM 0 FOR 6)||'...'||SUBSTRING(pg_stat_activity.datname FROM '........$') ELSE pg_stat_activity.datname END AS database,
			pg_stat_activity.client_addr AS client,
			EXTRACT(epoch FROM (NOW() - pg_stat_activity.query_start)) AS duration,
			pg_stat_activity.waiting AS wait,
			pg_stat_activity.usename AS user,
			pg_stat_activity.current_query AS query
		FROM
			pg_stat_activity
		WHERE
			current_query <> '<IDLE>'
		ORDER BY
			EXTRACT(epoch FROM (NOW() - pg_stat_activity.query_start)) DESC
		"""
	cur = conn.cursor()
	res = cur.execute(query)
	ret = cur.fetchall()
	return ret

"""
Get system informations (CPU, memory, IO read & write)
for each process PID -- retreived previously from pg_stat_activity -- with
psutil module.
"""
def sys_get_proc(queries):
	process = {}
	for sq in queries:
		try:
			p = psutil.Process(sq['pid'])
			p._meminfo = p.get_memory_info()
			p._io_counters = p.get_io_counters()
			p._io_time = time.time()
			p._mem_percent = p.get_memory_percent()
			p._cpu_percent = p.get_cpu_percent(interval=0)
			p._cpu_times = p.get_cpu_times()
			p._pg_user = sq['user']
			p._pg_database = sq['database']
			p._pg_client = sq['client']
			p._pg_duration = sq['duration']
			p._pg_wait = sq['wait']
			p._pg_query = clean_str(sq['query'])
			p._io_read_delta = 0
			p._io_write_delta = 0
			process[sq['pid']] = p
		except psutil.error.NoSuchProcess as e:
			pass
		except psutil.error.AccessDenied as e:
			pass
	return process

"""
PAUSE mode
"""
def pause():
	(y,x) = win.getmaxyx()
	msg = "PAUSE"
	line = ""
	line += " " * (int(x/2) - len(msg))
	line += msg
	line += " " * (x - len(line) - 0)
	win.addstr(4, 0, line, curses.color_pair(C_RED_BLACK)|curses.A_REVERSE|curses.A_BOLD)
	current_pos = -1
	while 1:
		try:
			k = win.getch()
		except KeyboardInterrupt as e:
			raise e
		if k == ord('q'):
			curses.endwin()
			exit()
		if k == ord(' '):
			curses.flushinp()
			return 0
		curses.flushinp()

"""
Display interactive mode menu bar
"""
def help_key_interactive(y):
	colno = print_string((y - 1), 0, "<k>", curses.color_pair(0))
	colno += print_string((y - 1), colno, "Backend cancel    ", curses.color_pair(C_CYAN)|curses.A_REVERSE)
	colno += print_string((y - 1), colno, "<Space>", curses.color_pair(0))
	colno += print_string((y - 1), colno, "Back to activity    ", curses.color_pair(C_CYAN)|curses.A_REVERSE)
	colno += print_string((y - 1), colno, "<q>", curses.color_pair(0))
	colno += print_string((y - 1), colno, "Quit    ", curses.color_pair(C_CYAN)|curses.A_REVERSE)
	colno += print_string((y - 1), colno, add_blank(" "), curses.color_pair(C_CYAN)|curses.A_REVERSE)

"""
Ask for canceling a backend
"""
def ask_cancel_backend(y, pid, connector):
	colno = print_string((y - 1), 0, "Cancel the backend with PID %s ? <Y/N>" % (str(pid),), curses.color_pair(0))
	colno += print_string((y - 1), colno, add_blank(" "), curses.color_pair(C_CYAN)|curses.A_REVERSE)
	while 1:
		try:
			k = win.getch()
		except KeyboardInterrupt as e:
			raise e
		# quit
		if k == ord('q'):
			curses.endwin()
			exit()
		# yes
		if k == ord('y') or k == ord('Y'):
			res = pg_cancel_backend(connector, str(pid))
			time.sleep(1)
			return 1
		# no
		if k == ord('n') or k == ord('N') or k == ord(' '):
			return 0

"""
Interactive mode trigged on KEY_UP or KEY_DOWN key press
If no key hit during 3 seconds, exit this mode
"""
def interactive(process, flag, indent, connector):
	global lines
	global PGTOP_VERBOSE_MODE
	# Force truncated display
	old_verbose_mode = PGTOP_VERBOSE_MODE
	PGTOP_VERBOSE_MODE = PGTOP_TRUNCATE

	# Refresh lines with this verbose mode	
	(y,x) = win.getmaxyx()
	scroll_window(process, flag, indent, 0)
	
	help_key_interactive(y)

	current_pos = 0
	# refresh_line(process[current_pos], flag, indent, 'cursor', lines[current_pos][1])
	offset = 0
	refresh_line(process[current_pos], flag, indent, 'cursor', lines[current_pos] - offset)
	t_active = time.time()
	# raise Exception(str(lines))
	while 1:
		try:
			k = win.getch()
		except KeyboardInterrupt as e:
			raise e
		# quit
		if k == ord('q'):
			curses.endwin()
			exit()
		# cancel the backend attached to this PID
		if k == ord('k'):
			res = ask_cancel_backend(y, process[current_pos].pid, connector)
			PGTOP_VERBOSE_MODE = old_verbose_mode
			curses.flushinp()
			return 0
		# Move cursor
		if k == curses.KEY_DOWN or k == curses.KEY_UP:
			t_active = time.time()
			if k == curses.KEY_UP and current_pos > 0:
				if (lines[current_pos] - offset) < 7:
					offset -= 1
					scroll_window(process, flag, indent, offset)
					help_key_interactive(y)

				if current_pos < len(process):
					refresh_line(process[current_pos], flag, indent, 'default', lines[current_pos] - offset)
				current_pos -= 1
			if k == curses.KEY_DOWN and current_pos < (len(process) - 1):
				if (lines[current_pos] - offset) >= (y - 2):
					offset += 1
					scroll_window(process, flag, indent, offset)
					help_key_interactive(y)

				if current_pos >= 0:
					refresh_line(process[current_pos], flag, indent, 'default', lines[current_pos] - offset)
				current_pos += 1
			refresh_line(process[current_pos], flag, indent, 'cursor', lines[current_pos] - offset)
		# quit interactive mode
		if k == ord(' '):
			PGTOP_VERBOSE_MODE = old_verbose_mode
			curses.flushinp()
			return 0
		curses.flushinp()
		if time.time() - t_active > 3:
			PGTOP_VERBOSE_MODE = old_verbose_mode
			return 0

"""
Poll activities.
"""
def poll(interval, connector, flag, indent, process = None, pg_num_version = None, disp_proc = None):
	global PGTOP_VERBOSE_MODE
	global PGTOP_COLOR
	global PGTOP_SORT

	# Keyboard interactions
	win.timeout(1000 * interval)
	try:
		k = win.getch()
	except KeyboardInterrupt as e:
		raise e
	if k == ord('q'):
		curses.endwin()
		exit()
	# PAUSE mode
	if k == ord(' '):
		pause()
	# interactive mode
	if k == curses.KEY_DOWN or k == curses.KEY_UP:
		interactive(disp_proc, flag, indent, connector)
	# change verbosity
	if k == ord('v'):
		PGTOP_VERBOSE_MODE += 1
		if PGTOP_VERBOSE_MODE > 3:
			PGTOP_VERBOSE_MODE = 1
	# turnoff/on colors
	if k == ord('C'):
		if PGTOP_COLOR is True:
			set_nocolor()
		else:
			set_color()
	# sorts
	if k == ord('c') and (flag & PGTOP_FLAG_CPU):
		PGTOP_SORT = 'c'
	if k == ord('m') and (flag & PGTOP_FLAG_MEM):
		PGTOP_SORT = 'm'
	if k == ord('r') and (flag & PGTOP_FLAG_READ):
		PGTOP_SORT = 'r'
	if k == ord('w') and (flag & PGTOP_FLAG_WRITE):
		PGTOP_SORT = 'w'
	if k == ord('t'):
		PGTOP_SORT = 't'
		
	curses.flushinp()
	
	# poll postgresql activity
	queries =  pg_get_activities(connector, pg_num_version)
	# get resource usage for each process
	new_procs = sys_get_proc(queries)

	procs_status = {}
	procs = []
	for pid,p in new_procs.items():
		try:
			if process.has_key(pid):
				n_pg_duration = p._pg_duration
				n_pg_query = p._pg_query
				n_pg_client = p._pg_client
				n_io_counters = p.get_io_counters()
				n_io_time = time.time()
				p = process[pid]
				p._pg_duration = n_pg_duration
				p._pg_query = n_pg_query
				p._pg_client = n_pg_client
				p._io_read_delta = (n_io_counters.read_bytes - p._io_counters.read_bytes)/(n_io_time - p._io_time)
				p._io_write_delta = (n_io_counters.write_bytes - p._io_counters.write_bytes)/(n_io_time - p._io_time)
				p._io_counters = n_io_counters
				p._io_time = n_io_time
			else:
				p._io_counters = p.get_io_counters()
			
			p._mempercent = p.get_memory_percent()
			p._cpu_percent = p.get_cpu_percent(interval=0)
			p._name = p.name
			new_procs[pid] = p
			procs.append(p)
			try:
				procs_status[str(p.status)] += 1
			except KeyError:
				procs_status[str(p.status)] = 1
			except Exception as e:
				raise e

		except psutil.NoSuchProcess:
			pass
		except psutil.error.AccessDenied as e:
			pass
		except Exception as e:
			raise e

	# return processes sorted by query duration
	if PGTOP_SORT == 't':
		# TIME
		disp_procs = sorted(procs, key=lambda p: p._pg_duration, reverse=True)
	elif PGTOP_SORT == 'c':
		# CPU
		disp_procs = sorted(procs, key=lambda p: p._cpu_percent, reverse=True)
	elif PGTOP_SORT == 'm':
		# MEM
		disp_procs = sorted(procs, key=lambda p: p._mem_percent, reverse=True)
	elif PGTOP_SORT == 'r':
		# READ
		disp_procs = sorted(procs, key=lambda p: p._io_read_delta, reverse=True)
	elif PGTOP_SORT == 'w':
		# WRITE
		disp_procs = sorted(procs, key=lambda p: p._io_write_delta, reverse=True)
	else:
		disp_procs = sorted(procs, key=lambda p: p._pg_duration, reverse=True)
	
	return (disp_procs, procs_status, new_procs)

"""
Simple curses wrapper for word printing.
"""
def print_word(word, color = 0):
	win.addstr(word, color)

"""
Print a string at position (lineno, colno) and returns its length.
"""
def print_string(lineno, colno, word, color = 0):
	try:
		win.addstr(lineno, colno, word, color)
	except curses.error:
		pass
	return len(word)


"""
Complete string with white spaces from the end of string to the end of line.
"""
def add_blank(line, offset = 0): 
	line += " " * (win.getmaxyx()[1] - (len(line) + offset))
	return line

"""
Convert a size into a human readable format.
"""
def bytes2human(n):
	symbols = ('K', 'M', 'G', 'T', 'P', 'E', 'Z', 'Y')
	prefix = {}
	for i, s in enumerate(symbols):
		prefix[s] = 1 << (i+1)*10
	for s in reversed(symbols):
		if n >= prefix[s]:
			value = '%.2f' % float(float(n) / float(prefix[s]))
			return '%s%s' % (value, s)
	return "%.2fB" % n

"""
Returns identation for Query column.
"""
def get_indent(flag):
	indent = ''
	r = [0] * 11
	for key, val in PGTOP_COLS.items():
		if val['mandatory'] or (not val['mandatory'] and val['flag'] & flag):
			r[int(val['n'])] = val
	for val in r:
		if val is not 0:
			if val['name'] is not 'Query':
				indent += val['template_h'] % ' '
	return indent

"""
Print columns headers
"""
def print_cols_header(flag):
	global PGTOP_SORT
	global lineno
	line = ''
	disp = ''
	x = 0
	r = [0] * 11
	for key, val in PGTOP_COLS.items():
		if val['mandatory'] or (not val['mandatory'] and val['flag'] & flag):
			r[int(val['n'])] = val
	for val in r:
		if val is not 0:
			disp = val['template_h'] % val['name']
			if (
				(val['name'] == 'CPU%' and PGTOP_SORT == 'c') or
				(val['name'] == 'MEM%' and PGTOP_SORT == 'm') or
				(val['name'] == 'READ/s' and PGTOP_SORT == 'r') or
				(val['name'] == 'WRITE/s' and PGTOP_SORT == 'w') or
				(val['name'] == 'TIME+' and PGTOP_SORT == 't')):
				color_highlight = curses.color_pair(C_BLACK_CYAN)
			else:
				color_highlight = curses.color_pair(C_BLACK_GREEN)
			if val['name'] == 'Query':
				disp += " " * (win.getmaxyx()[1] - (len(line) + len(disp)))
			line += disp
			win.addstr(lineno, x, disp, color_highlight)
			x += len(disp)
	
	lineno += 1

"""
Print window header
"""
def print_header(procs_status, pg_version, hostname, user, host, port, io):
	global PGTOP_MAX_IOPS
	global lineno
	
	lineno = 0
	colno = 0
	version = " %s" % (pg_version)
	colno = print_string(lineno, colno, version)
	colno += print_string(lineno, colno, " - ")
	colno += print_string(lineno, colno, hostname, curses.A_BOLD)
	colno += print_string(lineno, colno, " - ")
	colno += print_string(lineno, colno, user, curses.color_pair(C_CYAN))
	colno += print_string(lineno, colno, "@")
	colno += print_string(lineno, colno, host, curses.color_pair(C_CYAN))
	colno += print_string(lineno, colno, ":")
	colno += print_string(lineno, colno, port, curses.color_pair(C_CYAN))

	try:
		# psutil >= 0.6.0
		phymem = psutil.virtual_memory()
		buffers = psutil.virtual_memory().buffers
		cached = psutil.virtual_memory().cached
		vmem = psutil.swap_memory()
	except AttributeError:
		# psutil > 0.4.0 and < 0.6.0
		phymem = psutil.phymem_usage()
		buffers = getattr(psutil, 'phymem_buffers', lambda: 0)()
		cached = getattr(psutil, 'cached_phymem', lambda: 0)()
		vmem = psutil.virtmem_usage()

	used = phymem.total - (phymem.free + buffers + cached)
	lineno += 1
	line = "  Mem.: %5s%% %9s/%s " % (phymem.percent, str(int(used / 1024 / 1024)) + "M", str(int(phymem.total / 1024 / 1024)) + "M")
	colno_io = print_string(lineno, 0, line)
	
	if (int(io['read_count'])+int(io['write_count'])) > PGTOP_MAX_IOPS:
		PGTOP_MAX_IOPS = (int(io['read_count'])+int(io['write_count']))
	
	line_io = " | IO Disks Max : %s IOPs" % (PGTOP_MAX_IOPS,)
	colno = print_string(lineno, colno_io, line_io)

	# swap usage
	line = "  Swap: %5s%% %9s/%s\n" % (vmem.percent, str(int(vmem.used / 1024 / 1024)) + "M", str(int(vmem.total / 1024 / 1024)) + "M")
	lineno += 1
	colno = print_string(lineno, 0, line)
	line_io = " | Read : %8s/s - %6s IOPs" % (bytes2human(io['read_bytes']), int(io['read_count']),)
	colno = print_string(lineno, colno_io, line_io)

	# load average, uptime
	av1, av2, av3 = os.getloadavg()
	line = "  Load:   %.2f %.2f %.2f" % (av1, av2, av3)
	lineno += 1
	colno = print_string(lineno, 0, line)
	line_io = " | Write: %8s/s - %6s IOPs" % (bytes2human(io['write_bytes']), int(io['write_count']),)
	colno = print_string(lineno, colno_io, line_io)

"""
Refresh the window
"""
def refresh_window(procs, procs_status, extras, flag, indent, io):
	global lineno
	global lines

	lines = []
	
	(pg_version, hostname, user, host, port) = extras
	win.erase()
	print_header(procs_status, pg_version, hostname, user, host, port, io)
	lineno += 2
	line_trunc = lineno
	print_cols_header(flag)
	colno = 0
	for p in procs:
		try:
			refresh_line(p, flag, indent, 'default')
			line_trunc += 1
			lines.append(line_trunc)
		except curses.error as e:
			break
		
		for l in range(lineno, (win.getmaxyx()[0]-1)):
			print_string(l, 0, add_blank(' '))

"""
Scroll the window
"""
def scroll_window(procs, flag, indent, offset = 0):
	global lineno
	lineno = 6
	pos = 0
	for p in procs:
		if pos >= offset:
			try:
				refresh_line(p, flag, indent, 'default')
			except curses.error as e:
				break
		
			for l in range(lineno, (win.getmaxyx()[0]-1)):
				print_string(l, 0, add_blank(' '))
		pos += 1

"""
Refresh a line
"""
def refresh_line(p, flag, indent, typecolor = 'default', line = None):
	global lineno
	if line is not None:
		l_lineno = line
	else:
		l_lineno = lineno

	(maxy, maxx) = win.getmaxyx()
	colno = 0
	if l_lineno > maxy:
		return
	
	colno += print_string(l_lineno, colno, "%-6s " % (p.pid,), PGTOP_LINE_COLORS['pid'][typecolor])
	if flag & PGTOP_FLAG_DATABASE:
		colno += print_string(l_lineno, colno, "%-16s " % (p._pg_database[:16],), PGTOP_LINE_COLORS['database'][typecolor])
 	if flag & PGTOP_FLAG_CLIENT:
		colno += print_string(l_lineno, colno, "%16s " % (str(p._pg_client)[:16],), PGTOP_LINE_COLORS['client'][typecolor])
	if flag & PGTOP_FLAG_CPU:
		colno += print_string(l_lineno, colno, "%5s " % (p._cpu_percent,), PGTOP_LINE_COLORS['cpu'][typecolor])
	if flag & PGTOP_FLAG_MEM:
		colno += print_string(l_lineno, colno, "%4s " % (round(p._mem_percent, 1),), PGTOP_LINE_COLORS['mem'][typecolor])
	if flag & PGTOP_FLAG_READ:
		colno += print_string(l_lineno, colno, "%8s " % (bytes2human(p._io_read_delta),), PGTOP_LINE_COLORS['read'][typecolor])
	if flag & PGTOP_FLAG_WRITE:
		colno += print_string(l_lineno, colno, "%8s " % (bytes2human(p._io_write_delta),), PGTOP_LINE_COLORS['write'][typecolor])
	if flag & PGTOP_FLAG_TIME:
		if p._pg_duration < 0:
			p._pg_duration = 0
		if p._pg_duration >= 1:
			ctime = timedelta(seconds=float(p._pg_duration))
			mic = '%.6d' % (ctime.microseconds)
			ctime = "%s:%s.%s" % (str((ctime.seconds // 60)).zfill(2), str((ctime.seconds % 60)).zfill(2), str(mic)[:2])
		if p._pg_duration < 1:
			colno += print_string(l_lineno, colno, " %.6f " % (p._pg_duration,), PGTOP_LINE_COLORS['time_green'][typecolor])
		elif p._pg_duration >= 1 and p._pg_duration < 3:
			colno += print_string(l_lineno, colno, "%9s " % (ctime,), PGTOP_LINE_COLORS['time_yellow'][typecolor])
		else:
			colno += print_string(l_lineno, colno, "%9s " % (ctime,), PGTOP_LINE_COLORS['time_red'][typecolor])
	if flag & PGTOP_FLAG_WAIT:
		if p._pg_wait:
			colno += print_string(l_lineno, colno, "%2s " % ('Y',), PGTOP_LINE_COLORS['wait_red'][typecolor])
		else:
			colno += print_string(l_lineno, colno, "%2s " % ('N',), PGTOP_LINE_COLORS['wait_green'][typecolor])

	dif = maxx - len(indent) - 1
	if PGTOP_VERBOSE_MODE == PGTOP_TRUNCATE:
		query = p._pg_query[:dif]
		colno += print_string(l_lineno, colno, " %s" % (add_blank(query,len(indent)+1),), PGTOP_LINE_COLORS['query'][typecolor])
	elif PGTOP_VERBOSE_MODE == PGTOP_WRAP or  PGTOP_VERBOSE_MODE == PGTOP_WRAP_NOINDENT:
		query = p._pg_query
		query_wrote = ''
		offset = 0
		if len(query) > dif:
			query_part = query[offset:dif]
			print_string(l_lineno, colno, " %s" % (add_blank(query_part,len(indent)+1),), PGTOP_LINE_COLORS['query'][typecolor])
			query_wrote += query_part
			offset = len(query_wrote)
			if PGTOP_VERBOSE_MODE == PGTOP_WRAP_NOINDENT:
				dif = maxx
				p_indent = ''
			else:
				p_indent = indent
			while (len(query) - offset > 0):	
				query_part = query[offset:(dif+offset)]
				l_lineno += 1
				lineno += 1
				print_string(l_lineno, 0, "%s" % (add_blank(p_indent + ' '+ query_part, len(indent)+1)), PGTOP_LINE_COLORS['query'][typecolor])
				query_wrote += query_part
				offset = len(query_wrote)
		else:
			colno += print_string(l_lineno, colno, " %s" % (add_blank(query, len(indent)),), PGTOP_LINE_COLORS['query'][typecolor])
	lineno += 1

"""
main
"""
def main():
	try:
		parser = OptionParser(add_help_option=False, version="%prog "+PGTOP_VERSION)
		parser.add_option("-U", "--username", dest="username", default="postgres", help="Database user name (default: \"postgres\").", metavar="USERNAME")
		parser.add_option("-p", "--port", dest="port", default="5432", help="Database server port (default: \"5432\").", metavar="PORT")
		parser.add_option("-h", "--host", dest="host", help="Database server host or socket directory (default: \"local socket\").", metavar="HOSTNAME", default="localhost")
		parser.add_option("-C", "--no-color", dest="nocolor", action="store_true", help="Disable color usage.", default="false")
		group = OptionGroup(parser, "Display Options, you can exclude some columns by using them ")
		group.add_option("--no-database", dest="nodb", action="store_true", help="Disable DATABASE.", default="false")
		group.add_option("--no-client", dest="noclient", action="store_true", help="Disable CLIENT.", default="false")
		group.add_option("--no-cpu", dest="nocpu", action="store_true", help="Disable CPU%.", default="false")
		group.add_option("--no-mem", dest="nomem", action="store_true", help="Disable MEM%.", default="false")
		group.add_option("--no-read", dest="noread", action="store_true", help="Disable READ/s.", default="false")
		group.add_option("--no-write", dest="nowrite", action="store_true", help="Disable WRITE/s.", default="false")
		group.add_option("--no-time", dest="notime", action="store_true", help="Disable TIME+.", default="false")
		group.add_option("--no-wait", dest="nowait", action="store_true", help="Disable W.", default="false")
		parser.add_option_group(group)
		parser.add_option("--help", dest="help", action="store_true", help="Show this help message and exit.", default="false")
    
		(options, args) = parser.parse_args()
		if options.help is True:
			at_exit_curses()
			print(parser.format_help().strip())
			sys.exit(0)
		
		password = os.environ.get('PGPASSWORD')
		
		try:
			conn = pg_connect(host = options.host, port = options.port, user = options.username, password = password)
		except psycopg2.Error as e:
			at_exit_curses()
			sys.exit("FATAL: %s" % (clean_str(str(e),)))
		pg_version = pg_get_version(conn)
		(pg_short_version, pg_num_version) = pg_get_num_version(pg_version)
		# top part
		interval = 0
		queries =  pg_get_activities(conn, pg_num_version)
		procs = sys_get_proc(queries)
		hostname = socket.gethostname()
		has_color = True
		# color ?
		if options.nocolor == True:
			set_nocolor()
			has_color = False
		else:
			set_color()	
		# default flag
		flag = PGTOP_FLAG_DATABASE | PGTOP_FLAG_CLIENT | PGTOP_FLAG_CPU | PGTOP_FLAG_MEM | PGTOP_FLAG_READ | PGTOP_FLAG_WRITE | PGTOP_FLAG_TIME | PGTOP_FLAG_WAIT
		if options.nodb is True:
			flag -= PGTOP_FLAG_DATABASE
		if options.nocpu is True:
			flag -= PGTOP_FLAG_CPU
		if options.noclient is True:
			flag -= PGTOP_FLAG_CLIENT
		if options.nomem is True:
			flag -= PGTOP_FLAG_MEM
		if options.noread is True:
			flag -= PGTOP_FLAG_READ
		if options.nowrite is True:
			flag -= PGTOP_FLAG_WRITE
		if options.notime is True:
			flag -= PGTOP_FLAG_TIME
		if options.nowait is True:
			flag -= PGTOP_FLAG_WAIT

		# indentation
		indent = get_indent(flag)
		# main loop
		disks_io = None
		disp_procs = None
		while 1:
			if len(indent) > (win.getmaxyx()[1] - 2):
				raise Exception('Window\'s too small, %s chars are required.' % (len(indent),))
			
			# poll process
			(disp_procs, procs_status, new_procs) = poll(interval, conn, flag, indent, procs, pg_num_version, disp_procs)
			new_disks_io = {'io':psutil.disk_io_counters(), 'time':time.time()}
			delta_disks_io = {'read_count': 0, 'write_count': 0, 'read_bytes': 0.0, 'write_bytes': 0.0}
			if disks_io is not None:
				delta_disks_io['read_count'] = (new_disks_io['io'].read_count - disks_io['io'].read_count)/(new_disks_io['time'] - disks_io['time'])
				delta_disks_io['write_count'] = (new_disks_io['io'].write_count - disks_io['io'].write_count)/(new_disks_io['time'] - disks_io['time'])
				delta_disks_io['read_bytes'] = (new_disks_io['io'].read_bytes - disks_io['io'].read_bytes)/(new_disks_io['time'] - disks_io['time'])
				delta_disks_io['write_bytes'] = (new_disks_io['io'].write_bytes - disks_io['io'].write_bytes)/(new_disks_io['time'] - disks_io['time'])
			disks_io = new_disks_io
	
			procs = new_procs
			# refresh the winodw
			refresh_window(disp_procs, procs_status, (pg_short_version, hostname, options.username, options.host, options.port), flag, indent, delta_disks_io)
			interval = 1

	except psutil.error.AccessDenied as e:
		at_exit_curses()
		sys.exit("FATAL: Can't acces IO informations for process %s" % (str(e),))
	except curses.error as e:
		at_exit_curses()
		sys.exit("FATAL: %s" % (str(e),))
	except KeyboardInterrupt as e:
		at_exit_curses()
		sys.exit(0)
	except Exception as e:
		at_exit_curses()
		# DEBUG
		# import traceback
		# exc_type, exc_value, exc_traceback = sys.exc_info()
		# print repr(traceback.format_tb(exc_traceback))
		# traceback.print_exception(exc_type, exc_value, exc_traceback, limit=2, file=sys.stdout)
		sys.exit("FATAL: %s" % (str(e),))

"""
Call the main function
"""
if __name__ == '__main__':
	signal.signal(signal.SIGTERM, signal_handler)
	main()
