"""S3 connector for handling pre-signed URLs and multipart uploads."""

from datetime import datetime, timedelta, timezone
from typing import Any, Optional, TypedDict

import boto3
from botocore.client import Config
from botocore.exceptions import ClientError

from abacus_file_upload.constants import PRESIGNED_URL_EXPIRATION_SECS


class ObjectMetadata(TypedDict, total=False):
    """Metadata for an S3 object."""

    etag: Optional[str]
    last_modified: Optional[datetime]
    metadata: dict[str, str]
    size: int


class PresignedUrlPart(TypedDict):
    """Presigned URL for a single part."""

    expires_at: datetime
    part_number: int
    url: str


class S3Connector:
    """Connector for S3 operations including pre-signed URLs and multipart uploads."""

    def __init__(self, s3_client):
        """Initialize S3 connector.

        Args:
            s3_client: Boto3 S3 client instance
        """
        self.s3_client = s3_client

    def abort_multipart_upload(self, bucket: str, key: str, upload_id: str) -> None:
        """Abort a multipart upload.

        Args:
            bucket: S3 bucket name
            key: S3 object key
            upload_id: Multipart upload ID
        """
        self.s3_client.abort_multipart_upload(
            Bucket=bucket, Key=key, UploadId=upload_id
        )

    def copy_object(self, source_bucket: str, target_bucket: str, key: str) -> None:
        """Copy an object to the target s3 bucket.

        Args:
            source_bucket: S3 bucket name where the file currently exists
            target_bucket: S3 bucket name to which the file should be moved
            key: S3 object key
        """
        return self.s3_client.copy_object(
            Bucket=target_bucket,
            Key=key,
            CopySource={'Bucket': source_bucket, 'Key': key},
            MetadataDirective='COPY',
        )

    def delete_object(self, bucket: str, key: str) -> None:
        """Delete an object from S3.

        Args:
            bucket: S3 bucket name
            key: S3 object key
        """
        return self.s3_client.delete_object(Bucket=bucket, Key=key)

    def generate_complete_multipart_presigned_url(
        self,
        bucket: str,
        key: str,
        upload_id: str,
        expires_in: Optional[int] = None,
    ) -> str:
        """Generate a pre-signed URL for completing a multipart upload.

        Args:
            bucket: S3 bucket name
            key: S3 object key
            upload_id: Multipart upload ID
            expires_in: URL expiration in seconds

        Returns:
            Pre-signed URL string for CompleteMultipartUpload operation
        """
        if expires_in is None:
            expires_in = PRESIGNED_URL_EXPIRATION_SECS

        url = self.s3_client.generate_presigned_url(
            ClientMethod='complete_multipart_upload',
            ExpiresIn=expires_in,
            Params={
                'Bucket': bucket,
                'Key': key,
                'UploadId': upload_id,
            },
        )

        return url

    def generate_get_presigned_url(
        self, bucket: str, key: str, expires_in: Optional[int] = None
    ) -> str:
        """Generate a pre-signed URL for GET (download).

        Args:
            bucket: S3 bucket name
            key: S3 object key
            expires_in: URL expiration in seconds

        Returns:
            Pre-signed URL string
        """
        if expires_in is None:
            expires_in = PRESIGNED_URL_EXPIRATION_SECS

        return self.s3_client.generate_presigned_url(
            ClientMethod='get_object',
            ExpiresIn=expires_in,
            Params={'Bucket': bucket, 'Key': key},
        )

    def generate_multipart_presigned_urls(
        self,
        bucket: str,
        key: str,
        upload_id: str,
        total_parts: int,
        expires_in: Optional[int] = None,
    ) -> list[PresignedUrlPart]:
        """Generate pre-signed URLs for all parts of a multipart upload.

        Args:
            bucket: S3 bucket name
            key: S3 object key
            upload_id: Multipart upload ID
            total_parts: Total number of parts (part numbers are sequential 1..N)
            expires_in: URL expiration in seconds

        Returns:
            List of pre-signed URL parts
        """
        if expires_in is None:
            expires_in = PRESIGNED_URL_EXPIRATION_SECS

        presigned_urls = []
        expires_at = datetime.now(timezone.utc) + timedelta(seconds=expires_in)

        for part_number in range(1, total_parts + 1):
            url = self.s3_client.generate_presigned_url(
                'upload_part',
                Params={
                    'Bucket': bucket,
                    'Key': key,
                    'UploadId': upload_id,
                    'PartNumber': part_number,
                },
                ExpiresIn=expires_in,
            )

            presigned_urls.append(
                {
                    'part_number': part_number,
                    'url': url,
                    'expires_at': expires_at,
                }
            )

        return presigned_urls

    def generate_put_presigned_url(
        self,
        bucket: str,
        key: str,
        metadata: Optional[dict[str, str]] = None,
        content_md5: Optional[str] = None,
        content_type: Optional[str] = None,
        expires_in: Optional[int] = None,
    ) -> str:
        """Generate a pre-signed URL for single-part upload (PUT).

        Args:
            bucket: S3 bucket name
            key: S3 object key
            metadata: Optional metadata to attach to the object
            content_md5: Optional base64-encoded MD5 for S3 verification
            content_type: Optional MIME type for the object
            expires_in: URL expiration in seconds

        Returns:
            Pre-signed URL string
        """
        if expires_in is None:
            expires_in = PRESIGNED_URL_EXPIRATION_SECS

        params: dict[str, Any] = {'Bucket': bucket, 'Key': key}
        if metadata:
            params['Metadata'] = metadata
        if content_md5:
            params['ContentMD5'] = content_md5
        if content_type:
            params['ContentType'] = content_type

        url = self.s3_client.generate_presigned_url(
            ClientMethod='put_object',
            ExpiresIn=expires_in,
            Params=params,
        )

        return url

    def get_object_metadata(self, bucket: str, key: str) -> ObjectMetadata:
        """Get object metadata.

        Args:
            bucket: S3 bucket name
            key: S3 object key

        Returns:
            Dict with size, last_modified, etag, and metadata
        """
        response = self.s3_client.head_object(Bucket=bucket, Key=key)

        return {
            'size': response.get('ContentLength', 0),
            'last_modified': response.get('LastModified'),
            'etag': response.get('ETag', None),
            'metadata': response.get('Metadata', {}),
        }

    def initiate_multipart_upload(
        self,
        bucket: str,
        key: str,
        metadata: Optional[dict[str, str]] = None,
        content_type: Optional[str] = None,
    ) -> str:
        """Initiate a multipart upload.

        Args:
            bucket: S3 bucket name
            key: S3 object key
            metadata: Optional metadata to attach to the object
            content_type: Optional MIME type for the object

        Returns:
            Upload ID for the multipart upload
        """
        params: dict[str, Any] = {'Bucket': bucket, 'Key': key}
        if metadata:
            params['Metadata'] = metadata
        if content_type:
            params['ContentType'] = content_type

        response = self.s3_client.create_multipart_upload(**params)

        upload_id = response.get('UploadId')
        if not upload_id:
            raise ValueError(
                'Failed to initiate multipart upload: No upload ID returned'
            )

        return upload_id

    def object_exists(self, bucket: str, key: str) -> bool:
        """Check if an object exists in S3.

        Args:
            bucket: S3 bucket name
            key: S3 object key

        Returns:
            True if object exists, False otherwise
        """
        try:
            self.s3_client.head_object(Bucket=bucket, Key=key)
            return True
        except ClientError as e:
            if e.response['Error']['Code'] == '404':
                return False
            raise


def get_s3_client(config: Optional[Config] = None):
    """Get s3 client.

    Args:
        config: Optional botocore Config for the client

    Returns:
        Configured Boto3 S3 client instance
    """
    if config is None:
        #  SigV4 is required for KMS-encrypted buckets
        config = Config(signature_version='s3v4')
    return boto3.client('s3', config=config)


def get_s3_connector(s3_client=None) -> S3Connector:
    """Get S3 connector instance.

    Args:
        s3_client: Optional Boto3 S3 client instance

    Returns:
        Configured S3Connector instance
    """
    if not s3_client:
        s3_client = get_s3_client()
    return S3Connector(s3_client=s3_client)