from contextlib import asynccontextmanager
from typing import Optional, BinaryIO
import aioboto3
from botocore.exceptions import ClientError
import os

class S3Adapter:
    def __init__(self, 
                 endpoint_url: str,
                 aws_access_key_id: str,
                 aws_secret_access_key: str,
                 bucket_name: str):
        self.endpoint_url = endpoint_url
        self.aws_access_key_id = aws_access_key_id
        self.aws_secret_access_key = aws_secret_access_key
        self.bucket_name = bucket_name
        self.session = aioboto3.Session()

    @asynccontextmanager
    async def _get_client(self):
        async with self.session.client(
            "s3",
            endpoint_url=self.endpoint_url,
            aws_access_key_id=self.aws_access_key_id,
            aws_secret_access_key=self.aws_secret_access_key,
        ) as client:
            yield client

    async def upload_file(self, object_name: str, data: bytes, content_type: Optional[str] = None):
        """Uploads bytes data to S3."""
        try:
            extra_args = {}
            if content_type:
                extra_args["ContentType"] = content_type
            
            async with self._get_client() as client:
                await client.put_object(
                    Bucket=self.bucket_name,
                    Key=object_name,
                    Body=data,
                    **extra_args
                )
            return True
        except ClientError as e:
            # logging.error(e)
            print(f"Error uploading to S3: {e}")
            return False

    async def get_file(self, object_name: str) -> Optional[bytes]:
        """Downloads a file from S3 and returns bytes."""
        try:
            async with self._get_client() as client:
                response = await client.get_object(Bucket=self.bucket_name, Key=object_name)
                return await response['Body'].read()
        except ClientError as e:
            print(f"Error downloading from S3: {e}")
            return None

    async def stream_file(self, object_name: str, chunk_size: int = 65536):
        """Streams a file from S3 yielding chunks. Memory-efficient for large files."""
        try:
            async with self._get_client() as client:
                response = await client.get_object(Bucket=self.bucket_name, Key=object_name)
                async with response['Body'] as stream:
                    while True:
                        chunk = await stream.read(chunk_size)
                        if not chunk:
                            break
                        yield chunk
        except ClientError as e:
            print(f"Error streaming from S3: {e}")
            return

    async def delete_file(self, object_name: str):
        """Deletes a file from S3."""
        try:
            async with self._get_client() as client:
                await client.delete_object(Bucket=self.bucket_name, Key=object_name)
            return True
        except ClientError as e:
            print(f"Error deleting from S3: {e}")
            return False

    async def get_presigned_url(self, object_name: str, expiration: int = 3600) -> Optional[str]:
        """Generate a presigned URL to share an S3 object."""
        try:
            async with self._get_client() as client:
                response = await client.generate_presigned_url(
                    'get_object',
                    Params={'Bucket': self.bucket_name, 'Key': object_name},
                    ExpiresIn=expiration
                )
            return response
        except ClientError as e:
            print(f"Error generating presigned URL: {e}")
            return None