# AUTOGENERATED! DO NOT EDIT! File to edit: ../nbs/s3_loader2.ipynb.

# %% auto 0
__all__ = ['S3FileHandler']

# %% ../nbs/s3_loader2.ipynb 1
import boto3
import os
import datetime
from datetime import tzinfo
from dateutil.tz import tzutc
from . import stem, fname

# %% ../nbs/s3_loader2.ipynb 3
class S3FileHandler:
    def __init__(self, aws_access_key, aws_secret_access_key):
        self.s3_client = boto3.client(
            "s3",
            aws_access_key_id=aws_access_key_id,
            aws_secret_access_key=aws_secret_access_key,
        )

    def list_s3_buckets(self):
        """
        Lists all the s3 buckets in s3
        """
        try:
            # Call S3 to list current buckets
            response = self.s3_client.list_buckets()
            buckets = [bucket["Name"] for bucket in response["Buckets"]]
            return buckets
        except Exception as e:
            print(e)

    def list_s3_objects(self, bucket_name, key=""):
        """
        List all files in an S3 bucket or within a specific prefix.

        :param bucket_name: str. Name of the S3 bucket.
        :param key: str or None. Specific prefix to list files from, defaults to None.
        """
        try:
            # Initialize a paginator for listing objects
            paginator = self.s3_client.get_paginator("list_objects_v2")
            # Use the paginator to fetch all objects in the specified bucket and prefix if provided
            files = dict()
            for page in paginator.paginate(Bucket=bucket_name, Prefix=key):
                # Access the 'Contents' from the page, which lists the objects
                if "Contents" in page:
                    for obj in page["Contents"]:
                        files[obj["Key"]] = obj["Size"]
                        # print(f"{obj['Key']} ({obj['Size']} bytes)")
            return files
        except Exception as e:
            print(f"An error occurred: {e}")

    def download_s3_folder(self, bucket_name, local_dir, prefix="", verbose=0):
        """
        Download all files from an S3 bucket prefix to a local directory.

        :param bucket_name: str. Name of the S3 bucket.
        :param local_dir: str. Local directory to which files will be downloaded.
        :param prefix: str or None. Prefix path of the folder in the bucket. If None, the whole bucket is downloaded.
        """
        if not prefix.endswith("/"):
            prefix = prefix + "/"
            print(prefix)
        # Ensure local directory exists
        if prefix == "":
            local_dir = os.path.join(local_dir, bucket_name)
        else:
            local_dir = os.path.join(local_dir, stem(prefix))
        if not os.path.exists(local_dir):
            os.makedirs(local_dir)

        # List objects within the specified prefix
        paginator = self.s3_client.get_paginator("list_objects_v2")
        for page in paginator.paginate(Bucket=bucket_name, Prefix=prefix):
            for obj in page.get("Contents", []):
                key = obj["Key"]
                if not key.endswith("/"):  # skip directories/folders
                    # Define file path locally in same structure
                    local_file_path = os.path.join(local_dir, key[len(prefix) :])
                    local_file_dir = os.path.dirname(local_file_path)

                    # Ensure local file directory exists
                    if not os.path.exists(local_file_dir):
                        os.makedirs(local_file_dir)

                    # Download the file
                    self.s3_client.download_file(bucket_name, key, local_file_path)
                    if verbose:
                        print(f"Downloaded {key} to {local_file_path}")

    def download_s3_file(self, bucket_name, key, local_dir, metadata=False, verbose=0):
        """
        Download a specific file from an S3 bucket and optionally return its metadata.

        :param bucket_name: str. Name of the S3 bucket.
        :param key: str. The key of the file in the S3 bucket.
        :param local_dir: str. Local directory to which the file will be downloaded.
        :param metadata: bool. If True, return the file's metadata; otherwise, return None.
        :param verbose: bool.
        :return: dict or None. Returns metadata of the file if metadata is True, otherwise None.
        """
        # Define the local file path
        local_file_path = os.path.join(local_dir, os.path.basename(key))

        # Ensure the local directory exists
        if not os.path.exists(local_dir):
            os.makedirs(local_dir)

        # Download the file
        self.s3_client.download_file(bucket_name, key, local_file_path)
        if verbose:
            print(f"Downloaded {key} to {local_file_path}")

        # Optionally retrieve and return metadata
        if metadata:
            response = self.s3_client.head_object(Bucket=bucket_name, Key=key)
            return response  # Return the metadata dictionary
        return None

    def upload_file_to_s3(self, bucket_name, localfile_path, s3_key, metadata=None):
        """
        Upload a file to an S3 bucket with optional metadata.

        :param bucket_name: str. Name of the S3 bucket.
        :param localfile_path: str. Local path to the file to be uploaded.
        :param s3_key: str. S3 key (path within the bucket) where the file will be stored with file name included.
        :param metadata: dict or None. Optional metadata for the file. Defaults to None.
        """
        try:
            # Setup the file upload options
            extra_args = {}
            if metadata:
                extra_args["Metadata"] = metadata

            # Perform the file upload
            with open(localfile_path, "rb") as file_data:
                self.s3_client.upload_fileobj(
                    Fileobj=file_data,
                    Bucket=bucket_name,
                    Key=s3_key,
                    ExtraArgs=extra_args,
                )
            print(f"File uploaded successfully to {bucket_name}/{s3_key}")
        except Exception as e:
            print(f"Failed to upload file: {e}")

    def upload_folder_to_s3(
        self, bucket_name, local_folder_path, s3_prefix="", metadata=None, verbose=0
    ):
        """
        Upload all files in a local folder to an S3 bucket with optional metadata.

        :param bucket_name: str. Name of the S3 bucket.
        :param local_folder_path: str. Local path to the folder to be uploaded.
        :param s3_prefix: str. S3 prefix (folder path within the bucket) where the files will be stored.
                           Defaults to the root of the bucket.
        :param metadata: dict or None. Optional metadata for the files. Defaults to None.
        """
        try:
            # Ensure the local_folder_path ends with a slash to properly preserve folder structure
            local_folder_path = os.path.normpath(local_folder_path)

            # Iterate over all files in the folder and its subfolders
            for root, dirs, files in os.walk(local_folder_path):
                for filename in files:
                    local_file_path = os.path.join(root, filename)

                    # Generate the corresponding S3 key (prefix + relative file path)
                    relative_path = os.path.relpath(
                        local_file_path, os.path.dirname(local_folder_path)
                    )
                    s3_key = os.path.join(s3_prefix, relative_path).replace(
                        "\\", "/"
                    )  # Replace Windows path separators

                    extra_args = {}
                    if metadata:
                        extra_args["Metadata"] = metadata

                    with open(local_file_path, "rb") as file_data:
                        self.s3_client.upload_fileobj(
                            Fileobj=file_data,
                            Bucket=bucket_name,
                            Key=s3_key,
                            ExtraArgs=extra_args,
                        )
                    if verbose:
                        print(f"Uploaded {s3_key} to {bucket_name}/{s3_key}")

        except Exception as e:
            print(f"Failed to upload folder: {e}")

    def inmemory_download_s3(bucket_name, key):
        """
        Downloads a file from an Amazon S3 bucket and loads it directly into a pandas DataFrame.
        The function automatically detects the file format based on its extension.

        Parameters:
        key (str): The S3 object key of the file to download.
        bucket (str, optional): The name of the S3 bucket. Defaults to AWS_BUCKET from .env if not provided.
        """
        response = self.s3_client.get_object(Bucket=bucket_name, Key=key)
        file_content = response["Body"].read()
        return file_content
