#!/usr/bin/env python

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


def space_separated_arguments(string):
    arguments = string.strip().split(" ")
    return arguments


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 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
    # 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", "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")
    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,
                            type=str,
                            help="A comma-separted list of email recipients (ie 'user1@domain.com,user2@domain.com'). "
                                 "For AWS, this must be validated within SES console")
    run_parser.add_argument("--from",
                            dest="sender",
                            required=False,
                            type=str,
                            help="The email sender. For AWS, this must be validated within SES console")
    run_parser.add_argument("--aws",
                            required=False,
                            default=None,
                            type=string_as_bool,
                            help="Whether to use AWS SES or simple local SMTP server")
    run_parser.add_argument('--output_dir',
                            dest='output_dir',
                            required=False,
                            default=DefaultPaths["output"],
                            type=str,
                            help="Output folder to place memory/cpu/IO usage log")
    run_parser.add_argument('--source_email',
                            dest='source_email',
                            required=False,
                            default=None,
                            type=str,
                            help="gmail source email")
    run_parser.add_argument('--source_password',
                            dest='source_password',
                            required=False,
                            default=None,
                            type=str,
                            help="password to source email to send email via gmail")
    run_parser.add_argument('--interval',
                            dest='interval',
                            required=False,
                            default=None,
                            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=None,
                            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=None,
                            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="logs/resource_monitor/{date}_{instance_id}/",
                            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()
    if args.command == "run":
        # get arguments from config or command line
        config = get_task_manager_config()
        aws = select_not_none(args.aws, config.aws)
        sender = select_not_none(args.sender, config.sender)
        recipient = select_not_none(args.recipient, config.recipient)
        source_email = select_not_none(args.source_email, config.source_email)
        source_password = select_not_none(args.source_password, config.source_password)
        output_dir = select_not_none(args.output_dir, config.output_dir)
        interval = select_not_none(args.interval, config.interval)
        s3_upload_bucket = select_not_none(args.s3_upload_bucket, config.s3_upload_bucket)
        s3_upload_path = select_not_none(args.s3_upload_path, config.s3_upload_path)
        s3_upload_interval = select_not_none(args.s3_upload_interval, config.s3_upload_interval)
        # create resource monitoring class
        monitor = ResourceMonitor(output_dir=output_dir, interval=interval, aws=aws,
                                  s3_upload_bucket=s3_upload_bucket, s3_upload_path=s3_upload_path,
                                  s3_upload_interval=s3_upload_interval)
        # start process and track resources
        handler = ProcessHandler(aws=aws, email_sender=sender, email_recipients=recipient,
                                 source_email=source_email, source_password=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()
