#!/usr/bin/env python
#
#    Copyright 2018 EPAM Systems
#
#    Licensed under the Apache License, Version 2.0 (the "License");
#    you may not use this file except in compliance with the License.
#    You may obtain a copy of the License at
#
#        http://www.apache.org/licenses/LICENSE-2.0
#
#    Unless required by applicable law or agreed to in writing, software
#    distributed under the License 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.
#
import sys
import os
import argparse
import logging

import legion.docker_bootup
from legion.serving.pyserve import serve_model
from legion.external.edi import add_edi_arguments
from legion.edi.deploy import \
    build_model, \
    deploy_kubernetes, undeploy_kubernetes, scale_kubernetes, inspect_kubernetes, \
    VALID_INSPECT_FORMATS
import legion.utils

ROOT_LOGGER = logging.getLogger()

if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='legion Command-Line Interface')
    parser.add_argument('--verbose',
                        help='verbose log output',
                        action='store_true')
    subparsers = parser.add_subparsers()

    # --------- LOCAL DOCKER SECTION -----------
    build_parser = subparsers.add_parser('build', description='build model into new docker image (should be run '
                                                              'in the docker container)')
    build_parser.add_argument('model_file',
                              type=str, help='serialized model file name')
    build_parser.add_argument('--docker-image-tag',
                              type=str, help='docker image tag')
    build_parser.add_argument('--push-to-registry',
                              type=str, help='docker registry address')
    build_parser.set_defaults(func=build_model)

    # --------- KUBERNETES SECTION -----------
    deploy_k8s_parser = subparsers.add_parser('deploy',
                                              description='deploys a model into a kubernetes cluster')
    deploy_k8s_parser.add_argument('image',
                                   type=str, help='docker image')
    deploy_k8s_parser.add_argument('--image-for-k8s',
                                   type=str, help='docker image for kubernetes deployment')
    deploy_k8s_parser.add_argument('--scale',
                                   default=1,
                                   type=int, help='count of instances')
    deploy_k8s_parser.add_argument('--livenesstimeout',
                                   default=2,
                                   type=int, help='model startup timeout for liveness probe')
    deploy_k8s_parser.add_argument('--readinesstimeout',
                                   default=2,
                                   type=int, help='model startup timeout for readiness probe')
    deploy_k8s_parser.add_argument('--no-wait',
                                   action='store_true', help='no wait until pods will be ready')
    deploy_k8s_parser.add_argument('--timeout',
                                   default=300,
                                   type=int, help='timeout in s. for wait (if no-wait is off)')
    add_edi_arguments(deploy_k8s_parser)
    deploy_k8s_parser.set_defaults(func=deploy_kubernetes)

    inspect_k8s_parser = subparsers.add_parser('inspect',
                                               description='get information about currently deployed models')
    inspect_k8s_parser.add_argument('--model-id',
                                    type=str, help='model ID')
    inspect_k8s_parser.add_argument('--model-version',
                                    type=str, help='model version')
    inspect_k8s_parser.add_argument('--format',
                                    default=VALID_INSPECT_FORMATS[0],
                                    choices=VALID_INSPECT_FORMATS, help='output format')
    add_edi_arguments(inspect_k8s_parser)
    inspect_k8s_parser.set_defaults(func=inspect_kubernetes)

    scale_k8s_parser = subparsers.add_parser('scale',
                                             description='change count of model pods')
    scale_k8s_parser.add_argument('model_id',
                                  type=str, help='model ID')
    scale_k8s_parser.add_argument('scale',
                                  type=int, help='new count of replicas')
    scale_k8s_parser.add_argument('--model-version',
                                  type=str, help='model version')
    scale_k8s_parser.add_argument('--no-wait',
                                  action='store_true', help='no wait until scale will be finished')
    scale_k8s_parser.add_argument('--timeout',
                                  default=300,
                                  type=int, help='timeout in s. for wait (if no-wait is off)')
    add_edi_arguments(scale_k8s_parser)
    scale_k8s_parser.set_defaults(func=scale_kubernetes)

    undeploy_k8s_parser = subparsers.add_parser('undeploy',
                                                description='undeploy model deployment')
    undeploy_k8s_parser.add_argument('model_id',
                                     type=str, help='model ID')
    undeploy_k8s_parser.add_argument('--model-version',
                                     type=str, help='model version')
    undeploy_k8s_parser.add_argument('--grace-period',
                                     default=0,
                                     type=int, help='removal grace period')
    undeploy_k8s_parser.add_argument('--ignore-not-found',
                                     action='store_true', help='ignore if cannot found pod')
    undeploy_k8s_parser.add_argument('--no-wait',
                                     action='store_true', help='no wait until pods will be removed')
    undeploy_k8s_parser.add_argument('--timeout',
                                     default=300,
                                     type=int, help='timeout in s. for wait (if no-wait is off)')
    add_edi_arguments(undeploy_k8s_parser)
    undeploy_k8s_parser.set_defaults(func=undeploy_kubernetes)

    # --------- SERVING SECTION -----------
    pyserve_parser = subparsers.add_parser('pyserve', description='serve a python model')
    pyserve_parser.add_argument('--model_file',
                                type=str)
    pyserve_parser.add_argument('--model-id',
                                type=str)
    pyserve_parser.add_argument('--legion-addr',
                                type=str)
    pyserve_parser.add_argument('--legion-port',
                                type=int)
    pyserve_parser.add_argument('--debug',
                                type=legion.utils.string_to_bool)
    pyserve_parser.set_defaults(func=serve_model)

    # --------- END OF SECTIONS -----------
    args = parser.parse_args(sys.argv[1:])

    v = vars(args)

    if args.verbose or legion.utils.string_to_bool(os.getenv('VERBOSE', '')):
        log_level = logging.DEBUG
    else:
        log_level = logging.ERROR

    logging.basicConfig(level=log_level, 
                        format='%(asctime)s - %(levelname)s - %(message)s',
                        stream=sys.stderr)

    try:
        if 'func' in v:
            args.func(args)
        else:
            parser.print_help()
            sys.exit(1)
    except Exception as exception:
        ROOT_LOGGER.exception(exception)
        sys.exit(2)
