#!/usr/bin/env python2

# THIS FILE IS PART OF THE CYLC SUITE ENGINE.
# Copyright (C) 2008-2018 NIWA & British Crown (Met Office) & Contributors.
#
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
#
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
#
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <http://www.gnu.org/licenses/>.

"""
Display the state of live task proxies in a running suite.

For color terminal ASCII escape codes, see
http://ascii-table.com/ansi-escape-sequences.php
"""

import sys
if '--use-ssh' in sys.argv[1:]:
    # requires local terminal
    sys.exit("No '--use-ssh': this command requires a local terminal.")

import os
import re
from time import sleep, time


from parsec.OrderedDict import OrderedDict
from cylc.option_parsers import CylcOptionParser as COP
from cylc.network.httpclient import SuiteRuntimeServiceClient, ClientError
from cylc.cfgspec.glbl_cfg import glbl_cfg
from cylc.task_state import (
    TASK_STATUS_RUNAHEAD, TASK_STATUSES_ORDERED,
    TASK_STATUSES_RESTRICTED)
from cylc.task_state_prop import get_status_prop
from cylc.wallclock import get_time_string_from_unix_time


SUITE_STATUS_SPLIT_REC = re.compile('^([a-z ]+ at )(.*)$')


class SuiteMonitor(object):

    def __init__(self):
        self.parser = COP(
            """cylc [info] monitor [OPTIONS] ARGS

A terminal-based live suite monitor.  Exit with 'Ctrl-C'.

The USER_AT_HOST argument allows suite selection by 'cylc scan' output:
  cylc monitor $(cylc scan | grep <suite_name>)
""",
            argdoc=[('REG', 'Suite name'),
                    ('[USER_AT_HOST]', 'user@host:port, shorthand for --user, '
                     '--host & --port.')], comms=True, noforce=True)

        self.parser.add_option(
            "-a", "--align",
            help="Align task names. Only useful for small suites.",
            action="store_true", default=False, dest="align_columns")

        self.parser.add_option(
            "-r", "--restricted",
            help="Restrict display to active task states. "
            "This may be useful for monitoring very large suites. "
            "The state summary line still reflects all task proxies.",
            action="store_true", default=False, dest="restricted")

        def_sort_order = glbl_cfg().get(["monitor", "sort order"])

        self.parser.add_option(
            "-s", "--sort", metavar="ORDER",
            help="Task sort order: \"definition\" or \"alphanumeric\"."
            "The default is " + def_sort_order + " order, as determined by "
            "global config. (Definition order is the order that tasks appear "
            "under [runtime] in the suite definition).",
            action="store", default=def_sort_order, dest="sort_order")

        self.parser.add_option(
            "-o", "--once",
            help="Show a single view then exit.",
            action="store_true", default=False, dest="once")

        self.parser.add_option(
            "-u", "--runahead",
            help="Display task proxies in the runahead pool (off by default).",
            action="store_true", default=False, dest="display_runahead")

        self.parser.add_option(
            "-i", "--interval",
            help="Interval between suite state retrievals, "
                 "in seconds (default 1).",
            metavar="SECONDS", action="store", default=1,
            dest="update_interval")

    def reset(self, suite, owner, host, port, timeout):
        self.pclient = SuiteRuntimeServiceClient(
            suite, owner, host, port, timeout)

    def run(self):
        (options, args) = self.parser.parse_args()
        suite = args[0]

        if len(args) > 1:
            try:
                user_at_host, options.port = args[1].split(':')
                options.owner, options.host = user_at_host.split('@')
            except ValueError:
                print >> sys.stderr, ('USER_AT_HOST must take the form '
                                      '"user@host:port"')
                sys.exit(1)

        client_name = os.path.basename(sys.argv[0])
        if options.restricted:
            client_name += " -r"

        legend = ''
        for state in TASK_STATUSES_ORDERED:
            legend += get_status_prop(state, 'ascii_ctrl')
        legend = legend.rstrip()

        len_header = sum(len(s) for s in TASK_STATUSES_ORDERED)

        self.reset(suite, options.owner, options.host, options.port,
                   options.comms_timeout)

        is_cont = False
        while True:
            if is_cont:
                if options.once:
                    break
                else:
                    sleep(float(options.update_interval))
            is_cont = True
            try:
                glbl, task_summaries = (
                    self.pclient.get_suite_state_summary()[0:2])
            except ClientError as exc:
                print >> sys.stderr, "\033[1;37;41mERROR\033[0m", str(exc)
                self.reset(suite, options.owner, options.host, options.port,
                           options.comms_timeout)
            else:
                if not glbl:
                    print >> sys.stderr, (
                        "\033[1;37;41mWARNING\033[0m suite initialising")
                    continue
                states = [t["state"] for t in task_summaries.values() if (
                          "state" in t)]
                n_tasks_total = len(states)
                if options.restricted:
                    task_summaries = dict(
                        (i, j) for i, j in task_summaries.items() if (
                            j['state'] in TASK_STATUSES_RESTRICTED))
                if not options.display_runahead:
                    task_summaries = dict(
                        (i, j) for i, j in task_summaries.items() if (
                            j['state'] != TASK_STATUS_RUNAHEAD))
                try:
                    updated_at = get_time_string_from_unix_time(
                        glbl['last_updated'])
                except KeyError:
                    updated_at = time()
                except (TypeError, ValueError):
                    # Older suite.
                    updated_at = glbl['last_updated'].isoformat()

                run_mode = glbl['run_mode']
                ns_defn_order = glbl['namespace definition order']
                status_string = glbl['status_string']

                task_info = {}
                name_list = set()
                task_ids = task_summaries.keys()
                for task_id in task_ids:
                    name = task_summaries[task_id]['name']
                    point_string = task_summaries[task_id]['label']
                    state = task_summaries[task_id]['state']
                    name_list.add(name)
                    if point_string not in task_info:
                        task_info[point_string] = {}
                    task_info[point_string][name] = get_status_prop(
                        state, 'ascii_ctrl', subst=name)

                # Sort the tasks in each cycle point.
                if options.sort_order == "alphanumeric":
                    sorted_name_list = sorted(name_list)
                else:
                    sorted_name_list = ns_defn_order

                sorted_task_info = {}
                for point_str, info in task_info.items():
                    sorted_task_info[point_str] = OrderedDict()
                    for name in sorted_name_list:
                        if name in name_list:
                            # (Defn order includes family names.).
                            sorted_task_info[point_str][name] = info.get(name)

                # Construct lines to blit to the screen.
                blit = []

                suite_name = suite
                if run_mode != "live":
                    suite_name += " (%s)" % run_mode
                prefix = "%s - %d tasks" % (suite_name, int(n_tasks_total))
                suffix = "%s %s" % (client_name, self.pclient.my_uuid)
                title_str = ' ' * len_header
                title_str = prefix + title_str[len(prefix):]
                title_str = '\033[1;37;44m%s%s\033[0m' % (
                    title_str[:-len(suffix)], suffix)
                blit.append(title_str)
                blit.append(legend)

                updated_str = "updated: \033[1;38m%s\033[0m" % updated_at
                blit.append(updated_str)

                summary = 'state summary:'
                state_totals = glbl['state totals']
                for state, tot in state_totals.items():
                    subst = " %d " % tot
                    summary += get_status_prop(state, 'ascii_ctrl', subst)
                blit.append(summary)

                # Print a divider line containing the suite status string.
                try:
                    status1, status2 = (
                        SUITE_STATUS_SPLIT_REC.match(status_string)).groups()
                except AttributeError:
                    status1 = status_string
                    status2 = ''
                suffix = '_'.join(list(status1.replace(' ', '_'))) + status2
                divider_str = '_' * len_header
                divider_str = "\033[1;31m%s%s\033[0m" % (
                    divider_str[:-len(suffix)], suffix)
                blit.append(divider_str)

                blitlines = {}
                for point_str, val in sorted_task_info.items():
                    indx = point_str
                    line = "\033[1;34m%s\033[0m" % point_str
                    for name, info in val.items():
                        if info is not None:
                            line += " %s" % info
                        elif options.align_columns:
                            line += " %s" % (' ' * len(name))
                    blitlines[indx] = line

                if not options.once:
                    os.system("clear")
                print '\n'.join(blit)
                indxs = blitlines.keys()
                try:
                    int(indxs[1])
                except (ValueError, IndexError):
                    indxs.sort()
                else:
                    indxs.sort(key=int)
                for ix in indxs:
                    print blitlines[ix]


if __name__ == "__main__":
    try:
        SuiteMonitor().run()
    except KeyboardInterrupt:
        pass
