import boto3
from smart_open import open
from typing import Any, List, Generator, Optional

S3 = "s3"


class S3Client:
    def __init__(self, bucket_name: str):
        self.client = boto3.client(S3)
        self.bucket_name = bucket_name

    def get_object(self, file_name: str, format: Optional[str] = None) -> Any:
        doc = self.client.get_object(Bucket=self.bucket_name, Key=file_name).get("Body").read()
        if format:
            return doc.decode(format, errors="ignore")
        return doc

    def get_objects_list(self, folder: str = None) -> List[dict]:
        documents = self.client.list_objects(Bucket=self.bucket_name, Prefix=folder).get("Contents")
        return list(filter(lambda x: x.get("Size") > 0, documents))

    def get_object_stream(self, file_name: str) -> Generator[List, Any, None]:
        return open(f"s3://{self.bucket_name}/{file_name}", transport_params=dict(session=boto3.Session()))