#!/usr/bin/python
# v1.0.0  Kit Menlove    2016-12-21
from __future__ import print_function
import ecflow
import argparse

parser = argparse.ArgumentParser(description='Suspend or release all transfer jobs on a given system in a given environment (or all environments).')
parser.add_argument('system_phase', metavar='MACHINE_PHASE', choices=('p1', 'p2', 'xc40', 'p3'), help="The system phase of the nodes you wish to target (p1, p2, xc40, or p3)")
parser.add_argument('-e', '--envir', metavar='envir', choices=('prod', 'para', 'test', 'all'), default='all', help="Environment of the transfer jobs to suspend or release (prod, para, test, or all) (default: all)")
parser.add_argument('-r', '--resume', action='store_true', help="Resume the nodes instead of suspending them")
parser.add_argument('-f', '--force', action='store_true', help="Do not ask for confirmation")
args = parser.parse_args()

try:
    # Create the client
    ci = ecflow.Client()

    # Get the node tree suite definition as stored in the server
    # The definition is retrieved and stored on the variable 'ci'
    ci.sync_local()

    # access the definition retrieved from the server
    defs = ci.get_defs()

    if defs == None :
        print("The server has no definition")
        exit(1)

    if args.envir == "all":
        node_vec = []
        for envir in 'prod', 'para', 'test':
            family = defs.find_abs_node(envir+'/transfer')
            if family:
                node_vec += family.get_all_nodes()
    else:
        node_vec = defs.find_abs_node(args.envir+'/transfer').get_all_nodes()

    nodes_to_act_upon = []

    # iterate over tasks and print path and state
    for node in node_vec:
        if node.find_variable('MACHINE_PHASE').value() == args.system_phase:
            if args.resume:
                if node.is_suspended():
                    nodes_to_act_upon.append(node.get_abs_node_path())
                else:
                    print(node.get_abs_node_path() + " is not currently suspended")
            else:
                if node.is_suspended():
                    print(node.get_abs_node_path() + " is already suspended")
                else:
                    nodes_to_act_upon.append(node.get_abs_node_path())

    if not nodes_to_act_upon:
        if args.resume:
            print("\nNo nodes were found to resume.  Exiting...")
        else:
            print("\nNo nodes were found to suspend.  Exiting...")
        exit(1)

    if args.resume:
        print("\nThe following nodes will be resumed:")
    else:
        print("\nThe following nodes will be suspended:")

    for node_path in nodes_to_act_upon:
        print(node_path)

    # Print the nodes to suspend/resume and ask the user for confirmation
    if args.force:
        response = 'Y'
    else:
        response = raw_input("\nDoes thos look correct? (Y/N) ")

    # If the user confirmed the nodes are correct, suspend/resume them!
    if response.upper() == 'Y' or response.lower() == 'yes':
        if args.resume:
            ci.resume(nodes_to_act_upon)
            print("The nodes have been resumed.")
        else:
            ci.suspend(nodes_to_act_upon)
            print("The nodes have been suspended.")
    else:
        print("No changes were made.")

except RuntimeError, e:
    print("Failed: " + str(e))