#!/usr/bin/env python3

# THIS FILE IS PART OF THE CYLC SUITE ENGINE.
# Copyright (C) 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/>.

"""cylc [control] stop|shutdown [OPTIONS] ARGS

Tell a suite server program to shut down. In order to prevent failures going
unnoticed, suites only shut down automatically at a final cycle point if no
failed tasks are present. There are several shutdown methods:

  1. (default) stop after current active tasks finish
  2. (--now) stop immediately, orphaning current active tasks
  3. (--kill) stop after killing current active tasks
  4. (with STOP as a cycle point) stop after cycle point STOP
  5. (with STOP as a task ID) stop after task ID STOP has succeeded
  6. (--wall-clock=T) stop after time T (an ISO 8601 date-time format e.g.
     CCYYMMDDThh:mm, CCYY-MM-DDThh, etc).

Tasks that become ready after the shutdown is ordered will be submitted
immediately if the suite is restarted.  Remaining task event handlers and job
poll and kill commands, however, will be executed prior to shutdown, unless
--now is used.

This command exits immediately unless --max-polls is greater than zero, in
which case it polls to wait for suite shutdown."""

import sys
if '--use-ssh' in sys.argv[1:]:
    sys.argv.remove('--use-ssh')
    from cylc.flow.remote import remrun
    if remrun():
        sys.exit(0)

from cylc.flow.command_polling import Poller
from cylc.flow.exceptions import ClientError, ClientTimeout
from cylc.flow.network.client import SuiteRuntimeClient
from cylc.flow.option_parsers import CylcOptionParser as COP
from cylc.flow.task_id import TaskID
from cylc.flow.terminal import prompt, cli_function


class StopPoller(Poller):
    """A polling object that checks if a suite has stopped yet."""

    def __init__(self, pclient, condition, interval, max_polls):
        Poller.__init__(self, condition, interval, max_polls, None)
        self.pclient = pclient

    def check(self):
        """Return True if suite has stopped (success) else False"""
        try:
            self.pclient('ping_suite')
        except (ClientError, ClientTimeout):
            # failed to ping - suite stopped
            return True
        else:
            # pinged - suite must be alive
            return False


def get_option_parser():
    parser = COP(
        __doc__, comms=True,
        argdoc=[("REG", "Suite name"),
                ("[STOP]", """a/ task POINT (cycle point), or
                            b/ ISO 8601 date-time (clock time), or
                            c/ TASK (task ID).""")])

    parser.add_option(
        "-k", "--kill",
        help="Shut down after killing currently active tasks.",
        action="store_true", default=False, dest="kill")

    parser.add_option(
        "-n", "--now",
        help=(
            "Shut down without waiting for active tasks to complete." +
            " If this option is specified once," +
            " wait for task event handler, job poll/kill to complete." +
            " If this option is specified more than once," +
            " tell the suite to terminate immediately."),
        action="count", default=0, dest="now")

    parser.add_option(
        "-w", "--wall-clock", metavar="STOP",
        help="Shut down after time STOP (ISO 8601 formatted)",
        action="store", dest="wall_clock")

    StopPoller.add_to_cmd_options(parser, d_max_polls=0)

    return parser


@cli_function(get_option_parser)
def main(parser, options, suite, shutdown_arg=None):
    if shutdown_arg is not None and options.kill:
        parser.error("ERROR: --kill is not compatible with [STOP]")

    if options.kill and options.now:
        parser.error("ERROR: --kill is not compatible with --now")

    pclient = SuiteRuntimeClient(
        suite, options.owner, options.host, options.port,
        options.comms_timeout)

    if int(options.max_polls) > 0:
        # (test to avoid the "nothing to do" warning for # --max-polls=0)
        spoller = StopPoller(
            pclient, "suite stopped", options.interval, options.max_polls)

    if options.wall_clock:
        prompt(
            'Set shutdown at wall clock %s for %s' % (
                options.wall_clock, suite),
            options.force)
        pclient('set_stop_after_clock_time',
                {'datetime_string': options.wall_clock})
    elif shutdown_arg is not None and TaskID.is_valid_id(shutdown_arg):
        # STOP argument detected
        prompt(
            'Set shutdown after task %s for %s' % (shutdown_arg, suite),
            options.force)
        pclient('set_stop_after_task', {'task_id': shutdown_arg})
    elif shutdown_arg is not None:
        # not a task ID, may be a cycle point
        prompt(
            'Set shutdown at cycle point %s for %s' % (shutdown_arg, suite),
            options.force)
        pclient('set_stop_after_point', {'point_string': shutdown_arg})
    elif options.now > 1:
        prompt('Shut down and terminate %s now' % suite, options.force)
        pclient('stop_now', {'terminate': True})
    elif options.now:
        prompt('Shut down %s now' % suite, options.force)
        pclient('stop_now')
    else:
        prompt('Shut down %s' % suite, options.force)
        pclient('set_stop_cleanly', {'kill_active_tasks': options.kill})

    if int(options.max_polls) > 0:
        # (test to avoid the "nothing to do" warning for # --max-polls=0)
        if not spoller.poll():
            sys.exit(1)


if __name__ == "__main__":
    main()
