#!/usr/bin/env python
# -*- coding: utf-8 -*-
##
#
# Vaisala software source code file
#
# Copyright (c) Vaisala Oyj 2015. All rights reserved.
#
##
"""
Tool for radar system GIS database maintenance.
"""
from __future__ import print_function

import ConfigParser
import argparse
import collections
import contextlib
import errno
import glob
import json
import logging
import os
import platform
import pwd
import string
import subprocess
import sys

from os import environ as env
from os import path
from hashlib import sha512


logger = logging.getLogger('rsw-gis-db-tool')


@contextlib.contextmanager
def working_directory(path):
    starting_directory = os.getcwd()
    try:
        os.chdir(path)
        yield
    finally:
        os.chdir(starting_directory)


DEFAULT_CONFIG_DIRECTORY = env.get("VAISALA_RADARSW_CONFIG_DIR",
                                   '/etc/vaisala/radarsw/configuration')


class DbToolException(Exception):
    pass


DBConfig = collections.namedtuple('DBConfig', ['username', 'password', 'database'])


def resolve_config(config_directory=DEFAULT_CONFIG_DIRECTORY):
    if 'GIS_DATABASE_USER' in env and \
       'GIS_DATABASE_USER_PASSWORD' in env and \
       'GIS_DATABASE_NAME' in env:
        return DBConfig(env['GIS_DATABASE_USER'],
                        env['GIS_DATABASE_USER_PASSWORD'],
                        env['GIS_DATABASE_NAME'])
    else:
        return read_config(config_directory=config_directory)


def read_config(config_directory=DEFAULT_CONFIG_DIRECTORY):
    '''Reads in configuration INI files such as:

    [GIS_DATABASE]
    gis.database.user = user
    gis.database.user.password = password
    gis.database.name = database

    Returns:
        A DBConfig instance.
    '''
    override_config_file = path.join(config_directory, 'gis-override.ini')
    default_config_file = path.join(config_directory, 'default', 'gis-default.ini')
    configs = [override_config_file, default_config_file]

    if any(map(path.isfile, configs)) is not True:
        raise DbToolException(
            u"No valid configuration files found, tried {0}".format(
                ", ".join(configs)))

    username, password, database = None, None, None
    for config_path in configs:
        if not path.isfile(config_path):
            continue

        logger.debug("Reading config file %s" % config_path)
        config = ConfigParser.RawConfigParser()
        config.read(config_path)

        username_option = 'GIS_DATABASE', 'gis.database.user'
        if username is None and config.has_option(*username_option):
            username = config.get(*username_option)

        password_option = 'GIS_DATABASE', 'gis.database.user.password'
        if password is None and config.has_option(*password_option):
            password = config.get(*password_option)

        database_option = 'GIS_DATABASE', 'gis.database.name'
        if database is None and config.has_option(*database_option):
            database = config.get(*database_option)

    if password is None:
        raise DbToolException("No password found for database!")

    result = DBConfig(username=username, password=password, database=database)
    logger.info("Using database configuration: %s" % repr(result))
    return result


def handle_error(step_description, tpl):
    """Error handler and reporter used in calling external executables.

    Args:
        step_description: A human-readable description of the step being
            executed fit for inclusion into an error message, in the form of
            'creating user' or 'dropping database'.
        tpl: A four-tuple consisting of 
            (command executed,
             command return code,
             stdout as string,
             stderr as string); basically something _run_as_superuser returns.
    """
    command, retcode, out, err = tpl
    if retcode == 0:
        return

    logger.error("Command output:")
    logger.error("----")
    logger.error(out)
    logger.error("----")
    logger.error("Command error output:")
    logger.error("----")
    logger.error(err)
    logger.error("----")

    raise DbToolException(
        u"Encountered error while {0}. Command exited with code {1}.".format(
            step_description, retcode
    ))


def _run_as_superuser(sql, database="postgres"):
    """Execute SQL as database super-user.

    Args:
        sql: SQL to execute as string.
        database: Database to run the SQL in. Default is 'postgres' as this
            tool is mainly used for managing databases.

    Returns:
        A four-tuple of
            (command executed,
             command return code,
             stdout as string,
             stderr as string)
    """
    if platform.system() != 'Linux':
        raise DbToolException("This tool only supports Linux")

    cmd = ['su', '-', 'postgres', '-c', "psql -t -v ON_ERROR_STOP=1 -d " + database + "; exit $?"]
    logger.info(u"Attempting to run {0}".format(' '.join(cmd)))
    p = subprocess.Popen(cmd, stdin=subprocess.PIPE, stdout=subprocess.PIPE, stderr=subprocess.PIPE)
    out_data, err_data = p.communicate(sql)
    return cmd, p.returncode, out_data, err_data


def create_db_cli(db_config):
    sql = '''
CREATE DATABASE %s
  WITH OWNER = %s
       ENCODING = 'UTF8'
       TABLESPACE = pg_default
       CONNECTION LIMIT = -1;
''' % (db_config.database, db_config.username)
    return _run_as_superuser(sql)


def install_extensions_cli(db_config):
    sql = '''
CREATE EXTENSION postgis;
CREATE EXTENSION hstore;
'''
    return _run_as_superuser(sql, database=db_config.database)


def create_user_cli(db_config):
    sql = '''
CREATE ROLE %s
  WITH
    LOGIN PASSWORD '%s'
    NOSUPERUSER
    INHERIT
    NOCREATEDB
    NOCREATEROLE
    NOREPLICATION;
''' % (db_config.username, db_config.password)
    return _run_as_superuser(sql)


def drop_db_cli(db_config):
    sql = 'DROP DATABASE IF EXISTS %s;' % db_config.database
    return _run_as_superuser(sql)


def vacuum_cli(db_config):
    sql = 'VACUUM ANALYZE;'
    return _run_as_superuser(sql, database=db_config.database)


def drop_user_cli(db_config):
    sql = 'DROP USER IF EXISTS %s;' % db_config.username
    return _run_as_superuser(sql)


def create_db(db_config, args):
    handle_error("creating database", create_db_cli(db_config))
    handle_error("installing PostgreSQL extensions", install_extensions_cli(db_config))


def create_user(db_config, args):
    return handle_error("creating user", create_user_cli(db_config))


def drop_db(db_config, args):
    return handle_error("dropping (removing) database", drop_db_cli(db_config))


def drop_user(db_config, args):
    return handle_error("dropping (removing) user", drop_user_cli(db_config))


def create(db_config, args):
    create_user(db_config, args)
    create_db(db_config, args)
    logger.info("User and database created successfully!")


def recreate(db_config, args):
    if env.get('RECREATE', None) != 'yes':
        print("Run with RECREATE environment variable set to yes", file=sys.stderr)
        print("  In Linux, try RECREATE=yes %s recreate" % sys.argv[0], file=sys.stderr)
        sys.exit(1)
    else:
        drop_db(db_config, args)
        logger.info("Database dropped (removed) succesfully!")
        drop_user(db_config, args)
        logger.info("User dropped (removed) successfully!")
        create(db_config, args)


def show(db_config, args):
    print(json.dumps(db_config._asdict()))


def random_password(length=32):
    chars = string.ascii_uppercase + string.digits + string.ascii_lowercase
    password = ''
    rng = random.SystemRandom()
    for i in range(length):
        password += rng.choice(chars)
    return password


def restore_dump(db_config, args):
    if not os.path.exists(args.gis_db_dump):
        raise DbToolException("GIS database dump must exist!")

    logger.info(u"Starting pg_restore from {} to database {}".format(
        args.gis_db_dump, db_config.database))

    part_paths = []
    checksum_path = None
    if os.path.isdir(args.gis_db_dump):
        for filename in os.listdir(args.gis_db_dump):
            filepath = os.path.join(args.gis_db_dump, filename)
            if 'part' in filename:
                part_paths.append(filepath)
            elif 'md5sums.txt' in filename:
                checksum_path = filename
            else:
                pass

    if len(part_paths) > 0 and checksum_path is not None:
        if not args.skip_checksum_check:
            logger.info(u"Checking DB dump file fragment integrity...")
            with working_directory(args.gis_db_dump):
                subprocess.check_call(["md5sum", "-c", checksum_path])

        logger.info("Restoring from a split PostgreSQL dump file...")
        command = u"cat {} | PGPASSWORD='{}' pg_restore -U '{}' -d '{}'".format(
            " ".join(sorted(part_paths)), db_config.password, db_config.username, db_config.database)
    else:
        logger.info("Restoring from standard (non-split) PostgreSQL dump....")
        command = u"PGPASSWORD='{}' pg_restore -U '{}' -d '{}' {}".format(
            db_config.password, db_config.username, db_config.database, args.gis_db_dump)

    subprocess.check_call(command, shell=True)
    logger.info("Restore done.")


def vacuum(db_config, args):
    handle_error("running database maintenance", vacuum_cli(db_config))


def main():
    parser = argparse.ArgumentParser(formatter_class=argparse.ArgumentDefaultsHelpFormatter)
    parser.add_argument("-d", "--debug", dest="debug", action="store_true",
                        help="Enable debug messages.")
    subparsers = parser.add_subparsers(title="supported commands")

    commands = {
        'create': ("Create user and database.", create),
        'create-db': ("Create database.", create_db),
        'create-user': ("Create database user.", create_user),
        'drop-db': ("Drop (remove completely) database.", drop_db),
        'drop-user': ("Drop (remove completely) database user", drop_user),
        'recreate': ("Recreate (drop + create) user and database.", recreate),
        'restore-dump': ("Restore database dump.", restore_dump),
        'show': ("Show database credentials.", show),
        'vacuum': ("Run maintenance tasks on the database.", vacuum)
    }

    for command, tpl in commands.iteritems():
        description, function = tpl
        cmd_parser = subparsers.add_parser(command, help=description,
                                           formatter_class=argparse.ArgumentDefaultsHelpFormatter)
        cmd_parser.set_defaults(func=function)

        if command == 'restore-dump':
            cmd_parser.add_argument("--gis-db-dump", dest="gis_db_dump", required=True,
                                    help="GIS database dump.", metavar="PATH")
            cmd_parser.add_argument("--skip-checksum-check", dest="skip_checksum_check",
                                    action="store_true", default=False,
                                    help="Don't check the integrity of GIS database dump parts.")

    args = parser.parse_args()

    def setup_logging(debug_enabled):
        root = logging.getLogger()
        if args.func == globals()['show']:
            root.setLevel(logging.ERROR)
        elif debug_enabled:
            root.setLevel(logging.DEBUG)
        else:
            root.setLevel(logging.INFO)
        ch = logging.StreamHandler(sys.stdout)
        ch.setLevel(logging.DEBUG)
        formatter = logging.Formatter('%(asctime)s - %(name)s - %(levelname)s - %(message)s')
        ch.setFormatter(formatter)
        root.addHandler(ch)
    setup_logging(args.debug)

    if pwd.getpwuid(os.getuid()).pw_name != 'root':
        sys.exit("This script must be run as root.")

    args = parser.parse_args()
    args.func(resolve_config(), args)


if __name__ == '__main__':
    main()
