from __future__ import print_function
import os
import sys
import json
import signal
import logging
import argparse
import functools
import subprocess
try:
    from termcolor import cprint
except ImportError:
    def cprint(*args):
        print(args[0])
from datetime import timedelta
from gpstime import gpstime, GPSTimeException
from dateutil.tz import tzutc, tzlocal

from . import systemd
from . import svlogd
from . import util

##################################################

PROG = 'guardctrl'

parser = argparse.ArgumentParser(
    prog=PROG,
    description="""Guardian daemon supervision control interface.

Control guardian daemon processes managed by the site systemd
supervision system.
""",
    formatter_class=argparse.RawDescriptionHelpFormatter,
    epilog="""
Add '-h' after individual commands for command help.

Node names may be specified with wildcard/globbing, e.g. 'SUS_*'.
A single '*' will act on all configured nodes (where appropriate).
""",
)

parser.add_argument('-d', '--debug', action='store_true',
                    help="print debug information to stderr")
# parser.add_argument('-v', '--version', action='version',
#                     version='Guardctrl {}'.format(__version__),
#                     help="print guardctrl version and exit")

subparser = parser.add_subparsers(
    title='Commands',
    metavar='<command>',
    dest='cmd',
    #help=argparse.SUPPRESS,
)
subparser.required = True

def gen_subparser(cmd, func, node_nargs=None, **kwargs):
    help = func.__doc__.split('\n')[0].lower().strip('.')
    desc = func.__doc__
    p = subparser.add_parser(
        cmd,
        #formatter_class=argparse.RawDescriptionHelpFormatter,
        help=help,
        description=desc,
        **kwargs
    )
    p.set_defaults(func=func)
    if node_nargs:
        p.add_argument('nodes', metavar='node', nargs=node_nargs,
                       help="guardian node name (may be wildcard/glob pattern)")
    return p

##################################################

def command_help(args):
    """Show command help."""
    if args.cmd:
        try:
            p = subparser.choices[args.cmd]
        except KeyError:
            sys.exit("Unknown command: {}".format(args.cmd))
    else:
        p = parser
    p.print_help()

p = gen_subparser("help", command_help)
p.add_argument('cmd', metavar='<command>', nargs='?', help="command")


def list_nodes(args):
    """List nodes and node supervision state."""
    nodes = args.nodes
    include_status = args.include_status
    only_active = args.only_active
    if not nodes:
        nodes = '*'
    nodes = list(systemd.list_nodes(nodes))
    if not nodes:
        return
    if not include_status and not only_active:
        for node in nodes:
            print(node)
        return
    node_w = max([len(node) for node in nodes]) + 2
    if include_status:
        fmt = '{{node:<{node_w}}} {{enabled:<{enabled_w}}} {{active:<{active_w}}} {{execstart}}'.format(
            node_w=node_w,
            enabled_w=10,
            active_w=10,
        )
    else:
        fmt = '{node}'
    for node, enabled, active, execstart in systemd.list_nodes(
            nodes,
            only_active=only_active,
            include_status=True):
        color = None
        attrs = []
        if active == 'active':
            color = 'green'
        elif active == 'failed':
            color = 'red'
            attrs = ['bold']
        else:
            color = 'yellow'
        line = fmt.format(node=node, enabled=enabled, active=active, execstart=execstart)
        cprint(line, color, attrs=attrs)

p = gen_subparser('list', list_nodes, node_nargs='*')
p.add_argument('-s', '--status', dest='include_status', action='store_true',
               help="include node status columns")
p.add_argument('-r', '--running', dest='only_active', action='store_true',
               help="only list running (active) nodes")


@functools.wraps(systemd.print_node_status)
def print_node_status(args):
    nodes = args.nodes
    try:
        return systemd.print_node_status(nodes, print_logs=True)
    except util.GuardCtrlError as e:
        sys.exit(e)

gen_subparser('status', print_node_status, node_nargs='*')


@functools.wraps(systemd.enable_nodes)
def enable_nodes(args):
    nodes = args.nodes
    now = args.now
    try:
        return systemd.enable_nodes(nodes, now=now)
    except util.GuardCtrlError as e:
        sys.exit(e)

p = gen_subparser('enable', enable_nodes, node_nargs='+')
p.add_argument('-n', '--now', action='store_true',
               help="start nodes immediately")


@functools.wraps(systemd.start_nodes)
def start_nodes(args):
    nodes = args.nodes
    try:
        return systemd.start_nodes(nodes)
    except util.GuardCtrlError as e:
        sys.exit(e)

gen_subparser('start', start_nodes, node_nargs='+')


@functools.wraps(systemd.restart_nodes)
def restart_nodes(args):
    nodes = args.nodes
    try:
        return systemd.restart_nodes(nodes)
    except util.GuardCtrlError as e:
        sys.exit(e)

gen_subparser('restart', restart_nodes, node_nargs='+')


@functools.wraps(systemd.stop_nodes)
def stop_nodes(args):
    nodes = args.nodes
    try:
        return systemd.stop_nodes(nodes)
    except util.GuardCtrlError as e:
        sys.exit(e)

gen_subparser('stop', stop_nodes, node_nargs='+')


@functools.wraps(systemd.disable_nodes)
def disable_nodes(args):
    nodes = args.nodes
    now = args.now
    node_list = list(systemd.list_nodes(nodes))
    if not node_list:
        sys.exit("Unknown nodes: {}".format(' '.join(nodes)))
    print("Really disable the following nodes?", file=sys.stderr)
    for node in node_list:
        print("  {}".format(node), file=sys.stderr)
    print("This will remove these nodes from the site list.", file=sys.stderr)
    ret = input("Type 'yes' to disable: ")
    if ret != 'yes':
        sys.exit('Aborted.')
    try:
        return systemd.disable_nodes(node_list, now=now)
    except util.GuardCtrlError as e:
        sys.exit(e)

p = gen_subparser('disable', disable_nodes, node_nargs='+')
p.add_argument('-n', '--now', action='store_true',
               help="stop nodes immediately")

##################################################

def datetime(ds):
    """wrapper function to parse datetime strings to gpstime objects.

    Throws TypeError needed for argparse type checking.

    """
    try:
        return gpstime.parse(ds)
    except GPSTimeException as e:
        raise TypeError(e)


def print_logs(args):
    """View node logs.

Logs will be followed by default, unless the +f flag is specified.
Log timestamps are in UTC by default.
    """
    node = args.nodes
    time_format = args.time_format
    search_old_logs = args.search_old_logs
    # ddelta = timedelta(seconds=float(args.duration))

    # if carg('hours'):
    #     args.after = gpstime.now(tz=tzlocal()) - timedelta(seconds=3600*float(args.hours))
    # elif carg('after'):
    #     args.after = _parse_time_arg('after')
    #     if carg('before'):
    #         args.before = _parse_time_arg('before')
    #         if args.before < args.after:
    #             raise util.GuardCtrlError("Before time is earlier than after time.")
    #     elif carg('duration'):
    #         args.before = args.after + ddelta
    # elif carg('center'):
    #     dur = ddelta/2
    #     args.after = _parse_time_arg('center') - dur
    #     args.before = _parse_time_arg('center') + dur
    # elif carg('before'):
    #     args.before = _parse_time_arg('before')
    #     args.after = args.before - ddelta
    # # else:
    # #     args.follow = True

    logging.debug("search after : %s" % args.after)
    logging.debug("search before: %s" % args.before)

    if time_format == 'utc':
        def dt2s(dt):
            return dt.astimezone(tzutc()).strftime(util.ISO_FMT+'Z')
    elif time_format == 'local':
        def dt2s(dt):
            return dt.strftime(util.ISO_FMT)
    elif time_format == 'gps':
        def dt2s(dt):
            return '{:f}'.format(gpstime.fromdatetime(dt).gps())

    if search_old_logs:
        kwargs = {k: getattr(args, k) for k in ['nodes', 'after', 'before']}
        gen = svlogd.node_log_search(**kwargs)
    else:
        kwargs = {k: getattr(args, k) for k in ['nodes', 'after', 'before', 'nlines', 'follow']}
        gen = systemd.node_logs(**kwargs)

    for dt, msg in gen:
        if dt:
            print('{} {}'.format(dt2s(dt), msg))
        else:
            print(msg)

p = gen_subparser('log', print_logs, node_nargs='*', prefix_chars='-+')
p.epilog = """
Date/time strings for filters are handled by a date/time string parser
and therefore may be specified in various formats.  If --after and
--before are not both specified, --duration will be used.  --lines
will ultimately limit the number of lines returned.

Examples:

  Show all ISC_LOCK logs from Sept. 14, 2015:

  $ {prog} --after 2015-09-14 --before 2015-09-15 ISC_LOCK

  Show the last 1000 lines of all SUS node logs:

  $ {prog} -n1000 SUS_*

  Quote compound date strings, and specify 'UTC' for times in UTC:

  $ {prog} -a "2016-08-20 06:23:22 UTC" -b "2016-08-20 06:23:52 UTC" ISC_LOCK

  GPS times are accepted as well:

  $ {prog} 1155709419 ISC_DRMI

""".format(prog=PROG+' log')
p.formatter_class = argparse.RawDescriptionHelpFormatter
p.add_argument('-f', '--follow', action='store_true',
               help="follow logs")
p.add_argument('-a', '--after', metavar='DATETIME', type=datetime,
               help="filter for logs after specified date/time")
p.add_argument('-b', '--before', metavar='DATETIME', type=datetime,
               help="filter for logs before specified date/time")
# p.add_argument('-c', '--center', metavar='DATETIME', type=datetime,
#                help="filter for logs centered around specified time")
# p.add_argument('-o', '--hours', 'hours', type=int, metavar='HOURS',
#                help="filter for logs from the last H hours")
# p.add_argument('-t', '--day', 'hours', flag_value=24,
#                help="filter for logs from the last 24 hours")
# p.add_argument('-d', '--duration', metavar='SECONDS', type=float, default=60)
p.add_argument('-n', '--lines', metavar='N', dest='nlines', type=int, default=100,
               help="limit logs to N lines [100]")
p.add_argument('-u', '--utc', dest='time_format', action='store_const', const='utc',
               help="show log lines with UTC timestamps (default)")
p.add_argument('-l', '--local', dest='time_format', action='store_const', const='local',
               help="show log lines with local timestamps")
p.add_argument('-g', '--gps', dest='time_format', action='store_const', const='gps',
               help="show log lines with GPS timestamps")
p.set_defaults(time_format='utc')
p.add_argument('--old', dest='search_old_logs', action='store_true',
               help="search old logs (<=O2, some arguments ignored)")

##################################################

def system_command(cmd, argv):
    cmd = [cmd] + argv
    logging.debug("cmd: {}".format(' '.join(cmd)))
    try:
        sys.exit(subprocess.run(cmd).returncode)
    except KeyboardInterrupt:
        sys.exit()

subparser.add_parser('env')
subparser.add_parser('systemctl')
subparser.add_parser('journalctl')

##################################################

def main_local():
    # handle invocation as ssh ForceCommand
    ssh_args = os.getenv('SSH_ORIGINAL_COMMAND')
    if ssh_args and len(sys.argv) == 1:
        # assume json argument encoding if the first character is the
        # json bracket list start character
        if ssh_args[0] == '[':
            sys.argv += json.loads(ssh_args)
        else:
            sys.argv += os.getenv('SSH_ORIGINAL_COMMAND').split()

    # parse known args first, just to get debug flag and handle direct
    # sys commands
    args, remainder = parser.parse_known_args()

    log_level = os.getenv('LOG_LEVEL', 'INFO').upper()
    if args.debug:
        log_level = 'DEBUG'
    FORMAT = '%(levelname)s: %(module)s: %(message)s'
    logging.basicConfig(format=FORMAT, level=log_level)

    logging.debug(sys.argv)

    # handle direct systemctl/journalctl invocation
    if args.cmd in ['env', 'systemctl', 'journalctl']:
        system_command(args.cmd, remainder)

    else:
        # re-parse all args to properly fail for unknown args
        args = parser.parse_args()
        args.func(args)


def main_remote():
    parser = argparse.ArgumentParser(prefix_chars='-+', add_help=False)
    parser.add_argument('+t', dest='terminal', action='store_false')
    wargs, args = parser.parse_known_args()

    IFO = os.getenv('IFO')
    GUARDCTRL_USER = os.getenv('GUARDCTRL_USER', 'guardian')
    GUARDCTRL_HOST = os.getenv('GUARDCTRL_HOST', '{}guardian'.format(IFO.lower()))

    SSH_CMD = [
        'ssh',
        '-S', '~/.ssh/controls/%r@%h:%p',
        '-o', 'ControlMaster=auto',
        '-o', 'ControlPersist=600',
    ]

    # force TTY by default, with "+t" option to unset and allow
    # separate stdout/stderr streams
    if wargs.terminal:
        SSH_CMD += ['-t']

    SSH_CMD += [
        '{}@{}'.format(GUARDCTRL_USER, GUARDCTRL_HOST),
        '--',
    ]

    # create a directory for the ssh control sockets
    try:
        os.makedirs(os.path.expanduser('~/.ssh/controls/'))
    except:
        pass

    # json encode commands for ssh, to pass spaces in a reasonable way
    args_enc = json.dumps(args)

    # launch in foreground
    try:
        subprocess.run(
            SSH_CMD + [args_enc],
            stderr=subprocess.PIPE,
            check=True,
        )
    except subprocess.CalledProcessError as e:
        sys.exit("cmd: {}\n{}".format(
            ' '.join(e.cmd), e.stderr.decode().strip()))


def main():
    signal.signal(signal.SIGINT, signal.SIG_DFL)
    signal.signal(signal.SIGPIPE, signal.SIG_DFL)

    if os.path.exists(os.path.expanduser('~/.guardctrl-home')) \
       or os.getenv('GUARDCTRL_LOCAL'):
        main_local()

    elif os.getenv('GUARDCTRL_HOST'):
        main_remote()

    else:
        main_local()


if __name__ == '__main__':
    main()
