import base64
import hashlib
from unittest import mock

import pytest

from aws_testing_utils import s3_handler
from tests.conftest import get_client_error, mock_session_setup


@mock.patch('boto3.session.Session')
def test_check_object_exists(mock_session_class: mock.MagicMock) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.head_object.return_value = None
    handler = s3_handler.S3Handler()
    assert handler.check_object_exists('testBucket', 'object')


@mock.patch('boto3.session.Session')
def test_check_object_exists_error(mock_session_class: mock.MagicMock) -> None:
    error = get_client_error()
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.head_object.side_effect = error
    handler = s3_handler.S3Handler()
    assert not handler.check_object_exists('testBucket', 'object')


@mock.patch('aws_testing_utils.s3_handler.S3Handler.check_object_exists')
def test_assert_object_exists(mock_check_object_exists: mock.MagicMock) -> None:
    mock_check_object_exists.return_value = True
    handler = s3_handler.S3Handler()
    bucket = 'bucket'
    obj_name = 'object'
    handler.assert_object_exists(bucket, obj_name)
    mock_check_object_exists.assert_called_with(bucket, obj_name)


@mock.patch('aws_testing_utils.s3_handler.S3Handler.check_object_exists')
def test_assert_object_exists_error(mock_check_object_exists: mock.MagicMock) -> None:
    mock_check_object_exists.return_value = False
    handler = s3_handler.S3Handler()
    bucket = 'bucket'
    obj_name = 'object'
    with pytest.raises(AssertionError):
        handler.assert_object_exists(bucket, obj_name)
    mock_check_object_exists.assert_called_with(bucket, obj_name)


@mock.patch('boto3.session.Session')
def test_put(mock_session_class: mock.MagicMock) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.put_object.return_value = None
    bucket = 'bucket'
    file_path = 'path'
    file_data = bytearray(123)
    handler = s3_handler.S3Handler()
    handler.put(bucket, file_path, file_data)
    mock_client.put_object.assert_called_with(
        Body=file_data, Bucket=bucket, Key=file_path
    )


@mock.patch('aws_testing_utils.s3_handler.S3Handler.check_object_exists')
@mock.patch('aws_testing_utils.s3_handler.S3Handler.put')
def test_put_if_object_not_present_obj_present(
    mock_put: mock.MagicMock, mock_check_object_exists: mock.MagicMock
) -> None:
    mock_check_object_exists.return_value = True
    handler = s3_handler.S3Handler()
    data = bytearray(123)
    handler.put_if_object_not_present('bucket', 'path', data)
    assert not mock_put.called


@mock.patch('aws_testing_utils.s3_handler.S3Handler.check_object_exists')
@mock.patch('aws_testing_utils.s3_handler.S3Handler.put')
def test_put_if_object_not_present_obj_not_present(
    mock_put: mock.MagicMock, mock_check_object_exists: mock.MagicMock
) -> None:
    mock_check_object_exists.return_value = False
    handler = s3_handler.S3Handler()
    bucket = 'bucket'
    file_path = 'path'
    data = bytearray(123)
    handler.put_if_object_not_present(bucket, file_path, data)
    mock_put.assert_called_with(bucket, file_path, data)


@mock.patch('aws_testing_utils.s3_handler.S3Handler.check_object_exists')
@mock.patch('boto3.session.Session')
def test_delete_if_object_present_obj_present(
    mock_session_class: mock.MagicMock, mock_check_object_exists: mock.MagicMock
) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_client.delete_object.return_value = None
    mock_check_object_exists.return_value = True
    handler = s3_handler.S3Handler()
    bucket = 'bucket'
    file_path = 'path'
    handler.delete_if_object_present(bucket, file_path)
    mock_client.delete_object.assert_called_with(Bucket=bucket, Key=file_path)


@mock.patch('aws_testing_utils.s3_handler.S3Handler.check_object_exists')
@mock.patch('boto3.session.Session')
def test_delete_if_object_present_obj_not_present(
    mock_session_class: mock.MagicMock, mock_check_object_exists: mock.MagicMock
) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    mock_check_object_exists.return_value = False
    handler = s3_handler.S3Handler()
    bucket = 'bucket'
    file_path = 'path'
    handler.delete_if_object_present(bucket, file_path)
    assert not mock_client.delete_object.called


@mock.patch('boto3.session.Session')
def test_verify_s3_object_sha256_with_metadata_success(
    mock_session_class: mock.MagicMock,
) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    file_data = b'hello world'
    sha256 = hashlib.sha256(file_data).digest()
    sha256_b64 = base64.b64encode(sha256).decode('utf-8')

    mock_client.head_object.return_value = {'ChecksumSHA256': sha256_b64}
    mock_body = mock.Mock()
    mock_body.read.return_value = file_data
    mock_client.get_object.return_value = {'Body': mock_body}

    handler = s3_handler.S3Handler()
    # Should not raise
    handler.verify_s3_object_sha256('bucket', 'path')


@mock.patch('boto3.session.Session')
def test_verify_s3_object_sha256_with_metadata_mismatch(
    mock_session_class: mock.MagicMock,
) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    file_data = b'hello world'
    mock_client.head_object.return_value = {'ChecksumSHA256': 'invalid_checksum'}
    mock_body = mock.Mock()
    mock_body.read.return_value = file_data
    mock_client.get_object.return_value = {'Body': mock_body}

    handler = s3_handler.S3Handler()
    with pytest.raises(AssertionError):
        handler.verify_s3_object_sha256('bucket', 'path')


@mock.patch('boto3.session.Session')
def test_verify_s3_object_sha256_with_stored_checksum_success(
    mock_session_class: mock.MagicMock,
) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    file_data = b'another content'
    sha256 = hashlib.sha256(file_data).digest()
    sha256_b64 = base64.b64encode(sha256).decode('utf-8')

    # Simulate no checksum in S3 metadata
    mock_client.head_object.return_value = {}
    mock_body = mock.Mock()
    mock_body.read.return_value = file_data
    mock_client.get_object.return_value = {'Body': mock_body}

    handler = s3_handler.S3Handler()
    # Pass the expected checksum explicitly
    handler.verify_s3_object_sha256('bucket', 'path', stored_sha256=sha256_b64)


@mock.patch('boto3.session.Session')
def test_verify_s3_object_sha256_with_stored_checksum_mismatch(
    mock_session_class: mock.MagicMock,
) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    file_data = b'another content'

    # Simulate no checksum in S3 metadata
    mock_client.head_object.return_value = {}
    mock_body = mock.Mock()
    mock_body.read.return_value = file_data
    mock_client.get_object.return_value = {'Body': mock_body}

    handler = s3_handler.S3Handler()
    with pytest.raises(AssertionError):
        handler.verify_s3_object_sha256(
            'bucket', 'path', stored_sha256='wrong_checksum'
        )


@mock.patch('boto3.session.Session')
def test_verify_s3_object_sha256_no_checksum_available(
    mock_session_class: mock.MagicMock,
) -> None:
    mock_session_class, mock_session_object, mock_client = mock_session_setup(
        mock_session_class
    )
    file_data = b'some content'
    mock_client.head_object.return_value = {}
    mock_body = mock.Mock()
    mock_body.read.return_value = file_data
    mock_client.get_object.return_value = {'Body': mock_body}

    handler = s3_handler.S3Handler()
    with pytest.raises(ValueError):
        handler.verify_s3_object_sha256('bucket', 'path')