#!/usr/bin/env python

"""
pg_activity
author: Julien Tachoires <julmon@gmail.com>
license: PostgreSQL License

Copyright (c) 2012 - 2015, Julien Tachoires

Permission to use, copy, modify, and distribute this software and its
documentation for any purpose, without fee, and without a written
agreement is hereby granted, provided that the above copyright notice
and this paragraph and the following two paragraphs appear in all copies.

IN NO EVENT SHALL JULIEN TACHOIRES BE LIABLE TO ANY PARTY FOR DIRECT,
INDIRECT, SPECIAL, INCIDENTAL, OR CONSEQUENTIAL DAMAGES, INCLUDING LOST
PROFITS, ARISING OUT OF THE USE OF THIS SOFTWARE AND ITS DOCUMENTATION,
EVEN IF JULIEN TACHOIRES HAS BEEN ADVISED OF THE POSSIBILITY OF SUCH DAMAGE.

JULIEN TACHOIRES SPECIFICALLY DISCLAIMS ANY WARRANTIES, INCLUDING, BUT
NOT LIMITED TO, THE IMPLIED WARRANTIES OF MERCHANTABILITY AND FITNESS
FOR A PARTICULAR PURPOSE. THE SOFTWARE PROVIDED HEREUNDER IS ON AN "AS IS"
BASIS, AND JULIEN TACHOIRES HAS NO OBLIGATIONS TO PROVIDE MAINTENANCE,
SUPPORT, UPDATES, ENHANCEMENTS, OR MODIFICATIONS.
"""

from __future__ import print_function

PGTOP_VERSION = "1.3.0"

import os
import sys
if os.name != 'posix':
    sys.exit("FATAL: Platform not supported.")
import signal
from optparse import OptionParser, OptionGroup
import socket
import getpass
import psutil
import curses
import psycopg2
from psycopg2 import errorcodes

from pgactivity import UI

# Customized OptionParser
class ModifiedOptionParser(OptionParser):
    """
    ModifiedOptionParser
    """
    def error(self, msg):
        raise OptionParsingError(msg)

class OptionParsingError(RuntimeError):
    """
    OptionParsingError
    """
    def __init__(self, msg):
        self.msg = msg

# Create the UI
PGAUI = UI.UI(PGTOP_VERSION)

def main():
    """
    Main function
    """
    try:
        try:
            parser = ModifiedOptionParser(
                add_help_option = False,
                version = "%prog "+PGTOP_VERSION,
                description = "htop like application for PostgreSQL \
server activity monitoring.")
            # -U / --username
            parser.add_option(
                '-U',
                '--username',
                dest = 'username',
                default = os.environ.get('PGUSER') or getpass.getuser(),
                help = "Database user name (default: \"%s\")."
                    % (getpass.getuser(),),
                metavar = 'USERNAME')
            # -p / --port
            parser.add_option(
                '-p',
                '--port',
                dest = 'port',
                default = os.environ.get('PGPORT') or '5432',
                help = "Database server port (default: \"5432\").",
                metavar = 'PORT')
            # -h / --host
            parser.add_option(
                '-h',
                '--host',
                dest = 'host',
                help = "Database server host or socket directory \
                            (default: \"localhost\").",
                metavar = 'HOSTNAME',
                default = os.environ.get('PGHOST') or 'localhost')
            # -d / --dbname
            parser.add_option(
                '-d',
                '--dbname',
                dest = 'dbname',
                help = "Database name to connect to (default: \"postgres\").",
                metavar = 'DBNAME',
                default = 'postgres')
            # -C / --no-color
            parser.add_option(
                '-C',
                '--no-color',
                dest = 'nocolor',
                action = 'store_true',
                help = "Disable color usage.",
                default = 'false')
            # --blocksize
            parser.add_option(
                '--blocksize',
                dest = 'blocksize',
                help = "Filesystem blocksize (default: 4096)",
                metavar = 'BLOCKSIZE',
                default = 4096)
            # --rds
            parser.add_option(
                '--rds',
                dest = 'rds',
                action = 'store_true',
                help = "Enable support for AWS RDS",
                default = 'false')
            group = OptionGroup(
                parser,
                "Display Options, you can exclude some columns by using them ")
            # --no-database
            group.add_option(
                '--no-database',
                dest = 'nodb',
                action = 'store_true',
                help = "Disable DATABASE.",
                default = 'false')
            # --no-user
            group.add_option(
                '--no-user',
                dest = 'nouser',
                action = 'store_true',
                help = "Disable USER.",
                default = 'false')
            # --no-client
            group.add_option(
                '--no-client',
                dest = 'noclient',
                action = 'store_true',
                help = "Disable CLIENT.",
                default = 'false')
            # --no-cpu
            group.add_option(
                '--no-cpu',
                dest = 'nocpu',
                action = 'store_true',
                help = "Disable CPU%.",
                default = 'false')
            # --no-mem
            group.add_option(
                '--no-mem',
                dest = 'nomem',
                action = 'store_true',
                help = "Disable MEM%.",
                default = 'false')
            # --no-read
            group.add_option(
                '--no-read',
                dest = 'noread',
                action = 'store_true',
                help = "Disable READ/s.",
                default = 'false')
            # --no-write
            group.add_option(
                '--no-write',
                dest = 'nowrite',
                action = 'store_true',
                help = "Disable WRITE/s.",
                default = 'false')
            # --no-time
            group.add_option(
                '--no-time',
                dest = 'notime',
                action = 'store_true',
                help = "Disable TIME+.",
                default = 'false')
            # --no-wait
            group.add_option(
                '--no-wait',
                dest = 'nowait',
                action = 'store_true',
                help = "Disable W.",
                default = 'false')
            parser.add_option_group(group)
            # --help
            parser.add_option(
                '--help',
                dest = 'help',
                action = 'store_true',
                help = "Show this help message and exit.",
                default = 'false')
            # --debug
            parser.add_option(
                '--debug',
                dest = 'debug',
                action = 'store_true',
                help = "Enable debug mode for traceback tracking.",
                default = 'false')
            (options, _) = parser.parse_args()
        except OptionParsingError as err:
            print('pg_activity: error: %s' % err.msg)
            print('Try "pg_activity --help" for more information.')
            sys.exit(1)
        if options.help is True:
            print(parser.format_help().strip())
            sys.exit(1)
        password = os.environ.get('PGPASSWORD')
        if password is None:
            # pgpass file handling
            try:
                pgpa_file = os.environ.get('PGPASSFILE')
                for (pgpa_h, pgpa_p, pgpa_d, pgpa_u, pgpa_pwd) in \
                    PGAUI.data.get_pgpass(pgpa_file):
                    if (pgpa_h == options.host or pgpa_h == '*') and \
                        (pgpa_p == options.port or pgpa_p == '*') and \
                        (pgpa_u == options.username or pgpa_u == '*') and \
                        (pgpa_d == options.dbname or pgpa_d == '*'):
                        password = pgpa_pwd
                        continue
            except Exception as err:
                pass
        debug = options.debug
        nb_try = 0
        while nb_try < 2:
            try:
                PGAUI.data.pg_connect(
                    host = options.host,
                    port = options.port,
                    user = options.username,
                    password = password,
                    database = options.dbname,
                    rds_mode = options.rds)
                break
            except psycopg2.OperationalError as err:
                if nb_try < 1 and (err.pgcode == errorcodes.INVALID_PASSWORD or
                    str(err).strip().startswith(
                        "FATAL:  password authentication failed for user")):
                    nb_try += 1
                    password = PGAUI.ask_password()
                elif nb_try < 1 and str(err).strip() == "fe_sendauth: no password supplied":
                    nb_try += 1
                    password = PGAUI.ask_password()
                else:
                    sys.exit("pg_activity: FATAL: %s" % (PGAUI.clean_str(str(err),)))

        pg_version = PGAUI.data.pg_get_version()
        PGAUI.data.pg_get_num_version(pg_version)
        hostname = socket.gethostname()
        # reduce DATABASE column length
        PGAUI.set_max_db_length(16)
        # blocksize
        PGAUI.set_blocksize(int(options.blocksize))
        # does pg_activity runing against local PG instance
        if not PGAUI.data.pg_is_local():
            PGAUI.set_is_local(False)
            PGAUI.set_start_line(2)
            hostname = options.host
        # if not connected to a local pg server, then go to degraded mode
        elif not PGAUI.data.pg_is_local_access():
            PGAUI.set_is_local(False)
            PGAUI.set_start_line(2)
            hostname = options.host
        # top part
        interval = 0
        if PGAUI.get_mode() == 'activities':
            queries =  PGAUI.data.pg_get_activities()
            procs = PGAUI.data.sys_get_proc(queries, PGAUI.get_is_local())
        elif PGAUI.get_mode() == 'waiting':
            procs = PGAUI.data.pg_get_waiting()
        elif PGAUI.get_mode() == 'blocking':
            procs = PGAUI.data.pg_get_blocking()
        # draw the flag
        flag = PGAUI.get_flag_from_options(options)
        # main loop
        disp_procs = None
        delta_disk_io = None
        # get DB informations
        db_info = PGAUI.data.pg_get_db_info(None, using_rds = options.rds)
        PGAUI.set_max_db_length(db_info['max_length'])
        # indentation
        indent = PGAUI.get_indent(flag)
        # Init curses
        PGAUI.init_curses()
        # color ?
        if options.nocolor == True:
            PGAUI.set_nocolor()
        else:
            PGAUI.set_color()
        while 1:
            PGAUI.check_window_size()
            old_pgtop_mode = PGAUI.get_mode()
            # poll process
            (disp_procs, new_procs) = PGAUI.poll(
                                        interval,
                                        flag,
                                        indent,
                                        procs,
                                        disp_procs)
            if PGAUI.get_mode() != old_pgtop_mode:
                indent = PGAUI.get_indent(flag)
            if PGAUI.get_is_local():
                delta_disk_io = PGAUI.data.get_global_io_counters()
            procs = new_procs
            # refresh the winodw
            db_info = PGAUI.data.pg_get_db_info(db_info, using_rds = options.rds)
            PGAUI.set_max_db_length(db_info['max_length'])
            # bufferize
            PGAUI.set_buffer({
                'procs': disp_procs,
                'extras': (
                    PGAUI.data.get_pg_version(),
                    hostname,
                    options.username,
                    options.host,
                    options.port,
                    options.dbname),
                'flag': flag,
                'indent': indent,
                'io': delta_disk_io,
                'tps': db_info['tps'],
                'size_ev': db_info['size_ev'],
                'total_size': db_info['total_size']
            })
            # refresh
            PGAUI.refresh_window(
                disp_procs,
                (
                    PGAUI.data.get_pg_version(),
                    hostname,
                    options.username,
                    options.host,
                    options.port,
                    options.dbname),
                flag,
                indent,
                delta_disk_io,
                db_info['tps'],
                db_info['size_ev'],
                db_info['total_size']
            )
            interval = 1

    except psutil.AccessDenied as err:
        PGAUI.at_exit_curses()
        sys.exit(
            "FATAL: Acces denied, can't acces system informations for PID %s"
            % (str(err),))
    except curses.error as err:
        PGAUI.at_exit_curses()
        if debug is True:
            import traceback
            exc_type, exc_value, exc_traceback = sys.exc_info()
            traceback.print_exception(
                exc_type,
                exc_value,
                exc_traceback,
                file=sys.stdout)
        sys.exit("FATAL: %s" % (str(err),))
    except KeyboardInterrupt as err:
        PGAUI.at_exit_curses()
        sys.exit(1)
    except Exception as err:
        PGAUI.at_exit_curses()
        # DEBUG
        if debug is True:
            import traceback
            exc_type, exc_value, exc_traceback = sys.exc_info()
            traceback.print_exception(
                exc_type,
                exc_value,
                exc_traceback,
                file=sys.stdout)
        sys.exit("FATAL: %s" % (str(err),))

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