#!/usr/bin/env python

from taskManager.ProcessHandler import ProcessHandler
from taskManager.utils.config 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()

    # skip parsing args if run mode is "configure"
    if len(sys.argv) > 1 and sys.argv[1] == "configure":
        create_task_manager_config()
        exit()

    # 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=load_config_argument(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=load_config_argument(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=load_config_argument(config, "resource_monitor"),
                            type=bool,
                            dest="resource_monitor",
                            help="Monitor compute resources while running command. Default: True")
    run_parser.add_argument("--attach_log",
                            required=False,
                            default=load_config_argument(config, "attach_log"),
                            type=bool,
                            dest="attach_log",
                            help="Attach full tsv log of resource usage to the email notification. Default: False")
    run_parser.add_argument("--aws",
                            required=False,
                            default=load_config_argument(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=load_config_argument(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=load_config_argument(config, "source_email"),
                            type=str,
                            help="gmail source email")
    run_parser.add_argument('--source_password',
                            dest='source_password',
                            required=False,
                            default=load_config_argument(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=load_config_argument(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=load_config_argument(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=load_config_argument(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=load_config_argument(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()

    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)

    # initialize process handler
    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,
                             attach_full_log=args.attach_log)

    # update signal handling rule with the process handler's member functions
    signal.signal(signal.SIGTERM, handler.handle_exit)
    signal.signal(signal.SIGINT, handler.handle_exit)

    # Launch process via the process handler
    handler.launch_process(args.run_command,
                           redirect_stdout=args.redirect_stdout,
                           redirect_stderr=args.redirect_stderr)


if __name__ == "__main__":
    main()
