#!/usr/bin/env python3
#                           PUBLIC DOMAIN NOTICE
#              National Center for Biotechnology Information
#  
# This software is a "United States Government Work" under the
# terms of the United States Copyright Act.  It was written as part of
# the authors' official duties as United States Government employees and
# thus cannot be copyrighted.  This software is freely available
# to the public for use.  The National Library of Medicine and the U.S.
# Government have not placed any restriction on its use or reproduction.
#   
# Although all reasonable efforts have been taken to ensure the accuracy
# and reliability of the software and data, the NLM and the U.S.
# Government do not and cannot warrant the performance or results that
# may be obtained by using this software or data.  The NLM and the U.S.
# Government disclaim all warranties, express or implied, including
# warranties of performance, merchantability or fitness for any particular
# purpose.
#   
# Please cite NCBI in any work or product based on this material.
"""
bin/elastic-blast - See DESC constant below

Author: Christiam Camacho (camacho@ncbi.nlm.nih.gov)
Created: Wed 22 Apr 2020 06:31:30 AM EDT
"""
import os
import sys
import signal
import argparse
import logging
from pprint import pformat
from typing import List
from elastic_blast import VERSION
from elastic_blast.commands.submit import submit as elb_submit
from elastic_blast.commands.status import create_arg_parser as create_status_arg_parser
from elastic_blast.commands.delete import delete as elb_delete
from elastic_blast.commands.run_summary import create_arg_parser as create_run_summary_arg_parser
from elastic_blast.util import check_positive_int, config_logging, UserReportError, SafeExecError
from elastic_blast.util import ElbSupportedPrograms, clean_up
from elastic_blast import constants
from elastic_blast.gcp import check_prerequisites
from elastic_blast.config import configure
from elastic_blast.elb_config import ElasticBlastConfig
from elastic_blast.constants import ElbCommand
from elastic_blast.constants import ELB_DFLT_LOGLEVEL, ELB_DFLT_LOGFILE
from elastic_blast.constants import CFG_CLOUD_PROVIDER, CFG_CP_GCP_PROJECT


DESC = r"""This application facilitates running BLAST on large amounts of query sequence data
on the cloud"""

# error message for missing Elastic-BLAST task on the command line
NO_TASK_MSG =\
"""Elastic-BLAST task was not specified. Please, use submit, status, delete, or run-summary.
usage: elastic-blast [-h] [--version] {submit,status,delete,run-summary} --cfg <config file> [options]"""

def main():
    """Local main entry point which sets up arguments, undo stack,
    and processes exceptions """
    try:
        signal.signal(signal.SIGINT, signal.default_int_handler)
        clean_up_stack = []
        # Check parameters for Unicode letters and reject if codes higher than 255 occur
        reject_cli_args_with_unicode(sys.argv[1:])
        parser = create_arg_parser()
        args = parser.parse_args()
        if not args.subcommand:
            # report missing command line task
            raise UserReportError(returncode=constants.INPUT_ERROR,
                                  message=NO_TASK_MSG)
        config_logging(args)
        cfg = configure(args)
        logging.info(f"ElasticBLAST {args.subcommand} {VERSION}")
        if CFG_CP_GCP_PROJECT in cfg[CFG_CLOUD_PROVIDER]:
            check_prerequisites()
        task = ElbCommand(args.subcommand.lower())
        cfg = ElasticBlastConfig(cfg, args.dry_run, task=task)
        logging.debug(pformat(cfg.asdict()))
        #TODO: use cfg only when args.wait, args.sync, and args.run_label are replicated in cfg
        return args.func(args, cfg, clean_up_stack)
    except (SafeExecError, UserReportError) as e:
        logging.error(e.message)
        if 'ELB_DEBUG' in os.environ:
            import traceback
            traceback.print_exc(file=sys.stderr)
        # SafeExecError return code is the exit code from command line
        # application ran via subprocess
        if isinstance(e, SafeExecError):
            return constants.DEPENDENCY_ERROR
        return e.returncode
    except KeyboardInterrupt:
        return constants.INTERRUPT_ERROR
    #TODO: process filehelper.TarReadError here
    finally:
        messages = clean_up(clean_up_stack)
        if messages:
            for msg in messages:
                logging.error(msg)
            sys.exit(constants.UNKNOWN_ERROR)


def reject_string_with_unicode(content: str) -> None:
    for c in content:
        if ord(c) > 255:
            raise UserReportError(returncode=constants.INPUT_ERROR,
                                  message=f"Command line has Unicode letters in argument '{content}', can't be processed")


def reject_cli_args_with_unicode(args: List[str]) -> None:
        for arg in args:
            reject_string_with_unicode(arg)


def file_must_exist(path: str) -> str:
    """Check if given  path exists and is a file, helper function for
    argparse.ArgumentParser. If used for a command line argument, the
    application will exit with an error if file is not found or path is not a
    file."""
    if os.path.isfile(path):
        return path
    raise argparse.ArgumentTypeError(f'File {path} was not found')


def positive_int(arg: str) -> int:
    """Positive integer type for argparse.ArgumentParser. Raises
    argparse.ArgumentTypeError if the supplied string is not a positive integer."""
    retval = None
    try:
        retval = check_positive_int(arg)
    except Exception as err:
        raise argparse.ArgumentTypeError(str(err))
    return retval


def create_arg_parser():
    """ Create the command line options parser object for this script. """
    prog = os.path.splitext(os.path.basename(sys.argv[0]))[0]
    parser = ElbArgumentParser(prog=prog, description=DESC,
        epilog="To get help about specific command run %(prog)s command --help")
    parser.add_argument('--version', action='version', version='%(prog)s ' + VERSION)

    common_opts_parser = ElbArgumentParser(add_help=False)

    csp_opts = common_opts_parser.add_argument_group('Cloud Service Provider options')
    csp_opts.add_argument("--aws-region", help="AWS region to run ElasticBLAST")
    csp_opts.add_argument("--gcp-project", help="GCP project to run ElasticBLAST")
    csp_opts.add_argument("--gcp-region", help="GCP region to run ElasticBLAST")
    csp_opts.add_argument("--gcp-zone", help="GCP zone to run ElasticBLAST")

    elb_opts = common_opts_parser.add_argument_group('ElasticBLAST configuration options')
    elb_opts.add_argument("--cfg", metavar="FILE",
                                    type=file_must_exist,
                                    help="ElasticBLAST configuration file")
    elb_opts.add_argument(f"--{constants.CFG_BLAST_RESULTS}", type=str,
                        help="Bucket URI where to save the output from ElasticBLAST")

    app_opts = common_opts_parser.add_argument_group('Application options')
    app_opts.add_argument("--logfile", default=argparse.SUPPRESS, type=str,
                                    help=f"Default: {ELB_DFLT_LOGFILE}")
    app_opts.add_argument("--loglevel", default=argparse.SUPPRESS,
                                    help=f"Default: {ELB_DFLT_LOGLEVEL}",
                                    choices=["DEBUG", "INFO", "WARNING", "ERROR", "CRITICAL"])
    app_opts.add_argument("--dry-run", action='store_true', 
                                    help="Do not perform any actions")

    sp = parser.add_subparsers(dest='subcommand')#title='Subcommands')
    create_submit_arg_parser(sp, common_opts_parser)
    create_status_arg_parser(sp, common_opts_parser)
    create_delete_arg_parser(sp, common_opts_parser)
    create_run_summary_arg_parser(sp, common_opts_parser)
    return parser


def create_submit_arg_parser(subparser, common_opts_parser):
    """ Create the command line options subparser for the submit command. """
    parser = subparser.add_parser('submit', help='Submit an ElasticBLAST search',
                                  parents=[common_opts_parser])
    # BLAST configuration parameters
    blast_cli_opts = parser.add_argument_group('BLAST options')
    blast_cli_opts.add_argument("--program", type=str, help="BLAST program to run",
                        choices=ElbSupportedPrograms().get())
    blast_cli_opts.add_argument("--query", type=str,
                        help="Query sequence data, can be provided as a local path or GCS bucket URI to a single file/tarball")
    blast_cli_opts.add_argument("--db", type=str, help="BLAST database to search")
    blast_cli_opts.add_argument('blast_opts', nargs=argparse.REMAINDER,
                        metavar='BLAST_OPTS',
                        help="Options to pass to BLAST program")

    # ElasticBLAST search configuration parameters
    elb_opts = parser.add_argument_group('ElasticBLAST configuration options')
    elb_opts.add_argument("--batch-len", type=positive_int,
                        help="Query size for each BLAST job")
    elb_opts.add_argument("--machine-type", type=str,
                        help="Instance type to use")
    elb_opts.add_argument("--num-nodes", type=positive_int,
                        help="Number of worker nodes to use")
    elb_opts.add_argument("--num-cpus", type=positive_int,
                        help="Number of threads to run in each BLAST job")
    elb_opts.add_argument("--mem-limit", type=str,
                        help="Memory limit for each BLAST job")

    #parser.add_argument("--sync", action='store_true', 
    #                    help="Run in synchronous mode")
    # FIXME: EB-132
    parser.add_argument("--run-label", type=str,
                        help="Run-label to tag this ElasticBLAST search, format: key:value")
    parser.set_defaults(func=elb_submit)
    return parser


def create_delete_arg_parser(subparser, common_opts_parser):
    """ Create the command line options subparser for the status command. """
    parser = subparser.add_parser('delete',
                                  parents=[common_opts_parser],
                                  help='Delete resources associated with an ElasticBLAST search')
    # FIXME: EB-132
    parser.add_argument("--run-label", type=str,
                        help="Run-label for the ElasticBLAST search to delete, format: key:value")
    parser.set_defaults(func=elb_delete)


class ElbArgumentParser(argparse.ArgumentParser):
    """Custom argument parser to override application exit code"""
    def exit(self, status=0, message=None):
        """Custom exit function that overrides ArgumentParser application
        exit code"""
        if status:
            super().exit(constants.INPUT_ERROR, message)
        else:
            super().exit()

    def error(self, message):
        """Custom error message that does not print usage on errors"""
        self.exit(constants.INPUT_ERROR, f'{self.prog}: error: {message}\n')


if __name__ == "__main__":
    sys.exit(main())
    import traceback
    try:
        sys.exit(main())
    except Exception as e:
        traceback.print_exc(file=sys.stderr)
        sys.exit(constants.UNKNOWN_ERROR)
