#!/usr/bin/env python

from KubeJobSub.settings import AZURE_STORAGE_ACCOUNT, AZURE_SHARE_NAME, AZURE_STORAGE_KEY
from azure.storage.file import FileService
import azure.storage
import KubeJobSub.settings
from termcolor import colored
import argparse
import glob
import re
import os


def set_credentials():
    """
    Sets credentials for Azure File Storage. Credentials are saved to settings.py and loaded automatically every
    time this script gets run. Quits once done so updated credentials get loaded.
    """
    storage_account = input('Enter Azure storage account: ')
    share_name = input('Enter Azure share name: ')
    storage_key = input('Enter Azure storage key: ')
    with open(KubeJobSub.settings.__file__, 'w') as f:
        f.write("AZURE_STORAGE_ACCOUNT = '{}'\n".format(storage_account))
        f.write("AZURE_SHARE_NAME = '{}'\n".format(share_name))
        f.write("AZURE_STORAGE_KEY = '{}'\n".format(storage_key))
    print('Credentials set!')
    quit()


def check_credentials_set():
    """
    Checks all credentials necessary for interacting with azure file storage are set.
    DOES NOT check if credentials are actually valid.
    TODO: Add a check that credentials are, in fact, valid if all are present.
    :return: False if any credential is unset, True if all credentials are set.
    """
    if AZURE_SHARE_NAME is None or AZURE_STORAGE_ACCOUNT is None or AZURE_STORAGE_KEY is None:
        return False
    else:
        return True


def print_ls(file_list_generator, expression=None):
    """
    Prints output of azure's file list command in an ls-ish fashion with some nice highlighting
    provided by termcolor.
    :param file_list_generator: Generator object created by list_directories_and_files
    :param expression: Expression containing * to be matched to do selective ls-ing
    """
    for item in file_list_generator:
        filename = item.name
        match_found = True
        if expression is not None:
            expression = expression.replace('*', '.*')
            if re.match(expression, filename) is None:
                match_found = False
        if match_found:
            if filename.endswith('.gz') or filename.endswith('.bz2') or filename.endswith('.zip'):
                print(colored(filename, 'red', attrs=['bold']))
            elif type(item) is azure.storage.file.models.Directory:
                print(colored(filename, 'blue', attrs=['bold']))
            else:
                print(filename)


def find_files_matching_expression(file_list_generator, expression):
    """
    Looks for files in a directory that match a regex provided with *
    :param file_list_generator: Generator object created by list_directories_and_files
    :param expression: Expression containing a * of things you want to match
    :return: List of files that match the expression given
    """
    expression = expression.replace('*', '.*')
    files_matching_expression = list()
    for item in file_list_generator:
        filename = item.name
        if re.match(expression, filename) is not None and type(item) is azure.storage.file.models.File:
            files_matching_expression.append(filename)
    return files_matching_expression


def recursive_delete(file_service, directory_name):
    """
    Recursively deletes a directory.
    :param file_service: Instantiated file_service object.
    :param directory_name: Full path to directory that you want to delete.
    """
    generator = file_service.list_directories_and_files(share_name=AZURE_SHARE_NAME, directory_name=directory_name)
    for item in generator:
        if type(item) is azure.storage.file.models.Directory:
            recursive_delete(file_service=file_service,
                             directory_name=os.path.join(directory_name, item.name))
        else:
            file_service.delete_file(share_name=AZURE_SHARE_NAME,
                                     directory_name=directory_name,
                                     file_name=item.name)
    file_service.delete_directory(share_name=AZURE_SHARE_NAME, directory_name=directory_name)


def recursive_upload(file_service, directory_to_upload, upload_path, no_progress_bar=True):
    """
    Uploads a file on your system recursively to Azure share.
    :param file_service: Instatiated file_service object.
    :param directory_to_upload: The directory you want to upload to Azure
    :param upload_path: Path to upload to on Azure. For root of share, needs to be set to None
    :param no_progress_bar: Set to True if you do NOT want to see progress bars/info, False if seeing progress
    bar is what you want.
    """
    dir_name = os.path.split(directory_to_upload)[-1]
    if upload_path is None:
        dir_to_create = dir_name
    else:
        dir_to_create = os.path.join(upload_path, dir_name)
    if no_progress_bar is False:
        print('Creating directory {}...'.format(dir_to_create))
    file_service.create_directory(share_name=AZURE_SHARE_NAME,
                                  directory_name=dir_to_create)
    things_to_upload = glob.glob(os.path.join(directory_to_upload, '*'))
    for item in things_to_upload:
        if upload_path is None:
            directory_name = dir_name
        else:
            directory_name = os.path.join(upload_path, dir_name)
        if os.path.isfile(item):
            item_name = os.path.split(item)[-1]
            if no_progress_bar:
                file_service.create_file_from_path(share_name=AZURE_SHARE_NAME,
                                                   local_file_path=item,
                                                   directory_name=directory_name,
                                                   file_name=item_name)
            else:
                print('Uploading {}...'.format(item_name))
                file_service.create_file_from_path(share_name=AZURE_SHARE_NAME,
                                                   local_file_path=item,
                                                   directory_name=directory_name,
                                                   file_name=item_name,
                                                   progress_callback=download_callback)
        else:
            recursive_upload(file_service=file_service,
                             directory_to_upload=item,
                             upload_path=directory_name,
                             no_progress_bar=no_progress_bar)


def recursive_download(file_service, directory_to_download, download_path, no_progress_bar=True):
    """
    Downloads a folder recursively from Azure to your machine.
    :param file_service: Instantiated file_service object
    :param directory_to_download: Path of directory you want to download from Azure
    :param download_path: Path to directory you want to download the folder to on your machine
    :param no_progress_bar: Set to True if you do NOT want to see progress bars/info, False if seeing progress is
    what you want.
    """
    local_dir_name = os.path.split(directory_to_download)[-1]
    local_dir_fullpath = os.path.join(download_path, local_dir_name)
    if not os.path.isdir(local_dir_fullpath):
        os.makedirs(local_dir_fullpath)
        if no_progress_bar is False:
            print('Created directory {}...'.format(local_dir_fullpath))
    else:
        print(colored('WARNING: Local folder {} already exists'.format(local_dir_fullpath), 'red', attrs=['bold']))

    generator = file_service.list_directories_and_files(share_name=AZURE_SHARE_NAME, directory_name=directory_to_download)
    for item in generator:
        if type(item) is azure.storage.file.models.Directory:
            recursive_download(file_service=file_service,
                               directory_to_download=os.path.join(directory_to_download, item.name),
                               download_path=local_dir_fullpath,
                               no_progress_bar=no_progress_bar)
        else:
            if no_progress_bar is False:
                print('Downloading {}...'.format(item.name))
                file_service.get_file_to_path(share_name=AZURE_SHARE_NAME,
                                              directory_name=directory_to_download,
                                              file_path=os.path.join(local_dir_fullpath, item.name),
                                              file_name=item.name,
                                              progress_callback=download_callback)
            else:
                file_service.get_file_to_path(share_name=AZURE_SHARE_NAME,
                                              directory_name=directory_to_download,
                                              file_path=os.path.join(local_dir_fullpath, item.name),
                                              file_name=item.name)


def download_callback(current, total):
    """
    Shows a progress bar for file uploads/downloads.
    :param current: Current number of bytes downloaded
    :param total: Total file size
    """
    if total == 0:  # Sometimes files have zero size - set total and current to one to show instant 100 percent completion
        total = 1
        current = 1
    percent_completed = float('%1.f' % (100.0 * float(current/total)))
    size_in_mb = float('%.1f' % float(total/(1024*1024)))  # TODO: Change the suffix depending on file size.
    progress_string = '['
    for i in range(0, 100, 2):
        if percent_completed >= i:
            progress_string += '#'
        else:
            progress_string += ' '
    progress_string += ']'
    if percent_completed == 100:
        print('{} {}%\tSize:{}MB'.format(progress_string, percent_completed, size_in_mb))
    else:
        print('{} {}%\tSize:{}MB'.format(progress_string, percent_completed, size_in_mb), end='\r')


if __name__ == '__main__':
    parser = argparse.ArgumentParser(description='StorageWrapper: Using azure file shares is kind of a pain.'
                                                 'This wraps a bunch of Azure CLI file share commands into a more '
                                                 'linux-esque environment.')
    subparsers = parser.add_subparsers(help='SubCommand Help', dest='subparsers')
    set_credentials_subparser = subparsers.add_parser('set_credentials', help='Sets the azure file share and account'
                                                                              ' key as environment variables.')
    # LS
    # TODO: Add further wildcard support. Would like to have support for all wildcards, anywhere in query.
    # Also TODO: add support for -l(ish) option. Would be good to at least be able to list size
    ls_subparser = subparsers.add_parser('ls', help='Lists files in a directory. Wildcard (*) can be used,'
                                                    ' but only in final part of expression. (you can ls foo/bar*.py, '
                                                    'but not foo*/bar.py)')
    ls_subparser.add_argument('ls_path',
                              nargs='?',  # This allows the argument to be optional so things behave like actual ls.
                              default='.',
                              type=str,
                              help='Name of the directory you want to list things for.')
    # MKDIR
    mkdir_subparser = subparsers.add_parser('mkdir', help='Makes a directory.')
    mkdir_subparser.add_argument('mkdir_path',
                                 type=str,
                                 help='Name of the directory you want to create. You cannot do nested dirs.')
    # UPLOAD
    upload_subparser = subparsers.add_parser('upload', help='Uploads a file to azure file storage. Can use'
                                                            'wildcard to upload multiple files.')
    upload_subparser.add_argument('local_file',
                                  nargs='+',
                                  type=str,
                                  help='Local file you want to upload.')
    upload_subparser.add_argument('-r', '--recursive',
                                  default=False,
                                  action='store_true',
                                  help='Set this flag if you want to upload entire directories.')
    upload_subparser.add_argument('-p', '--upload_path',
                                  type=str,
                                  default=None,
                                  help='Directory to upload to on cloud. Defaults to root of share. File will have'
                                       ' same name it does locally.')
    upload_subparser.add_argument('-n', '--no_progress_bar',
                                  default=False,
                                  action='store_true',
                                  help='A progress bar on file uploads. Useful for larger files, enabled by '
                                       'default. Setting flag will disable progress bar.')
    # DOWNLOAD
    download_subparser = subparsers.add_parser('download', help='Downloads files from cloud to your machine.')
    download_subparser.add_argument('cloud_download_file',
                                    type=str,
                                    help='Path to files on cloud you want to download.')
    download_subparser.add_argument('download_dir',
                                    type=str,
                                    nargs='?',
                                    default='.',
                                    help='Directory to download files to. Defaults to your current working directory')
    download_subparser.add_argument('-r', '--recursive',
                                    default=False,
                                    action='store_true',
                                    help='Set this flag if you want to upload entire directories.')
    download_subparser.add_argument('-n', '--no_progress_bar',
                                    default=False,
                                    action='store_true',
                                    help='A progress bar on file uploads. Useful for larger files, enabled by '
                                         'default. Setting flag will disable progress bar.')
    # RM
    rm_subparser = subparsers.add_parser('rm', help='Deletes a file. Can be run recursively to delete entire'
                                                    ' directories with the -r flag.')
    rm_subparser.add_argument('rm_path',
                              type=str,
                              help='Full path on cloud to file you want deleted. Wildcards (*) works on final'
                                   ' part of path (You can rm dir/*.txt, but not rm */asdf.txt')
    rm_subparser.add_argument('-r', '--recursive_delete',
                              default=False,
                              action='store_true',
                              help='Set to recursively delete. Use with caution!')
    args = parser.parse_args()

    credentials_set = check_credentials_set()

    # Used if people want to change the credentials (different storage accounts, or whatever.
    if 'set_credentials' == args.subparsers:
        set_credentials()
    # Any of the rest of the commands need credentials set, so force user to do them if they haven't been set previously
    if credentials_set is False:
        print('Credentials not set! Please set them before proceeding.')
        set_credentials()

    # Now go through our commands!
    # Create our file service object first, it gets used by everything.
    file_service = FileService(account_name=AZURE_STORAGE_ACCOUNT, account_key=AZURE_STORAGE_KEY)
    # LS
    if args.subparsers == 'ls':
        if '*' in args.ls_path:
            ls_path = os.path.split(args.ls_path)[0]
            wildcard_expression = os.path.split(args.ls_path)[-1]
        else:
            ls_path = args.ls_path
            wildcard_expression = None
        try:
            generator = file_service.list_directories_and_files(share_name=AZURE_SHARE_NAME, directory_name=ls_path)
            print_ls(file_list_generator=generator, expression=wildcard_expression)
        except azure.common.AzureMissingResourceHttpError:
            print('ERROR: {} does not exist on your Azure File share.'.format(ls_path))
    # MKDIR
    elif args.subparsers == 'mkdir':
        file_service.create_directory(share_name=AZURE_SHARE_NAME, directory_name=args.mkdir_path)
    # UPLOAD
    elif args.subparsers == 'upload':
        # Azure does not want any trailing slash in the upload path, or it gives an error. Deal with that here
        upload_path = args.upload_path
        if upload_path is not None:
            if upload_path.endswith('/'):
                upload_path = upload_path[:-1]
        if args.recursive:
            local_dir = os.path.abspath(args.local_file[0])
            recursive_upload(file_service=file_service,
                             directory_to_upload=local_dir,
                             upload_path=upload_path,
                             no_progress_bar=args.no_progress_bar)
            quit()
        for local_file in args.local_file:
            file_name = os.path.split(local_file)[-1]
            # Azure storage claims you get automatic progress notifications, apparently you need the progress_callback
            # set.
            if args.no_progress_bar:
                file_service.create_file_from_path(share_name=AZURE_SHARE_NAME,
                                                   local_file_path=local_file,
                                                   directory_name=upload_path,
                                                   file_name=file_name)
            else:
                print('Uploading {}...'.format(file_name))
                file_service.create_file_from_path(share_name=AZURE_SHARE_NAME,
                                                   local_file_path=local_file,
                                                   directory_name=upload_path,
                                                   file_name=file_name,
                                                   progress_callback=download_callback)

    # DOWNLOAD
    elif args.subparsers == 'download':
        if args.recursive:
            # Check for trailing slash and get rid of it.
            cloud_download_dir = args.cloud_download_file
            if cloud_download_dir.endswith('/'):
                cloud_download_dir = cloud_download_dir[:-1]
            recursive_download(file_service=file_service,
                               directory_to_download=cloud_download_dir,
                               download_path=args.download_dir,
                               no_progress_bar=args.no_progress_bar)
            quit()

        file_name = os.path.split(args.cloud_download_file)[-1]
        file_directory = os.path.split(args.cloud_download_file)[0]
        if file_directory == '':
            file_directory = None
        generator = file_service.list_directories_and_files(share_name=AZURE_SHARE_NAME, directory_name=file_directory)
        files_to_download = find_files_matching_expression(file_list_generator=generator,
                                                           expression=file_name)
        for file_to_download in files_to_download:
            if args.no_progress_bar:
                file_service.get_file_to_path(share_name=AZURE_SHARE_NAME,
                                              file_path=os.path.join(args.download_dir, file_to_download),
                                              file_name=file_to_download,
                                              directory_name=file_directory)
            else:
                print('Downloading {}...'.format(file_to_download))
                file_service.get_file_to_path(share_name=AZURE_SHARE_NAME,
                                              file_path=os.path.join(args.download_dir, file_to_download),
                                              file_name=file_to_download,
                                              directory_name=file_directory,
                                              progress_callback=download_callback)
    # DELETE/RM
    elif args.subparsers == 'rm':
        file_name = os.path.split(args.rm_path)[-1]
        file_directory = os.path.split(args.rm_path)[0]
        if file_directory == '':
            file_directory = None
        if args.recursive_delete:
            # Azure also apparently does not like trailing slashes for directory removal, so take care of that too.
            rm_path = args.rm_path
            if rm_path.endswith('/'):
                rm_path = rm_path[:-1]
            recursive_delete(file_service=file_service,
                             directory_name=rm_path)
        else:
            generator = file_service.list_directories_and_files(share_name=AZURE_SHARE_NAME, directory_name=file_directory)
            files_to_delete = find_files_matching_expression(file_list_generator=generator,
                                                             expression=file_name)
            for file_to_delete in files_to_delete:
                file_service.delete_file(share_name=AZURE_SHARE_NAME,
                                         directory_name=file_directory,
                                         file_name=file_to_delete)
