import boto3
from botocore.exceptions import ClientError, NoCredentialsError
from loguru import logger

from .service import StorageService


class S3StorageService(StorageService):
    """A service class for handling operations with AWS S3 storage."""

    def __init__(self, session_service, settings_service) -> None:
        """Initialize the S3 storage service with session and settings services."""
        super().__init__(session_service, settings_service)
        self.bucket = "langflow"
        self.s3_client = boto3.client("s3")
        self.set_ready()

    async def save_file(self, folder: str, file_name: str, data) -> None:
        """Save a file to the S3 bucket.

        :param folder: The folder in the bucket to save the file.
        :param file_name: The name of the file to be saved.
        :param data: The byte content of the file.
        :raises Exception: If an error occurs during file saving.
        """
        try:
            self.s3_client.put_object(Bucket=self.bucket, Key=f"{folder}/{file_name}", Body=data)
            logger.info(f"File {file_name} saved successfully in folder {folder}.")
        except NoCredentialsError:
            logger.exception("Credentials not available for AWS S3.")
            raise
        except ClientError:
            logger.exception(f"Error saving file {file_name} in folder {folder}")
            raise

    async def get_file(self, folder: str, file_name: str):
        """Retrieve a file from the S3 bucket.

        :param folder: The folder in the bucket where the file is stored.
        :param file_name: The name of the file to be retrieved.
        :return: The byte content of the file.
        :raises Exception: If an error occurs during file retrieval.
        """
        try:
            response = self.s3_client.get_object(Bucket=self.bucket, Key=f"{folder}/{file_name}")
            logger.info(f"File {file_name} retrieved successfully from folder {folder}.")
            return response["Body"].read()
        except ClientError:
            logger.exception(f"Error retrieving file {file_name} from folder {folder}")
            raise

    async def list_files(self, folder: str):
        """List all files in a specified folder of the S3 bucket.

        :param folder: The folder in the bucket to list files from.
        :return: A list of file names.
        :raises Exception: If an error occurs during file listing.
        """
        try:
            response = self.s3_client.list_objects_v2(Bucket=self.bucket, Prefix=folder)
        except ClientError:
            logger.exception(f"Error listing files in folder {folder}")
            raise

        files = [item["Key"] for item in response.get("Contents", []) if "/" not in item["Key"][len(folder) :]]
        logger.info(f"{len(files)} files listed in folder {folder}.")
        return files

    async def delete_file(self, folder: str, file_name: str) -> None:
        """Delete a file from the S3 bucket.

        :param folder: The folder in the bucket where the file is stored.
        :param file_name: The name of the file to be deleted.
        :raises Exception: If an error occurs during file deletion.
        """
        try:
            self.s3_client.delete_object(Bucket=self.bucket, Key=f"{folder}/{file_name}")
            logger.info(f"File {file_name} deleted successfully from folder {folder}.")
        except ClientError:
            logger.exception(f"Error deleting file {file_name} from folder {folder}")
            raise

    async def teardown(self) -> None:
        """Perform any cleanup operations when the service is being torn down."""
        # No specific teardown actions required for S3 storage at the moment.