from collections.abc import Generator
from io import BytesIO
from typing import Any
from unittest.mock import MagicMock, patch

import pytest
from botocore.exceptions import ClientError
from botocore.response import StreamingBody

from backfill.connectors.s3_connector import S3, FailedDownloadException


def fake_streaming_body(content: bytes) -> StreamingBody:
    stream = BytesIO(content)
    return StreamingBody(stream, len(content))


@pytest.fixture
def mock_boto3_client() -> Generator[MagicMock, None, None]:
    with patch("backfill.connectors.s3_connector.boto3") as mock_boto3:
        yield mock_boto3.client.return_value


@pytest.fixture()
def s3_connector(mock_boto3_client: MagicMock) -> S3:
    return S3()


@pytest.mark.parametrize(
    "expected, return_value, should_raise",
    [
        pytest.param(True, {"ResponseMetadata": {}}, False, id="File exists"),
        pytest.param(
            False,
            {"NotResponseMetadata": {}},
            False,
            id="There is no ResponseMetadata key",
        ),
        pytest.param(
            False,
            {"NotResponseMetadata": {}},
            False,
            id="There is no ResponseMetadata key",
        ),
        pytest.param(
            False,
            {"NotResponseMetadata": {}},
            False,
            id="There is no ResponseMetadata key",
        ),
        pytest.param(
            False,
            {"Error": {"Code": "404", "Message": "Not Found"}},
            False,
            id="File does not exist raised but 404 is controlled",
        ),
        pytest.param(
            False,
            {"Error": {"Code": "403", "Message": "Forbidden"}},
            True,
            id="Other ClientError raised",
        ),
    ],
)
def test_does_file_exist(
    s3_connector: S3,
    mock_boto3_client: MagicMock,
    expected: bool,
    return_value: dict[str, Any],
    should_raise: bool,
) -> None:
    mock_head_object = mock_boto3_client.head_object
    if should_raise:
        mock_head_object.side_effect = ClientError(
            error_response=return_value,  # type: ignore[arg-type]
            operation_name="head_object",
        )
        with pytest.raises(ClientError):
            s3_connector.does_file_exist("bucket", "key")
    else:
        mock_head_object.return_value = return_value
        assert s3_connector.does_file_exist("bucket", "key") == expected

    mock_head_object.assert_called_once_with(Bucket="bucket", Key="key")


@pytest.mark.parametrize(
    "bucket, key, local_path, should_raise",
    [
        pytest.param(
            "test-bucket",
            "test-key",
            "local-path",
            False,
            id="File downloaded successfully",
        ),
        pytest.param(
            "test-bucket",
            "test-key",
            "local-path",
            True,
            id="File download raises FailedDownloadException",
        ),
    ],
)
def test_download_file(
    s3_connector: S3,
    mock_boto3_client: MagicMock,
    bucket: str,
    key: str,
    local_path: str,
    should_raise: bool,
) -> None:
    mock_download_file = mock_boto3_client.download_file

    if should_raise:
        mock_download_file.side_effect = ClientError(
            error_response={"Error": {"Code": "400", "Message": "General error"}},
            operation_name="download_file",
        )
        with pytest.raises(FailedDownloadException):
            s3_connector.download_file(bucket, key, local_path)
    else:
        mock_download_file.return_value = None
        result = s3_connector.download_file(bucket, key, local_path)
        assert result == local_path

    mock_download_file.assert_called_once_with(bucket, key, local_path)


@pytest.mark.parametrize(
    "bucket, key, return_value, should_raise",
    [
        pytest.param(
            "test-bucket",
            "test-key",
            {"Body": MagicMock(read=lambda: b"file content")},
            False,
            id="File content retrieved successfully",
        ),
        pytest.param(
            "test-bucket",
            "test-key",
            {"Error": {"Code": "404", "Message": "Not Found"}},
            True,
            id="File retrieval raises ClientError",
        ),
    ],
)
def test_get_file_content(
    s3_connector: S3,
    mock_boto3_client: MagicMock,
    bucket: str,
    key: str,
    return_value: dict[str, Any],
    should_raise: bool,
) -> None:
    mock_get_object = mock_boto3_client.get_object

    if should_raise:
        mock_get_object.side_effect = ClientError(
            error_response=return_value,  # type: ignore[arg-type]
            operation_name="get_object",
        )
        with pytest.raises(ClientError):
            s3_connector.get_file_content(bucket, key)
    else:
        mock_get_object.return_value = return_value
        result = s3_connector.get_file_content(bucket, key)
        assert result == "file content"

    mock_get_object.assert_called_once_with(Bucket=bucket, Key=key)