import base64
import hashlib
from typing import TYPE_CHECKING

import boto3
from botocore.errorfactory import ClientError

from aws_testing_utils import config

from .logger import log

if TYPE_CHECKING:
    from mypy_boto3_s3 import S3Client


class S3Handler:
    """Handles interactions with S3 buckets and objects."""

    client: 'S3Client'

    def __init__(self) -> None:
        session = boto3.session.Session()
        self.client = session.client(service_name='s3', region_name=config.AWS_REGION)

    def check_object_exists(self, bucket: str, object_name: str) -> bool:
        """Returns True if the object exists in the bucket."""
        try:
            self.client.head_object(Bucket=bucket, Key=object_name)
            return True
        except ClientError:
            return False

    def assert_object_exists(self, bucket: str, object_name: str) -> None:
        """Asserts the object exists in the bucket."""
        assert self.check_object_exists(bucket, object_name), (
            '{} not found in {} bucket'.format(object_name, bucket)
        )

    def put(self, bucket: str, file_path: str, file_data: bytes | bytearray) -> None:
        """Uploads an object to S3.

        Args:
            bucket: S3 bucket name.
            file_path: S3 object key.
            file_data: Raw bytes to upload.
        """
        self.client.put_object(Body=bytes(file_data), Bucket=bucket, Key=file_path)

    def put_if_object_not_present(
        self, bucket: str, file_path: str, file_data: bytes | bytearray
    ) -> None:
        """Uploads an object to S3 only if it does not already exist.

        Args:
            bucket: S3 bucket name.
            file_path: S3 object key.
            file_data: Raw bytes to upload.
        """
        if not self.check_object_exists(bucket, file_path):
            self.put(bucket, file_path, file_data)
        else:
            log.info(f'{file_path} already in {bucket}')

    def delete_if_object_present(self, bucket: str, file_path: str) -> None:
        """Deletes an object from S3 if it exists.

        Args:
            bucket: S3 bucket name.
            file_path: S3 object key.
        """
        if self.check_object_exists(bucket, file_path):
            self.client.delete_object(Bucket=bucket, Key=file_path)
        else:
            log.info(f'{file_path} not found in {bucket}, no need to delete')

    @staticmethod
    def calculate_sha256_base64(file_data: bytes | bytearray) -> str:
        """Returns the base64-encoded SHA256 checksum of the given data."""
        sha256 = hashlib.sha256(file_data).digest()
        return base64.b64encode(sha256).decode('utf-8')

    def get_s3_object_data(self, bucket: str, file_path: str) -> bytes:
        """Downloads and returns the raw bytes of an S3 object.

        Args:
            bucket: S3 bucket name.
            file_path: S3 object key.
        """
        response = self.client.get_object(Bucket=bucket, Key=file_path)
        file_data = response['Body'].read()

        return file_data

    def verify_s3_object_sha256(
        self, bucket: str, file_path: str, stored_sha256: str | None = None
    ) -> None:
        """Verifies the SHA256 checksum of an S3 object.

        Checks integrity by comparing the object's computed checksum against
        the AWS-stored checksum metadata, or against stored_sha256 if metadata
        is unavailable.

        Args:
            bucket: S3 bucket name.
            file_path: S3 object key.
            stored_sha256: Expected base64-encoded checksum, used when AWS metadata is unavailable.

        Raises:
            ValueError: If no checksum is available for comparison.
            AssertionError: If the computed checksum does not match the expected value.

        Note:
            AWS errors during metadata fetch or object download are logged and cause
            an early return without asserting, rather than raising.
        """
        # Step 1: Try to get checksum metadata
        try:
            metadata = self.client.head_object(Bucket=bucket, Key=file_path)
            metadata_sha256 = metadata.get('ChecksumSHA256')  # base64-encoded string
        except ClientError as e:
            log.error(f'❌ Error fetching metadata for {bucket}/{file_path}: {e}')
            return None
        except Exception as e:
            log.error(
                f'❌ Unexpected error fetching metadata for: {bucket}/{file_path}: {e}'
            )
            return None

        if not metadata_sha256 and not stored_sha256:
            raise ValueError(
                f'Cannot verify {bucket}/{file_path}: no AWS checksum metadata and no stored_sha256 provided.'
            )

        # Step 2: Download the object
        try:
            file_data = self.get_s3_object_data(bucket, file_path)
        except ClientError as e:
            log.error(f'❌ Error downloading object {bucket}/{file_path}: {e}')
            return None
        except Exception as e:
            log.error(
                f'❌ Unexpected error downloading object: {bucket}/{file_path}: {e}'
            )
            return None

        # Step 3: Compute local SHA256 (base64-encoded)
        computed_sha256 = self.calculate_sha256_base64(file_data)

        # Step 4: Compare with metadata or expected
        if metadata_sha256:
            log.info(
                f'Verifying S3 object integrity for {bucket}/{file_path} using AWS S3 checksum.'
            )
            assert metadata_sha256 == computed_sha256, (
                f'❌ Checksum mismatch for {bucket}/{file_path}: expected metadata {metadata_sha256}, but got {computed_sha256}'
            )
        elif stored_sha256:
            log.info(
                f'Verifying S3 object integrity for {bucket}/{file_path} using provided checksum.'
            )
            assert stored_sha256 == computed_sha256, (
                f'❌ Checksum mismatch for {bucket}/{file_path}: stored checksum is {stored_sha256}, but got {computed_sha256}'
            )