#!/usr/bin/env python

from taskManager.ProcessHandler import ProcessHandler
from taskManager.config_helpers import *
from taskManager.ResourceMonitor import ResourceMonitor
import argparse
import signal


def string_as_bool(s):
    s = s.lower()
    boolean = None

    if s in {"t", "true", "1", "y", "yes"}:
        boolean = True
    elif s in {"f", "false", "0", "n", "no"}:
        boolean = False
    else:
        exit("Error: invalid argument specified for boolean flag: %s" % s)

    return boolean


def select_not_none(a, b):
    if a is None:
        return b
    else:
        return a


def comma_separated_list(s):
    tokens = s.strip().split(",")
    return tokens


def space_separated_list(s):
    tokens = s.strip().split(" ")
    return tokens


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    subparsers = parser.add_subparsers(dest="command")
    subparsers.required = True
    config = get_task_manager_config()

    # parsers for running the full pipeline
    run_parser = subparsers.add_parser("run", help="Run command, track memory, cpu and IO usage. "
                                                   "Send email when command exits")
    subparsers.add_parser("configure", help="Configure taskManager to send emails to "
                                            "the correct location")

    parser.register("type", "bool", string_as_bool)
    parser.register("type", "comma_separated_list", comma_separated_list)
    parser.register("type", "space_separated_list", space_separated_list)

    run_parser.add_argument("--command", "-c",
                            dest="run_command",
                            required=True,
                            type=space_separated_list,
                            help="The command to be run inside a subprocess: eg. -c 'sleep 10' ")
    run_parser.add_argument("--redirect_stdout", "-r",
                            dest="redirect_stdout",
                            required=False,
                            type=str,
                            help="Redirect stdout of command to a file")
    run_parser.add_argument('--redirect_stderr', '-e',
                            dest='redirect_stderr',
                            required=False,
                            default=None,
                            type=str,
                            help="Redirect stderr to a log file")
    run_parser.add_argument("--to",
                            dest="recipient",
                            required=False,
                            default=config["recipient"],
                            type=comma_separated_list,
                            help="A comma separted list of email recipients "
                                 "(ie '--to  user1@domain.com,user2@domain.com'). "
                                 "For AWS, this must be validated within SES console")
    run_parser.add_argument("--from",
                            dest="sender",
                            required=False,
                            default=config["sender"],
                            type=str,
                            help="The email sender. For AWS, this must be validated within SES console")
    run_parser.add_argument("--monitor",
                            required=False,
                            default=config["resource_monitor"],
                            type=bool,
                            dest="resource_monitor",
                            help="Monitor compute resources while running command. Default: True")
    run_parser.add_argument("--aws",
                            required=False,
                            default=config["aws"],
                            action='store_true',
                            help="If set, will use AWS SES if not set and source_email/source_password are not "
                                 "set will use simple local SMTP server")
    run_parser.add_argument('--output_dir',
                            dest='output_dir',
                            default=config["output_dir"],
                            type=str,
                            help="Output folder to place memory/cpu/IO usage log")
    run_parser.add_argument('--source_email',
                            dest='source_email',
                            required=False,
                            default=config["source_email"],
                            type=str,
                            help="gmail source email")
    run_parser.add_argument('--source_password',
                            dest='source_password',
                            required=False,
                            default=config["source_password"],
                            type=str,
                            help="password to source email to send email via gmail")
    run_parser.add_argument('--interval',
                            dest='interval',
                            required=False,
                            default=config["interval"],
                            type=int,
                            help="interval (in seconds) for sampling mem/cpu/IO usage")
    run_parser.add_argument('--s3_upload_interval',
                            dest='s3_upload_interval',
                            required=False,
                            default=config["s3_upload_interval"],
                            type=int,
                            help="Interval (in seconds) for uploading log to S3 bucket")
    run_parser.add_argument('--s3_upload_bucket',
                            dest='s3_upload_bucket',
                            required=False,
                            default=config["s3_upload_bucket"],
                            type=str,
                            help="bucket (s3:// is not required) for file uploading.  Setting this triggers upload.")
    run_parser.add_argument('--s3_upload_path',
                            dest='s3_upload_path',
                            required=False,
                            default=config["s3_upload_path"],
                            type=str,
                            help="s3 location (no bucket) where logs should be uploaded.  "
                                 "Can use custom python formatting parameters (need ':' prepended) including: "
                                 "'instance_id', 'timestamp', 'date'.  "
                                 "Default: 'logs/resource_monitor/{date}_{instance_id}/' ")
    args = parser.parse_args()

    if args.command == "configure":
        create_task_manager_config()
    elif args.command == "run":
        monitor = None
        assert args.sender is not None, "Must select an email sender. " \
                                        "Setup taskManager via 'taskManager configure' or set `--from` "
        assert args.recipient is not None, "Must select email recipient. " \
                                           "Setup taskManager via'taskManager configure' or set `--to` "

        if args.resource_monitor:
            # create resource monitoring class
            monitor = ResourceMonitor(output_dir=args.output_dir, interval=args.interval, aws=args.aws,
                                      s3_upload_bucket=args.s3_upload_bucket, s3_upload_path=args.s3_upload_path,
                                      s3_upload_interval=args.s3_upload_interval)
        # start process and track resources
        handler = ProcessHandler(aws=args.aws, email_sender=args.sender, email_recipients=args.recipient,
                                 source_email=args.source_email, source_password=args.source_password,
                                 resource_monitor=monitor)

        signal.signal(signal.SIGTERM, handler.handle_exit)
        signal.signal(signal.SIGINT, handler.handle_exit)

        handler.launch_process(args.run_command, redirect_stdout=args.redirect_stdout,
                               redirect_stderr=args.redirect_stderr)


if __name__ == "__main__":
    main()
