#!/usr/bin/env python
"""Script to benchmark several high level cli commands.

As of now this benchmarks `cp` and `rm` with test cases for multiple 4kb files
(default 10000 files) and a single large file (default 10gb, `cp` only).
"""
import os
import json
from subprocess import check_call, Popen, PIPE
from datetime import datetime
import random
import argparse
import inspect
import shutil
import platform

import awscli
import s3transfer


TEST_BUCKET = os.environ.get('PERF_TEST_BUCKET')
REPO_ROOT = os.path.dirname(
    os.path.dirname(os.path.dirname(os.path.abspath(__file__))))
WORKDIR = os.environ.get('PERF_WORKDIR', os.path.join(REPO_ROOT, 'workdir'))
MANY_FILES_DIR = 'many'
LARGE_FILE_DIR = 'large'


def run(command):
    return check_call(command, shell=True)


def generate_run_id():
    run_id = datetime.now().strftime("%Y-%m-%d-%H-%M-%S-")
    run_id += str(random.randint(1, 10000))
    return run_id


def initialize_files(num_files, file_size):
    # TODO: We probably need to recreate these files each time.
    # Because you can specify --num-files and --large-file-size
    # those arguments can be potentially ignored if you've run
    # this previously with different values.
    many_files_dir = os.path.join(WORKDIR, MANY_FILES_DIR)
    if not os.path.exists(many_files_dir):
        os.makedirs(many_files_dir)
        run('caf gen --file-size 4kb --max-files %s --directory %s' %
            (num_files, many_files_dir))

    large_file_dir = os.path.join(WORKDIR, LARGE_FILE_DIR)
    if not os.path.exists(large_file_dir):
        os.makedirs(large_file_dir)
        run('caf gen --file-size %s --max-files 1 --directory %s' %
            (file_size, large_file_dir))


def write_metadata_file(filename):
    metadata = _collect_metadata()
    with open(filename, 'w') as f:
        f.write(json.dumps(metadata, indent=2))


def _collect_metadata():
    # We want to track things like the machine where the benchmark was run,
    # and which version of the aws-cli/s3transfer code is being run.
    # This helps us make more meaningful comparison.
    metadata = {
        'python_version': platform.python_version(),
        'os': '%s/%s' % (platform.system(), platform.release()),
    }
    _inject_package_info(awscli, metadata)
    _inject_package_info(s3transfer, metadata)
    return metadata


def _inject_package_info(package, metadata):
    name = package.__name__
    metadata[name + '_version'] = package.__version__
    metadata[name + '_git_version'] = _get_git_version(package)


def _get_git_version(package):
    dname = os.path.dirname(inspect.getfile(package))
    git_sha = Popen(
        'git rev-parse HEAD',
        cwd=dname, shell=True, stdout=PIPE).communicate()[0].strip()
    git_branch = Popen(
        'git rev-parse --abbrev-ref HEAD',
        cwd=dname, shell=True, stdout=PIPE).communicate()[0].strip()
    return '%s (%s)' % (git_sha, git_branch)


def main(args):
    initialize_files(args.num_files, args.large_file_size)
    run_id = generate_run_id()
    results_dir = os.path.join(WORKDIR, 'results', run_id)
    os.makedirs(results_dir)
    write_metadata_file(os.path.join(results_dir, 'metadata.json'))
    try:
        benchmark(args.bucket, results_dir, args.num_iterations)
        print("RUN ID: " + run_id)
    except Exception:
        shutil.rmtree(results_dir)
        raise


def benchmark(bucket, results_dir, num_iterations=1):
    perf_dir = os.path.dirname(os.path.dirname(os.path.realpath(__file__)))
    perf_dir = os.path.join(perf_dir, 'performance')

    s3_location = bucket + '/' + MANY_FILES_DIR
    local_dir = os.path.join(WORKDIR, MANY_FILES_DIR)
    try:
        # 10k upload
        results = os.path.join(results_dir, 'upload-10k-small')
        os.makedirs(results)
        benchmark_cp = os.path.join(perf_dir, 'benchmark-cp')
        run(benchmark_cp + ' --recursive --num-iterations %s '
            '--source %s --dest %s --result-dir %s --no-cleanup' % (
                num_iterations, local_dir, s3_location, results))

        # 10k download
        results = os.path.join(results_dir, 'download-10k-small')
        os.makedirs(results)
        run(benchmark_cp + ' --recursive --num-iterations %s '
            '--source %s --dest %s --result-dir %s' % (
                num_iterations, s3_location, local_dir, results))

        # 10k rm
        results = os.path.join(results_dir, 'delete-10k-small')
        os.makedirs(results)
        benchmark_rm = os.path.join(perf_dir, 'benchmark-rm')
        run(benchmark_rm + ' --recursive --num-iterations %s '
            '--target %s --result-dir %s' % (
                num_iterations, s3_location, results))
    finally:
        # Note that the delete-10k-small benchmark restores
        # the files it's deleted once the script is finished.
        # Therefore we need to explicitly cleanup any files
        # we've created.
        run('aws s3 rm --recursive ' + s3_location)

    s3_location = bucket + '/' + LARGE_FILE_DIR
    local_dir = os.path.join(WORKDIR, LARGE_FILE_DIR)
    try:
        # 10gb upload
        results = os.path.join(results_dir, 'upload-10gb')
        os.makedirs(results)
        run(benchmark_cp + ' --recursive --num-iterations %s '
            '--source %s --dest %s --result-dir %s --no-cleanup' % (
                num_iterations, local_dir, s3_location, results))

        # 10gb download
        results = os.path.join(results_dir, 'download-10gb')
        os.makedirs(results)
        run(benchmark_cp + ' --recursive --num-iterations %s '
            '--source %s --dest %s --result-dir %s' % (
                num_iterations, s3_location, local_dir, results))
    finally:
        # Not benchmarking a single rm call since it's just a single call
        run('aws s3 rm --recursive ' + s3_location)


def s3_uri(value):
    if not value.startswith('s3://'):
        return 's3://' + value
    return value


if __name__ == "__main__":
    parser = argparse.ArgumentParser()
    parser.add_argument(
        '-n', '--num-iterations', type=int, default=10,
        help='The number of times to run each test.'
    )
    parser.add_argument(
        '-b', '--bucket', default=TEST_BUCKET, type=s3_uri,
        required=TEST_BUCKET is None,
        help='The bucket to use for testing as an s3 uri. This can also be '
             'set by the environment variable PERF_TEST_BUCKET. If the '
             'environment variable is not set, then this argument is required.'
    )
    parser.add_argument(
        '--num-files', default=10000, type=int,
        help='The number of files to use for the multiple file case.'
    )
    parser.add_argument(
        '--large-file-size', default='10gb',
        help='The file size for the large file case. This can be in the form '
             '10gb, 4kb, etc.'
    )
    main(parser.parse_args())