#!python
"""
S3 download plugin for CRDS.

Exit status:
0 - success
2 - Invalid arg or missing environment vars
10 - failed file size verification
11 - failed checksum verification
400 - An error occurred (400) when calling the HeadObject operation: Bad Request
404 - File does not exist on S3
An error occurred (ExpiredToken) when calling the GetObject operation: The provided token has expired.
other codes - aws-cli failure.  See https://docs.aws.amazon.com/cli/latest/topic/return-codes.html

As of 2021-04-23, copies of this script are maintained in the crds
and caldp repositories.  Please ensure that any bug fixes make it into
both!
"""
import argparse
import os
import re
import subprocess
import sys
from crds.core.utils import checksum
from crds.core import config
from crds.core import log

INVALID_ARGS_STATUS = 2 # aws: error: the following arguments are required:
BAD_REQUEST_STATUS = 400 # An error occurred (400) when calling the HeadObject operation: Bad Request
BAD_SIZE_STATUS = 10
BAD_CHECKSUM_STATUS = 11

try:
    import boto3
    import awscli
except ImportError:
    boto3 = None
    awscli = None

def check_aws_imports():
    return boto3 is not None


def format_uris(**kwargs):
    source, destination = kwargs.pop("source"), kwargs.pop("destination")
    fname = source.split("/")[-1]
    s3_uri = source if source.startswith("s3") else config.get_uri(fname)
    obs = s3_uri.split('/')[-2]
    crds_path = os.environ.get("CRDS_PATH", config.get_crds_path())
    if destination == crds_path:
        dest = config.locate_file(fname, obs)
        os.makedirs(os.path.dirname(dest), exist_ok=True)
    # destination is a directory that may or may not already exist
    elif not destination.endswith(fname):
        os.makedirs(destination, exist_ok=True)
        dest = f"{destination.rstrip('/')}/{fname}"
    else:
        # destination is already a filename path
        dest = destination
        os.makedirs(os.path.dirname(dest), exist_ok=True)
    return s3_uri, dest


def parse_args():
    parser = argparse.ArgumentParser("crds_s3_get", description="S3 download plugin for CRDS")
    parser.add_argument("source", help="filename to download or full S3 URI to the file")
    parser.add_argument("-d", "--destination", help="Destination path on local filesystem", default=os.environ.get("CRDS_PATH", None))
    parser.add_argument("-s", "--file-size", help="Expected file size in bytes", type=int, default=None)
    parser.add_argument("-c", "--file-sha1sum", help="Expected file SHA-1 checksum", default=None)
    parser.add_argument("-r", "--max-retries", help="Maximum number of retries on download failure", type=int, default=3)
    parser.add_argument("-I", "--ignore-cache", help="Ignore local cache and force download", action="store_true")
    parser.add_argument("--loglevel", default="INFO", help="set loglevel", type=str)
    return parser.parse_args()


def main():
    if not check_aws_imports():
        raise ImportError(
            "You must install awscli and boto3 for the crds_s3_get script to work. "
            "AWS dependencies for CRDS can be installed via `pip install crds[aws]`"
        )
    args = parse_args()
    log.THE_LOGGER.handlers[0].setLevel(args.loglevel)
    if args.destination is None:
        print("Destination path defaults to CRDS_PATH but no value was set. \n" \
        "Please set the CRDS_PATH variable e.g. `export CRDS_PATH=path/to/local/cache` \n" \
        "or pass an absolute path on local disk where you want the file to be downloaded: \n" \
        "`crds_s3_get myfile -d abs/path/to/download")
        sys.exit(1)
    kwargs = {**vars(args)}
    src, dest = format_uris(**kwargs)
    fname = src.split("/")[-1]
    if os.path.exists(dest):
        if not args.ignore_cache:
            log.debug(f"crds_s3_get - '{fname}' exists in cache: '{dest}' - Skipping", verbosity=60)
            sys.exit(0)
        else:
            log.debug(f"crds_s3_get - '{fname}' exists at '{dest}' but --ignore_cache flag was set - Re-syncing.", verbosity=60)

    result = subprocess.run([ "aws", "s3", "cp", src, dest], encoding="utf-8")

    if result.returncode != 0:
        log.error(f"crds_s3_get - Failed to download '{fname}' with return code {result.returncode}")
        sys.exit(result.returncode)
    else:
        log.info(f"crds_s3_get: Successfully downloaded '{fname}' to '{dest}'", file=sys.stdout)

    if args.file_size not in [None, ""]:
        downloaded_size = os.path.getsize(dest)
        if downloaded_size != args.file_size:
            log.error(f"crds_s3_get: '{src}' failed file size check.  Expected: {args.file_size} Received: {downloaded_size}")
            os.unlink(dest)
            sys.exit(BAD_SIZE_STATUS)

    if args.file_sha1sum not in [None, ""]:
        downloaded_sha1sum = checksum(dest)
        if downloaded_sha1sum != args.file_sha1sum:
            log.error(f"crds_s3_get: '{src}' failed checksum.  Expected: {args.file_sha1sum} Received: {downloaded_sha1sum}")
            os.unlink(dest)
            sys.exit(BAD_CHECKSUM_STATUS)


if __name__ == "__main__":
    main()
