import io
import sys

import boto3
import yaml
import os


class CDNCredentials:
    def __init__(self, host, downloader_access_key, downloader_access_secret, uploader_access_key, uploader_access_secret):
        self.host = host
        self.downloader_access_key = downloader_access_key
        self.downloader_access_secret = downloader_access_secret
        self.uploader_access_key = uploader_access_key
        self.uploader_access_secret = uploader_access_secret


class CDNManager:
    def __init__(self, credentials: CDNCredentials):
        self.creds = credentials
        self.download_client = boto3.client('s3', endpoint_url=self.creds.host,
                                         aws_access_key_id=self.creds.downloader_access_key,
                                         aws_secret_access_key=self.creds.downloader_access_secret)
        self.upload_client = boto3.client('s3', endpoint_url=self.creds.host,
                                                  aws_access_key_id=self.creds.uploader_access_key,
                                                  aws_secret_access_key=self.creds.uploader_access_secret)

    def upload(self, bucket: str, filename: str, file_bytes: bytearray):
        try:
            self.upload_client.upload_fileobj(io.BytesIO(file_bytes), bucket, filename)
            print(f'uploaded {filename} ({len(file_bytes)} bytes) to the {bucket}')
            return True
        except Exception as e:
            print(e)
            return False

    def download(self, bucket: str, filename: str):
        try:
            if filename is not None:
                self.download_client.download_file(bucket, filename, os.path.join(bucket, filename))
                print(f'downloaded {filename} from the {bucket} to current folder')
            else:
                response = self.download_client.list_objects_v2(Bucket=bucket)
                if 'Contents' in response:
                    for obj in response['Contents']:
                        object_key = obj['Key']
                        local_filepath = os.path.join(bucket, object_key)
                        local_dir = os.path.dirname(local_filepath)
                        if local_dir:
                            os.makedirs(local_dir, exist_ok=True)

                        if not object_key.endswith('/'):
                            try:
                                self.download_client.download_file(bucket, object_key, local_filepath)
                            except Exception as e_file:
                                all_successful = False  # Mark as failed if any file fails
            return True
        except Exception as e:
            print(e)
            return False


config_file = os.path.join('downloader_config.yaml')
with open(config_file, 'r') as f:
    config = yaml.safe_load(f)

cdn_manager = CDNManager(CDNCredentials(
    config['host'],
    config['downloader_access_key'],
    config['downloader_access_secret'],
    config['uploader_access_key'],
    config['uploader_access_secret']
))

input_action = sys.argv[1]
input_bucket = sys.argv[2]
input_filename = sys.argv[3] if len(sys.argv) > 3 else None
if len(sys.argv) > 4: # 0 is this script's path, hence 5 args is max
    input_path = sys.argv[4]

if input_action == 'upload':
    with open(input_path, 'rb') as f:
        cdn_manager.upload(input_bucket, input_filename, bytearray(f.read()))
elif input_action == 'download':
    cdn_manager.download(input_bucket, input_filename)