# Copyright 2017-2020 Amazon.com, Inc. or its affiliates. All Rights Reserved.
#
# Licensed under the Apache License, Version 2.0 (the "License"). You
# may not use this file except in compliance with the License. A copy of
# the License is located at
#
#     http://aws.amazon.com/apache2.0/
#
# or in the "license" file accompanying this file. This file is
# distributed on an "AS IS" BASIS, WITHOUT WARRANTIES OR CONDITIONS OF
# ANY KIND, either express or implied. See the License for the specific
# language governing permissions and limitations under the License.
"""Placeholder docstring"""
from __future__ import absolute_import

import argparse
import logging
import sys

import sagemaker
import sagemaker.cli.mxnet
import sagemaker.cli.tensorflow

logger = logging.getLogger(__name__)

DEFAULT_LOG_LEVEL = "info"
DEFAULT_BOTOCORE_LOG_LEVEL = "warning"


def parse_arguments(args):
    """
    Args:
        args:
    """
    parser = argparse.ArgumentParser(
        description="Launch SageMaker training jobs or hosting endpoints"
    )
    parser.set_defaults(func=lambda x: parser.print_usage())

    # common args for training/hosting/all frameworks
    common_parser = argparse.ArgumentParser(add_help=False)
    common_parser.add_argument(
        "--role-name", help="SageMaker execution role name", type=str, required=True
    )
    common_parser.add_argument(
        "--data", help="path to training data or model files", type=str, default="./data"
    )
    common_parser.add_argument("--script", help="path to script", type=str, default="./script.py")
    common_parser.add_argument("--job-name", help="job or endpoint name", type=str, default=None)
    common_parser.add_argument(
        "--bucket-name",
        help="S3 bucket for training/model data and script files",
        type=str,
        default=None,
    )
    common_parser.add_argument("--python", help="python version", type=str, default="py2")

    instance_group = common_parser.add_argument_group("instance settings")
    instance_group.add_argument(
        "--instance-type", type=str, help="instance type", default="ml.m4.xlarge"
    )
    instance_group.add_argument("--instance-count", type=int, help="instance count", default=1)

    # common training args
    common_train_parser = argparse.ArgumentParser(add_help=False)
    common_train_parser.add_argument(
        "--hyperparameters",
        help="path to training hyperparameters file",
        type=str,
        default="./hyperparameters.json",
    )

    # common hosting args
    common_host_parser = argparse.ArgumentParser(add_help=False)
    common_host_parser.add_argument(
        "--env", help="hosting environment variable(s)", type=str, nargs="*", default=[]
    )

    subparsers = parser.add_subparsers()

    # framework/algo subcommands
    mxnet_parser = subparsers.add_parser("mxnet", help="use MXNet", parents=[])
    mxnet_subparsers = mxnet_parser.add_subparsers()
    mxnet_train_parser = mxnet_subparsers.add_parser(
        "train", help="start a training job", parents=[common_parser, common_train_parser]
    )
    mxnet_train_parser.set_defaults(func=sagemaker.cli.mxnet.train)

    mxnet_host_parser = mxnet_subparsers.add_parser(
        "host", help="start a hosting endpoint", parents=[common_parser, common_host_parser]
    )
    mxnet_host_parser.set_defaults(func=sagemaker.cli.mxnet.host)

    tensorflow_parser = subparsers.add_parser("tensorflow", help="use TensorFlow", parents=[])
    tensorflow_subparsers = tensorflow_parser.add_subparsers()
    tensorflow_train_parser = tensorflow_subparsers.add_parser(
        "train", help="start a training job", parents=[common_parser, common_train_parser]
    )
    tensorflow_train_parser.add_argument(
        "--training-steps",
        help="number of training steps (tensorflow only)",
        type=int,
        default=None,
    )
    tensorflow_train_parser.add_argument(
        "--evaluation-steps",
        help="number of evaluation steps (tensorflow only)",
        type=int,
        default=None,
    )
    tensorflow_train_parser.set_defaults(func=sagemaker.cli.tensorflow.train)

    tensorflow_host_parser = tensorflow_subparsers.add_parser(
        "host", help="start a hosting endpoint", parents=[common_parser, common_host_parser]
    )
    tensorflow_host_parser.set_defaults(func=sagemaker.cli.tensorflow.host)

    log_group = parser.add_argument_group("optional log settings")
    log_group.add_argument(
        "--log-level", help="log level for this command", type=str, default=DEFAULT_LOG_LEVEL
    )
    log_group.add_argument(
        "--botocore-log-level",
        help="log level for botocore",
        type=str,
        default=DEFAULT_BOTOCORE_LOG_LEVEL,
    )

    return parser.parse_args(args)


def configure_logging(args):
    """
    Args:
        args:
    """
    log_format = "%(asctime)s %(levelname)s %(name)s: %(message)s"
    log_level = logging.getLevelName(args.log_level.upper())
    logging.basicConfig(format=log_format, level=log_level)
    logging.getLogger("botocore").setLevel(args.botocore_log_level.upper())


def main():
    """Placeholder docstring"""
    args = parse_arguments(sys.argv[1:])
    configure_logging(args)
    logger.debug("args: %s", args)
    args.func(args)


if __name__ == "__main__":
    main()
