import os
import boto3
import botocore
from botocore.client import Config
import re
import argparse
import glob
from ast import literal_eval
import sys


def get_args():
    parser = argparse.ArgumentParser()
    parser.add_argument("--bucket-name", dest="bucket_name", required=True)
    parser.add_argument(
        "--source-file-name-match-type",
        dest="source_file_name_match_type",
        default="exact_match",
        choices={"exact_match", "regex_match"},
        required=False,
    )
    parser.add_argument("--source-file-name", dest="source_file_name", required=True)
    parser.add_argument(
        "--source-folder-name", dest="source_folder_name", default="", required=False
    )
    parser.add_argument(
        "--destination-folder-name",
        dest="destination_folder_name",
        default="",
        required=False,
    )
    parser.add_argument(
        "--destination-file-name",
        dest="destination_file_name",
        default=None,
        required=False,
    )
    parser.add_argument("--s3-config", dest="s3_config", default=None, required=False)
    parser.add_argument("--aws-access-key-id", dest="aws_access_key_id", required=False)
    parser.add_argument(
        "--aws-secret-access-key", dest="aws_secret_access_key", required=False
    )
    parser.add_argument(
        "--aws-default-region", dest="aws_default_region", required=False
    )
    parser.add_argument("--extra-args", dest="extra_args", required=False)
    return parser.parse_args()


def set_environment_variables(args):
    """
    Set AWS credentials as environment variables if they're provided via keyword arguments
    rather than seeded as environment variables. This will override system defaults.
    """

    if args.aws_access_key_id:
        os.environ["AWS_ACCESS_KEY_ID"] = args.aws_access_key_id
    if args.aws_secret_access_key:
        os.environ["AWS_SECRET_ACCESS_KEY"] = args.aws_secret_access_key
    if args.aws_default_region:
        os.environ["AWS_DEFAULT_REGION"] = args.aws_default_region
    return


def connect_to_s3(s3_config=None):
    """
    Create a connection to the S3 service using credentials provided as environment variables.
    """
    s3_connection = boto3.client("s3", config=Config(s3_config))
    return s3_connection


def extract_file_name_from_source_full_path(source_full_path):
    """
    Use the file name provided in the source_full_path variable. Should be run only
    if a destination_file_name is not provided.
    """
    destination_file_name = os.path.basename(source_full_path)
    return destination_file_name


def enumerate_destination_file_name(destination_file_name, file_number=1):
    """
    Append a number to the end of the provided destination file name.
    Only used when multiple files are matched to, preventing the destination file from being continuously overwritten.
    """
    if re.search(r"\.", destination_file_name):
        destination_file_name = re.sub(
            r"\.", f"_{file_number}.", destination_file_name, 1
        )
    else:
        destination_file_name = f"{destination_file_name}_{file_number}"
    return destination_file_name


def determine_destination_file_name(
    *, source_full_path, destination_file_name, file_number=None
):
    """
    Determine if the destination_file_name was provided, or should be extracted from the source_file_name,
    or should be enumerated for multiple file downloads.
    """
    if destination_file_name:
        if file_number:
            destination_file_name = enumerate_destination_file_name(
                destination_file_name, file_number
            )
        else:
            destination_file_name = destination_file_name
    else:
        destination_file_name = extract_file_name_from_source_full_path(
            source_full_path
        )

    return destination_file_name


def clean_folder_name(folder_name):
    """
    Cleans folders name by removing duplicate '/' as well as leading and trailing '/' characters.
    """
    folder_name = folder_name.strip("/")
    if folder_name != "":
        folder_name = os.path.normpath(folder_name)
    return folder_name


def combine_folder_and_file_name(folder_name, file_name):
    """
    Combine together the provided folder_name and file_name into one path variable.
    """
    combined_name = os.path.normpath(
        f'{folder_name}{"/" if folder_name else ""}{file_name}'
    )
    combined_name = os.path.normpath(combined_name)

    return combined_name


def determine_destination_full_path(
    destination_folder_name, destination_file_name, source_full_path, file_number=None
):
    """
    Determine the final destination name of the file being downloaded.
    """
    destination_file_name = determine_destination_file_name(
        destination_file_name=destination_file_name,
        source_full_path=source_full_path,
        file_number=file_number,
    )
    destination_full_path = combine_folder_and_file_name(
        destination_folder_name, destination_file_name
    )
    return destination_full_path


def find_all_local_file_names(source_folder_name):
    """
    Returns a list of all files that exist in the current working directory,
    filtered by source_folder_name if provided.
    """
    cwd = os.getcwd()
    cwd_extension = os.path.normpath(f"{cwd}/{source_folder_name}/**")
    file_names = glob.glob(cwd_extension, recursive=True)
    return file_names


def find_all_file_matches(file_names, file_name_re):
    """
    Return a list of all file_names that matched the regular expression.
    """
    matching_file_names = []
    for file in file_names:
        if re.search(file_name_re, file):
            matching_file_names.append(file)

    return matching_file_names


def upload_s3_file(
    s3_connection, bucket_name, source_full_path, destination_full_path, extra_args=None
):
    """
    Uploads a single file to S3. Uses the s3.transfer method to ensure that files larger than 5GB are split up during the upload process.

    Extra Args can be found at https://boto3.amazonaws.com/v1/documentation/api/latest/guide/s3-uploading-files.html#the-extraargs-parameter
    and are commonly used for custom file encryption or permissions.
    """
    s3_upload_config = boto3.s3.transfer.TransferConfig()
    s3_transfer = boto3.s3.transfer.S3Transfer(
        client=s3_connection, config=s3_upload_config
    )

    s3_transfer.upload_file(
        source_full_path, bucket_name, destination_full_path, extra_args=extra_args
    )

    print(
        f"{source_full_path} successfully uploaded to {bucket_name}/{destination_full_path}"
    )


def main():
    args = get_args()
    set_environment_variables(args)
    bucket_name = args.bucket_name
    source_file_name = args.source_file_name
    source_folder_name = args.source_folder_name
    source_full_path = combine_folder_and_file_name(
        folder_name=f"{os.getcwd()}/{source_folder_name}", file_name=source_file_name
    )
    destination_folder_name = clean_folder_name(args.destination_folder_name)
    source_file_name_match_type = args.source_file_name_match_type
    s3_config = args.s3_config
    extra_args = literal_eval(args.extra_args if args.extra_args else "{}")

    s3_connection = connect_to_s3(s3_config)

    if source_file_name_match_type == "regex_match":
        file_names = find_all_local_file_names(source_folder_name)
        matching_file_names = find_all_file_matches(
            file_names, re.compile(source_file_name)
        )
        num_matches = len(matching_file_names)

        if num_matches == 0:
            print(f"No matches found for regex {source_file_name}")
            sys.exit(1)
        else:
            print(f"{num_matches} files found. Preparing to upload...")

        for index, key_name in enumerate(matching_file_names):
            destination_full_path = determine_destination_full_path(
                destination_folder_name=destination_folder_name,
                destination_file_name=args.destination_file_name,
                source_full_path=key_name,
                file_number=None if num_matches == 1 else index + 1,
            )
            print(f"Uploading file {index+1} of {len(matching_file_names)}")
            upload_s3_file(
                source_full_path=key_name,
                destination_full_path=destination_full_path,
                bucket_name=bucket_name,
                extra_args=extra_args,
                s3_connection=s3_connection,
            )

    else:
        destination_full_path = determine_destination_full_path(
            destination_folder_name=destination_folder_name,
            destination_file_name=args.destination_file_name,
            source_full_path=source_full_path,
        )
        upload_s3_file(
            source_full_path=source_full_path,
            destination_full_path=destination_full_path,
            bucket_name=bucket_name,
            extra_args=extra_args,
            s3_connection=s3_connection,
        )


if __name__ == "__main__":
    main()
