import sys
import os
import multiprocessing
import logging
import boto3
import datetime
import time
from botocore.exceptions import ClientError
#from botocore.client import Config
from boto3.s3.transfer import TransferConfig

import shutil
from pathlib import Path

# Set access keys as envs before running this script
# export AWS_SECRET_ACCESS_KEY=
# export AWS_ACCESS_KEY_ID=
# export AWS_S3_ENDPOINT=


paths = []
print(f"Reading input...")
try:
    counter = 0
    paths = [line.rstrip().split('\t') for line in sys.stdin]
except EOFError:
    pass
except Exception:
    pass

def upload_file(params):
    """Upload a file to an S3 bucket

    :param file_name: File to upload
    :param bucket: Bucket to upload to
    :param object_name: S3 object name. If not specified then file_name is used
    :return: True if file was uploaded, else False
    """

    bucket = 'local-media'
    data_dir = 'local-media'

    if len(params) != 2:
        print(f"!!! INVALID OR END OF INPUT: {params}")
        return ("None", "None", False)

    file_id = params[0]
    object_name = params[1]

    file_name = data_dir + '/' + file_id

    # If S3 object_name was not specified, use file_name
    #if object_name is None:
    #    object_name = os.path.basename(file_name)
    if not os.path.exists(file_name):
        return (file_id, object_name, False)

    # Upload the file

    s3_client = boto3.client('s3', endpoint_url=os.getenv('AWS_S3_ENDPOINT'))

    try:
        response = s3_client.upload_file(file_name, bucket, object_name, Config=TransferConfig(multipart_threshold=20*1024*1024))
    except ClientError as e:
        print(e)
        return (file_id, object_name, False)

    return (file_id, object_name, True)


print(f"Starting upload of {len(paths)} objects...")

with multiprocessing.Pool(40) as pool:
    result = pool.map(upload_file, paths);

report = open(f'upload-report-{int(time.mktime(datetime.datetime.now().timetuple()))}.csv', 'w')
for r in result:
    report.write(f'{r[0]}\t{r[1]}\t{r[2]}\n')

report.close()

print(f'Completed...')