#!/usr/bin/env python3
# -*- coding: utf-8 -*-

# Licensed to the Apache Software Foundation (ASF) under one
# or more contributor license agreements.  See the NOTICE file
# distributed with this work for additional information
# regarding copyright ownership.  The ASF licenses this file
# to you 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.

"""
Tool for uploading artifacts to the artifact repository
"""

__author__ = 'Per Goncalves da Silva'
__version__ = '0.1'

import argparse
import ctypes
import glob
import logging
import os
import re
import sys
from itertools import chain
from subprocess import CalledProcessError, check_output
from typing import Dict, List, Optional

import boto3
import botocore.exceptions
import yaml

s3 = boto3.client('s3')
logger = logging.getLogger(__name__)


def config_logging():
    """
    Configures default logging settings
    """
    logging.root.setLevel(logging.WARNING)
    logger.setLevel(logging.INFO)
    logging.basicConfig(format='%(levelname)s: %(message)s')


def s3_upload(bucket: str, s3_key_prefix: str, paths: List[str]):
    """
    Uploads a list of files to an S3 bucket with a particular S3 key prefix
    :param bucket: The name of the S3 bucket
    :param s3_key_prefix: The key prefix to apply to each of the files
    :param paths: A list of paths to files
    """
    for path in paths:
        s3_key = "{}/{}".format(s3_key_prefix, os.path.basename(path))
        logger.info('Uploading {}'.format(path))
        logger.debug("Uploading {} to s3://{}/{}".format(path, bucket, s3_key))
        with open(path, 'rb') as data:
            s3.upload_fileobj(Fileobj=data, Key=s3_key, Bucket=bucket)


def write_libmxnet_meta(args: argparse.Namespace, destination: str):
    """
    Writes a file called libmxnet.meta in the 'destination' folder that contains
    the libmxnet library information (commit id, type, etc.).
    :param args: A Namespace object containing the library
    :param destination: The folder in which to place the libmxnet.meta
    """
    with open(os.path.join(destination, 'libmxnet.meta'), 'w') as fp:
        fp.write(yaml.dump({
            "variant": args.variant,
            "os": args.os,
            "commit_id": args.git_sha,
            "dependency_linking": args.libtype,
        }))


def try_s3_download(bucket, s3_key_prefix, destination) -> bool:
    """
    Downloads a list of files to an S3 bucket with a particular S3 key prefix to 'destination'
    :param bucket: The name of the S3 bucket
    :param s3_key_prefix: The key prefix to apply to each of the files
    :param destination the path to which to download the files
    :return False if not artifacts were found, True otherwise
    """
    response = s3.list_objects_v2(Bucket=bucket, Prefix=s3_key_prefix)
    if not response:
        raise RuntimeError('Error listing S3 objects')

    if response.get('KeyCount') is None:
        logger.debug('Invalid S3 list objects response format')
        logger.debug(response)
        raise RuntimeError('Invalid response format.')

    key_count = response.get('KeyCount')
    if key_count == 0:
        logger.debug('No artifacts found')
        return False

    if not response.get('Contents'):
        logger.debug('Invalid S3 list objects response format')
        logger.debug(response)
        raise RuntimeError('Invalid response format.')

    for obj in response.get('Contents'):
        key = obj['Key']

        # extract file path with any subdirectories and remove the leading file separator
        output_path = os.path.join(destination, key[len(s3_key_prefix):].lstrip(os.sep))
        os.makedirs(os.path.dirname(output_path), exist_ok=True)
        logger.info('Downloading {}'.format(output_path))
        logger.debug("Downloading s3://{}/{} to {}".format(bucket, key, output_path))
        with open(output_path, 'wb') as fp:
            s3.download_fileobj(Fileobj=fp, Key=key, Bucket=bucket)

    return True


def get_commit_id_from_cmd() -> Optional[str]:
    """
    Returns the output of 'git rev-parse HEAD'
    :return: A commit id, or None if the command fails
    """
    try:
        logger.debug('Executing "git rev-parse HEAD"')
        commit_id = check_output("git rev-parse HEAD".split(" ")).decode('UTF-8').strip()
        logger.debug('Found commit id: {}'.format(commit_id))
        return commit_id
    except CalledProcessError as e:
        logger.debug('Error getting commit id:')
        logger.debug(format(e))
        return None


def probe_commit_id() -> str:
    """
    Probes the system in an attempt to ascertain the mxnet commit id
    :return: The commit id, or None if not found
    """
    logger.debug('Probing for commit id')
    commit_id = os.environ.get('MXNET_SHA')
    if not commit_id:
        logger.debug('MXNET_SHA environment variable not set. Trying GIT_COMMIT')
        commit_id = os.environ.get('GIT_COMMIT')
    if not commit_id:
        logger.debug('GIT_COMMIT environment variable not set. Trying git command')
        commit_id = get_commit_id_from_cmd()
    if not commit_id:
        logger.debug('Could not determine git commit id')
    else:
        logger.debug('Commit id is: {}'.format(commit_id))
    return commit_id


def get_linux_os_release_properties() -> Optional[Dict[str, str]]:
    """
    Makes a dictionary out of /etc/os-release
    :return: A dictionary of os release properties
    """
    logger.debug('Extracting operating system properties from /etc/os-release')
    if not os.path.isfile('/etc/os-release'):
        logger.debug('Error: /etc/os-release not found')
        return None

    try:
        with open('/etc/os-release', 'r') as fp:
            # removes empty spaces and quotation marks from line
            property_tuple_list = [line.strip().replace('"', '').split('=') for line in fp if line.strip()]
            return {key: value for (key, value) in property_tuple_list}
    except Exception as e:
        logger.debug('Error parsing /etc/os-release')
        logger.debug(e)
        return None


def get_linux_distribution_and_version() -> Optional[str]:
    """
    Returns the linux distribution and version by taking
    the values of ID and VERSION_ID from /etc/os-release and
    concatenating them. Eg. centos7, ubuntu16.04, etc.
    :return: The linux distribution and version string, or None if not found.
    """
    logger.debug('Getting linux distribution and version')
    os_properties = get_linux_os_release_properties()
    if os_properties:
        logger.debug('os properties: {}'.format(os_properties))
        distribution = os_properties['ID']
        version = os_properties['VERSION_ID']
        return "{}{}".format(distribution, version)

    logger.debug('Error getting linux distribution and version. Could not determine os properties.')
    return None


def probe_operating_system() -> str:
    """
    Probes the system to determine the operating system.
    :return: The name of the operating system, e.g. win32, darwin, ubuntu16.04, centos7, etc.
    """
    logger.debug('Determining operating system')
    operating_system = sys.platform
    logger.debug('Found platform: {}'.format(operating_system))
    if operating_system.startswith('linux'):
        operating_system = get_linux_distribution_and_version()

    logger.debug('Operating system is {}'.format(operating_system))
    return operating_system


def get_libmxnet_features(libmxnet_path: str) -> Optional[Dict[str, bool]]:
    """
    Returns a string -> boolean dictionary mapping feature name
    to whether it is enabled or not
    :param libmxnet_path: path to the libmxnet library
    :return: dictionary of features to whether they are enabled
    """
    logger.debug('Getting feature dictionary from {}'.format(libmxnet_path))

    class Feature(ctypes.Structure):
        _fields_ = [("_name", ctypes.c_char_p), ("enabled", ctypes.c_bool)]

        @property
        def name(self):
            return self._name.decode()

    # we are not using the mxnet python bindings here because we cannot assume
    # they are present and in the python path, or that they would point to the
    # specified libmxnet.so. Therefore, we load the libmxnet.so library independently
    # to extract its features
    try:
        libmxnet = ctypes.CDLL(libmxnet_path, ctypes.RTLD_LOCAL)
    except Exception as e:
        logger.error('Error loading {}. '
                     'Please check check path to libmxnet library is correct.'.format(libmxnet_path))
        logger.error(e)
        return None

    libmxnet.MXGetLastError.restype = ctypes.c_char_p
    feature_array = ctypes.POINTER(Feature)()
    feature_array_size = ctypes.c_size_t()
    if libmxnet.MXLibInfoFeatures(ctypes.byref(feature_array), ctypes.byref(feature_array_size)) != 0:
        logger.error('Could not determine features from {}. '
                     'Please specify the variant manually using the "--variant" argument.'.format(libmxnet_path))
        return None
    features = {feature_array[i].name: feature_array[i].enabled for i in range(feature_array_size.value)}
    logger.debug('Found features: {}'.format(features))
    return features


def get_cuda_version() -> Optional[str]:
    """
    Returns the major and minor cuda version without the '.'
    eg. 10.0 => 100, 9.2 => 92, etc.
    :return: CUDA version
    """
    logger.debug('Determining cuda version')

    try:
        logger.debug('Executing "nvcc -V"')
        nvcc_version = check_output("nvcc -V".split(" ")).decode('UTF-8').strip()
    except CalledProcessError as e:
        logger.error('Error getting nvcc version')
        logger.error(e)
        return None

    logger.debug('Extracting cuda version from {}'.format(nvcc_version))
    # eg. "Cuda compilation tools, release 10.0, V10.0.130"
    match = re.search(r' ([0-9]+.[0-9]+)', nvcc_version)
    if match:
        cuda_version = match.group(1).replace('.', '')
        logger.debug('Found cuda version: {}'.format(cuda_version))
        return cuda_version

    logger.debug('Could not determine cuda version from "{}"'.format(nvcc_version))
    return None


def probe_cpu_variant(mxnet_features: Dict[str, bool]) -> str:
    """
    Returns the mxnet cpu targeted variant depending on which mxnet features are enabled
    :param mxnet_features: An mxnet feature dictionary of feature to boolean (True = enabled)
    :return: Either cpu, or mkl as the variant
    """
    logger.debug('Determining cpu variant')
    if not mxnet_features['ONEDNN']:
        logger.debug('variant is: native')
        return 'native'

    logger.debug('variant is: cpu')
    return 'cpu'


def probe_gpu_variant(mxnet_features: Dict[str, bool]) -> Optional[str]:
    """
    Returns the mxnet gpu variant depending on which mxnet features are enabled
    :param mxnet_features: An mxnet feature dictionary of feature to boolean (True = enabled)
    :return: The mxnet gpu variant, eg. cu102, cu102mkl, etc.
    :raises RuntimeError is the CUDA feature is not enabled in the library
    """
    if not mxnet_features['CUDA']:
        raise RuntimeError('Cannot determine gpu variant. CUDA feature is disabled.')

    cuda_version = get_cuda_version()
    if cuda_version:
        variant = 'cu{}'.format(cuda_version)
        if not mxnet_features['ONEDNN']:
            RuntimeError('Error determining mxnet variant: oneDNN should be enabled for cuda variants')
        logger.debug('variant is: {}'.format(variant))
        return variant

    raise RuntimeError('Error determining mxnet variant: Could not retrieve cuda version')


def probe_mxnet_variant(limxnet_path: str) -> Optional[str]:
    """
    Probes the libmxnet library and environment to determine
    the mxnet variant, eg. cpu, cu102, etc.
    :return:
    """
    logger.debug('Probing for mxnet variant')
    features = get_libmxnet_features(limxnet_path)
    if not features:
        logger.debug('Error: could not determine variant. Features could not be extracted from libmxnet')
        return None

    if features['CUDA']:
        return probe_gpu_variant(features)
    return probe_cpu_variant(features)


def probe_artifact_repository_bucket() -> Optional[str]:
    """
    Probes environment variables in search of artifact repository bucket
    :return: string containing the artifact repository bucket name
    """
    logger.debug('Probing for artifact repository bucket name')
    bucket = os.environ.get('ARTIFACT_REPOSITORY_BUCKET')
    if not bucket:
        logger.debug('ARTIFACT_REPOSITORY_BUCKET environment variable not found')
    return bucket


def probe(args: argparse.Namespace) -> argparse.Namespace:
    """
    Probes the system to set any arguments that weren't manually set.
    Modifies the input Namespace object with the probed parameters.
    :param args: The namespace object given by argparse.parse()
    """
    logger.debug('Trying to auto-determine arguments from environment')
    if not args.git_sha:
        commit_id = probe_commit_id()
        if not commit_id:
            logger.error('Could not determine commit id. '
                         'Please set it manually with --git-sha, or ensure you are in a cloned '
                         'mxnet repository directory')
            sys.exit(1)
        args.git_sha = commit_id

    if not args.variant:
        variant = probe_mxnet_variant(args.libmxnet)
        if not variant:
            logger.error('Could not determine mxnet variant. Please set it manually with --variant')
            sys.exit(1)
        args.variant = variant

    if not args.os:
        operating_system = probe_operating_system()
        if not operating_system:
            logger.error('Could not determine operating system. Please set it manually with --os')
            sys.exit(1)
        args.os = operating_system

    if not args.bucket:
        artifact_repo_bucket = probe_artifact_repository_bucket()
        if not artifact_repo_bucket:
            logger.error('Could not determine artifact repository bucket. Please set it manually with --bucket')
            sys.exit(1)
        args.bucket = artifact_repo_bucket

    return args


def get_s3_key_prefix(args: argparse.Namespace, subdir: str = '') -> str:
    """
    Returns the S3 key prefix given the arguments namespace
    :param args: The arguments passed in by the user or derived by the script
    :param subdir: An optional subdirectory in which to store the files. Post-pended to the end of the prefix.
    :return: A string containing the S3 key prefix to be used to uploading and downloading files to the artifact repository
    """
    prefix = "{git_sha}/{libtype}/{os}/{variant}/".format(**vars(args))
    if subdir:
        return "{}{}/".format(prefix, subdir)
    return prefix


def push_artifact(args: argparse.Namespace):
    """
    Pushes the artifact to the artifact repository
    :param args: The arguments passed in to this script by the user
    :return 0 for success, non-zero for failure
    """

    args = probe(args)

    logger.info('Pushing artifact with: ')
    logger.info('COMMIT ID   : {}'.format(args.git_sha))
    logger.info('OS          : {}'.format(args.os))
    logger.info('VARIANT     : {}'.format(args.variant))
    logger.info("LIBMXNET    : {}".format(args.libmxnet))
    logger.info("LICENSES    : {}".format(args.licenses))
    logger.info("DEPENDENCIES: {}".format(args.dependencies))
    logger.info("")

    if not args.licenses:
        raise RuntimeError('No licenses defined. Please submit the licenses to be shipped with the binary.')

    # Upload mxnet
    try:
        logger.info('Uploading libmxnet library...')
        s3_upload(args.bucket, get_s3_key_prefix(args), [args.libmxnet])
        logger.info("")

        # Upload licenses
        logger.info('Uploading licenses...')
        s3_upload(args.bucket, get_s3_key_prefix(args, subdir='licenses'), args.licenses)
        logger.info("")

        # Upload dependencies, if necessary
        if args.dependencies:
            logger.info('Uploading dependencies...')
            s3_upload(args.bucket, get_s3_key_prefix(args, subdir='dependencies'), args.dependencies)
            logger.info("")
    except botocore.exceptions.BotoCoreError as e:
        logger.error('Error uploading artifact')
        logger.error(e)
        raise e

    logger.info('Successfully pushed artifact')


def pull_artifact(args: argparse.Namespace):
    """
    Pulls the artifact from the artifact repository
    :param args: The arguments passed in to this script by the user
    :return 0 for success, 1 for unexpected failure, 2 for no artifact found failure
    """
    if not args.variant:
        logger.warning('''variant not set. Using 'cpu' by default.''')
        args.variant = 'cpu'

    args = probe(args)

    logger.info('Pulling artifact with: ')
    logger.info('COMMIT ID   : {}'.format(args.git_sha))
    logger.info('OS          : {}'.format(args.os))
    logger.info('VARIANT     : {}'.format(args.variant))
    logger.info('To directory: {}'.format(args.destination))

    try:
        if not try_s3_download(args.bucket, get_s3_key_prefix(args), args.destination):
            raise RuntimeError('No artifacts found for this configuration.')
        write_libmxnet_meta(args=args, destination=args.destination)
    except botocore.exceptions.BotoCoreError as e:
        logger.error('Error downloading artifact')
        logger.error(e)
        raise e

    logger.info('Successfully pulled artifact')


def is_file(path: str) -> str:
    """
    Returns true or false if path points to an existing file
    :param path: A path to a file
    :return: True if file exists and is a file, False otherwise
    :raises FileNotFoundError if file does not exist
    """
    if not os.path.exists(path):
        raise FileNotFoundError('''File '{}' not found'''.format(path))
    return os.path.isfile(path)


def sanitize_path_array(paths: List[str]) -> List[str]:
    """
    Expands supplied paths and removes empty or non-file entries.
    :param paths: A list of paths
    :return: A sanitized list of paths
    :raises FileNotFoundError if a file does not exist
    """
    expanded_paths = list(chain.from_iterable(glob.glob(path.strip()) for path in paths if path.strip() != ''))
    return [path.strip() for path in expanded_paths if path.strip() != '' and is_file(path)]


def main() -> int:
    config_logging()

    logger.info("MXNet-CD Artifact Repository Tool")

    parser = argparse.ArgumentParser(description="Utility for uploading and downloading MXNet artifacts")

    parser.add_argument("--push",
                        help="Upload artifact to repository",
                        required=False,
                        action='store_true')

    parser.add_argument("--pull",
                        help="Download artifact from repository",
                        required=False,
                        action='store_true')

    parser.add_argument("--libmxnet",
                        help="Path to libmxnet library",
                        required=False,
                        type=str)

    parser.add_argument("--licenses",
                        help="Paths to license files",
                        required=False,
                        nargs=argparse.ZERO_OR_MORE,
                        default=[])

    parser.add_argument("--dependencies",
                        help="Paths to dependencies",
                        required=False,
                        nargs=argparse.ZERO_OR_MORE,
                        default=[])

    parser.add_argument("--os",
                        help="Target operating system",
                        type=str)

    parser.add_argument("--git-sha",
                        help="MXNet repository commit id",
                        required=False,
                        type=str)

    parser.add_argument("--variant",
                        help="MXNet binary variant. Eg. cpu, native, cu102, etc.",
                        required=False,
                        type=str)

    parser.add_argument("--libtype",
                        help="libmxnet dependency linking type",
                        choices=['static', 'dynamic'],
                        default='dynamic',
                        required=False)

    parser.add_argument('--bucket',
                        help="S3 bucket to store files",
                        type=str,
                        required=False)

    parser.add_argument('--destination',
                        help="Destination for downloaded library and supporting files",
                        type=str,
                        default=os.getcwd(),
                        required=False)

    parser.add_argument('--verbose',
                        help='Verbose',
                        action='store_true',
                        default=False)

    parser.add_argument('--debug',
                        help='Debug mode',
                        action='store_true',
                        default=False)

    args = parser.parse_args()

    if args.verbose:
        logger.setLevel(logging.DEBUG)

    if args.debug:
        # Set debug level on root logger (ie. for all other loggers)
        logging.getLogger().setLevel(logging.DEBUG)

    if not args.push and not args.pull:
        logger.info('''Mode not specified. Using 'push' by default.''')
        args.push = True

    # libmxnet argument is required for push mode
    if args.push and not args.libmxnet:
        logger.error('Path to libmxnet library must be specified when in push mode. '
                     'Please specify it with --libmxnet.')
        return 1

    # sanitize license and dependency arrays
    # Remove empty or directory entries
    args.licenses = sanitize_path_array(args.licenses)
    args.dependencies = sanitize_path_array(args.dependencies)

    # expand destination path
    args.destination = os.path.abspath(args.destination)

    try:
        if args.push:
            push_artifact(args)

        elif args.pull:
            pull_artifact(args)
    except RuntimeError as err:
        logger.error(err)
        return 1

    return 0


if __name__ == '__main__':
    sys.exit(main())
