#!/usr/bin/env python

from __future__ import print_function

import argparse
import subprocess
import socket
from collections import defaultdict
from clustertools import Monitor

__author__ = "Begon Jean-Michel <jm.begon@gmail.com>"
__copyright__ = "3-clause BSD License"


__LH__ = "localhost"


def reset(exp_name, hostname):
    try:
        histo = Monitor(exp_name)
        print(hostname,":", exp_name, histo.count_by_state(), end=" ")
        histo.to_launchables()
        histo.refresh()
        print(histo.count_by_state)
    except:
        print("No such experiment")


def aborted2launchable(exp_name, hostname):
    try:
        histo = Monitor(exp_name)
        print(hostname,":", exp_name, histo.count_by_state(), end=" ")
        histo.aborted_to_launchable()
        histo.refresh()
        print(histo.count_by_state)
    except:
        print("No such experiment")


if __name__ == "__main__":
    description = "Clustertools task monitoring script"
    doc = """The experiment(s) to process.
Use ssh-like syntax: 'luke@skywalker:exp_about_the_force'. The first part can
be omitted for local processing."""

    parser = argparse.ArgumentParser(description=description)
    parser.add_argument("queries", help=doc, nargs="*")
    action = parser.add_mutually_exclusive_group(required=True)
    action.add_argument("--reset", "-r",  action='store_true', default=False,
                        help="Reset the given experiments")
    parser.add_argument("--abr2lnch", "-a", action='store_true', default=False,
                        help="Reset the aborted jobs to launchable state")

    # =========== Preprocessing ============= #

    args = parser.parse_args()
    queries = args.queries

    # Sort query by host to group them
    sorted_queries = defaultdict(list)
    for query in queries:
        try:
            host, exp_name = query.split(":")
        except ValueError:
            host = __LH__
            exp_name = query
        sorted_queries[host] = exp_name

    # =========== Main ============= #
    if __LH__ in sorted_queries:
        try:
            hostname = socket.gethostname()
        except:
            hostname = "Unkown host"
        for exp_name in sorted_queries[__LH__]:
            if args.reset:
                reset(exp_name, hostname)
            elif args.abr2lnch:
                aborted2launchable(exp_name, hostname)
        del sorted_queries[__LH__]

    # Remote queries
    cmd = "ssh {host} {prog} {exp}".format(prog=args.prog)
    for host, exp_names in sorted_queries.items():
        exp_string = " ".join(exp_names)
        subprocess.call(cmd.format(host=host, exp=exp_string), shell=True)

