#!/usr/bin/env python2.7
# pylint: disable=C0301
from __future__ import (absolute_import, division, print_function,
                        unicode_literals)

from contextlib import contextmanager
from datetime import datetime
from hashlib import sha256
from json import dump as save
from json import load
from os import remove as _delete_file
from os import chdir, environ, stat
from os.path import abspath, basename, dirname, expandvars, realpath
from subprocess import check_call as run
from sys import argv

from boto.exception import S3ResponseError
from boto.s3.connection import S3Connection
from boto.s3.key import Key

CONFIG_FILE = './S3Cachefile.json'
UPLOAD_TODO_FILE = './S3CacheTodo.json'
BYTES_PER_MB = 1024 * 1024


@contextmanager
def timer():
    start = datetime.utcnow()
    yield
    end = datetime.utcnow()
    elapsed = end - start
    print("\tDone. Took", int(elapsed.total_seconds()), "second(s).")


@contextmanager
def todo_file(writeback=True):
    try:
        with open(UPLOAD_TODO_FILE, 'rt') as json_file:
            todo = load(json_file)
    except (IOError, OSError, ValueError):
        todo = {}

    yield todo

    if writeback:
        try:
            with open(UPLOAD_TODO_FILE, 'wt') as json_file:
                save(todo, json_file)
        except (OSError, IOError) as save_err:
            print("Error saving {}:".format(UPLOAD_TODO_FILE), save_err)


def _sha256_of_file(filename):
    hasher = sha256()
    with open(filename, 'rb') as input_file:
        hasher.update(input_file.read())
    file_hash = hasher.hexdigest()
    print('sha256({}) = {}'.format(filename, file_hash))
    return file_hash


def _delete_file_quietly(filename):
    try:
        _delete_file(filename)
    except (OSError, IOError):
        pass


def mark_needs_uploading(cache_name):
    with todo_file() as todo:
        todo[cache_name] = True


def mark_uploaded(cache_name):
    with todo_file() as todo:
        todo.pop(cache_name, None)


def need_to_upload(cache_name):
    with todo_file(writeback=False) as todo:
        return todo.get(cache_name, False)


def _tarball_size(directory):
    kib = stat(_tarball_filename_for(directory)).st_size // BYTES_PER_MB
    return "{} MiB".format(kib)


def _tarball_filename_for(directory):
    return abspath('./{}.tar.gz'.format(basename(directory)))


def _create_tarball(directory):
    print("Creating tarball of {}...".format(directory))
    with timer():
        run(['tar', '-czf', _tarball_filename_for(directory), '-C', dirname(directory), basename(directory)])


def _extract_tarball(directory):
    print("Extracting tarball of {}...".format(directory))
    with timer():
        run(['tar', '-xzf', _tarball_filename_for(directory), '-C', dirname(directory)])


def download(directory):
    mark_uploaded(cache_name)  # reset
    try:
        print("Downloading {} tarball from S3...".format(cache_name))
        with timer():
            key.get_contents_to_filename(_tarball_filename_for(directory))
    except S3ResponseError as err:
        mark_needs_uploading(cache_name)
        raise SystemExit("Cached {} download failed!".format(cache_name))
    print("Downloaded {}.".format(_tarball_size(directory)))
    _extract_tarball(directory)
    print("{} successfully installed from cache.".format(cache_name))


def upload(directory):
    _create_tarball(directory)
    print("Uploading {} tarball to S3... ({})".format(cache_name, _tarball_size(directory)))
    with timer():
        key.set_contents_from_filename(_tarball_filename_for(directory))
    print("{} cache successfully updated.".format(cache_name))
    mark_uploaded(cache_name)


if __name__ == '__main__':
    # Uses environment variables:
    #   AWS_ACCESS_KEY_ID -- AWS Access Key ID
    #   AWS_SECRET_ACCESS_KEY -- AWS Secret Access Key
    argv.pop(0)
    if len(argv) != 2:
        raise SystemExit("USAGE: s3_cache.py <download | upload> <cache name>")
    mode, cache_name = argv
    script_dir = dirname(realpath(__file__))
    chdir(script_dir)
    try:
        with open(CONFIG_FILE, 'rt') as config_file:
            config = load(config_file)
    except (IOError, OSError, ValueError) as config_err:
        print(config_err)
        raise SystemExit("Error when trying to load config from JSON file!")

    try:
        cache_info = config[cache_name]
        key_file = expandvars(cache_info["key"])
        fallback_cmd = cache_info["generate"]
        directory = expandvars(cache_info["cache"])
    except (TypeError, KeyError) as load_err:
        print(load_err)
        raise SystemExit("Config for cache named {!r} is missing or malformed!".format(cache_name))

    try:
        try:
            BUCKET_NAME = environ['TWBS_S3_BUCKET']
        except KeyError:
            raise SystemExit("TWBS_S3_BUCKET environment variable not set!")

        conn = S3Connection()
        bucket = conn.lookup(BUCKET_NAME)
        if bucket is None:
            raise SystemExit("Could not access bucket!")

        key_file_hash = _sha256_of_file(key_file)

        key = Key(bucket, key_file_hash)
        key.storage_class = 'REDUCED_REDUNDANCY'

        if mode == 'download':
            download(directory)
        elif mode == 'upload':
            if need_to_upload(cache_name):
                upload(directory)
            else:
                print("No need to upload anything.")
        else:
            raise SystemExit("Unrecognized mode {!r}".format(mode))
    except BaseException as exc:
        if mode != 'download':
            raise
        print("Error!:", exc)
        print("Unable to download from cache.")
        print("Running fallback command to generate cache directory {!r}: {}".format(directory, fallback_cmd))
        with timer():
            run(fallback_cmd, shell=True)
