import boto3
from ..provider import StorageProvider

class R2StorageProvider(StorageProvider):
    def __init__(self, access_key_id: str, access_key_secret: str, bucket_name: str, account_id: str, public_url: str, bucket_location: str = None):
        self.r2 = boto3.client("s3", aws_access_key_id=access_key_id, aws_secret_access_key=access_key_secret, endpoint_url=f'https://{account_id}.r2.cloudflarestorage.com')
        self.bucket_name = bucket_name
        if public_url:
            self.public_url = public_url
        else:
            print("Public URL not found in Cloudflare secrets")

    def upload_file(self, source_file, destination) -> str or None:
        destination = destination.lstrip("/")
        try:
            self.r2.upload_file(source_file, self.bucket_name, destination)
            return self.public_url + "/" + destination
        except Exception as e:
            print(e)
            return None

    def download_file(self, source, destination_file):
        source = source.lstrip("/")
        try:
            self.r2.download_file(self.bucket_name, source, destination_file)
        except Exception as e:
            print(e)
