import os
from concurrent.futures import ThreadPoolExecutor

import boto3
from boto3.s3.transfer import TransferConfig
from rich import print

from airflow_tools import config


class S3Connector:
    """Connect to S3."""

    def __init__(self):
        self._client = boto3.client("s3", region_name=config.AWS_REGION)

    def upload_file(
        self,
        bucket: str,
        filename: str,
        key: str,
        kms_key_id: str,
        transfer_config: TransferConfig = None,
    ):
        print(
            f"Uploading file {filename} to s3://{bucket}/{key} "
            f"with KMS key {kms_key_id}"
        )
        self._client.upload_file(
            Filename=filename,
            Bucket=bucket,
            Key=key,
            ExtraArgs={"ServerSideEncryption": "aws:kms", "SSEKMSKeyId": kms_key_id},
        )

    def upload_dir(self, bucket: str, folder: str, key_prefix: str, kms_key_id: str):
        print(
            f"Uploading directory {folder} to s3://{bucket}/{key_prefix} "
            f"with KMS key {kms_key_id}"
        )

        transfer_config = TransferConfig(use_threads=True)

        # Use ThreadPoolExecutor to upload files in parallel
        with ThreadPoolExecutor() as executor:
            futures = []
            for root, _, files in os.walk(folder):
                for file in files:
                    file_path = os.path.join(root, file)
                    relative_path = os.path.relpath(file_path, folder)
                    s3_key = f"{key_prefix}/{relative_path}"
                    futures.append(
                        executor.submit(
                            self.upload_file,
                            bucket=bucket,
                            filename=file_path,
                            key=s3_key,
                            kms_key_id=kms_key_id,
                            transfer_config=transfer_config,
                        )
                    )

            # Wait for all uploads to complete
            for future in futures:
                future.result()