from io import BytesIO
from unittest.mock import Mock
import brotli
import msgpack
import pydantic
import pytest
import zstandard
import json
from common.src.connectors import snowflake as sf
from common.src.enums import AWSPayload, Formats
from common.src.models.sql import BaseColumns
from common.src.typings import SqlString


class DummyModel(BaseColumns):
    uid: int
    name: str


class TestRunSQL:

    _class = sf._RunSQL

    # noinspection SqlNoDataSourceInspection
    sample_sql = SqlString("SELECT * FROM table")

    @pytest.fixture
    def instance(self):
        return self._class()

    def setup_method(self):
        self._class.lambda_client = None

    def test_run_invalid_format(self, instance):
        with pytest.raises(ValueError):
            instance.run(self.sample_sql, output_format="invalid")  # type: ignore

    def test_run(self, mocker, instance):
        mock_lambda_client = Mock()
        mocker.patch.object(
            sf.aws,
            "new_lambda_client",
            autospec=True,
            return_value=mock_lambda_client,
        )
        mocker.patch.object(instance, "_handle_response", autospec=True)
        instance.run(self.sample_sql, output_format=Formats.JSON)
        assert mock_lambda_client.invoke.call_count == 1
        assert instance._handle_response.called

    def test_handle_response_csv(self, instance):
        csv_data = "col1,col2\n1,A\n2,B"
        compressed = brotli.compress(csv_data.encode("utf-8"))
        response = {AWSPayload.PAYLOAD: BytesIO(compressed)}

        result = instance._handle_response(response, output_format=Formats.CSV)
        result_csv_data = "".join(result)
        assert (
            result_csv_data == csv_data
        ), "Expected generator outputted CSV data to match"

    mock_json_data = [{"uid": 1, "name": "Alice"}, {"uid": 2, "name": "Bob"}]

    def _mock_aws_response(self):
        packed = msgpack.packb(self.mock_json_data)
        compressed = brotli.compress(packed)
        return {AWSPayload.PAYLOAD: BytesIO(compressed)}

    def test_handle_response_json_no_model(self, instance):
        result = list(
            instance._handle_response(
                self._mock_aws_response(), output_format=Formats.JSON
            )
        )
        assert result == self.mock_json_data

    def test_handle_response_json_with_model(self, instance):
        result = list(
            instance._handle_response(
                self._mock_aws_response(),
                output_format=Formats.JSON,
                model=DummyModel,  # type: ignore
            )
        )
        assert all(isinstance(r, DummyModel) for r in result)
        assert result == [
            DummyModel(uid=1, name="Alice"),
            DummyModel(uid=2, name="Bob"),
        ]

    def test_handle_response_staged(self, mocker, instance):
        mock_response = self._mock_aws_response()
        mock_result = mocker.Mock()
        mocker.patch.object(
            instance, "_handle_zstd_csv", return_value=mock_result, autospec=True
        )
        result = instance._handle_response(mock_response, staged=True)
        assert result is mock_result, (
            "Expected staged response to call _handle_zstd_csv and return "
            "the result of that function"
        )

    def test_handle_response_staged_s3(self, instance):
        mock_s3_file_names = ["s3://bucket/file1.csv.zst", "s3://bucket/file2.csv.zst"]
        response_bytes = json.dumps(mock_s3_file_names).encode("utf-8")
        response = {AWSPayload.PAYLOAD: BytesIO(response_bytes)}
        result = instance._handle_response(response, staged="S3")
        assert list(result) == mock_s3_file_names, (
            "Expected staged response to return the list of S3 file names"
            " without any further processing"
        )

    @pytest.mark.parametrize(
        "csv_data,expected",
        [
            ('"COL1","COL2"\n"1","A"\n"2","B"', '"col1","col2"\n"1","A"\n"2","B"'),
            ('"COL1","COL2"\n"1","A"\n"2","B"\n', '"col1","col2"\n"1","A"\n"2","B"\n'),
        ],
    )
    def test_handle_zstd_csv(self, instance, csv_data, expected):
        compressed = zstandard.compress(csv_data.encode("utf-8"))
        result = instance._handle_zstd_csv(compressed)
        result_csv_data = "".join(result)
        assert (
            result_csv_data == expected
        ), "Expected generator outputted CSV data to match"

    def test_handle_csv(self, instance):
        csv_data = "COL1,COL2\n1,A\n2,B"
        expected_csv = "col1,col2\n1,A\n2,B"  # Normalize to lowercase
        result = instance._handle_csv(csv_data.encode("utf-8"))
        result_csv_data = "".join(result)
        assert result_csv_data == expected_csv, "Expected CSV header to be lowercase"

    @pytest.mark.parametrize(
        "raw,expect_error",
        [
            (b"valid compressed", False),
            (b"invalid data", True),
        ],
    )
    def test_decompress_payload_content(self, instance, raw, expect_error):
        data = brotli.compress(raw) if not expect_error else raw
        func = lambda: instance._decompress_payload_content(data)  # noqa: E731
        if expect_error:
            with pytest.raises(brotli.error):
                func()
        else:
            result = func()
            assert result == raw

    @pytest.mark.parametrize(
        "data,model,expect_error",
        [
            ([{"uid": 1, "name": "Alice"}, {"uid": 2, "name": "Bob"}], None, False),
            (
                [{"uid": 1, "name": "Alice"}, {"uid": 2, "name": "Bob"}],
                DummyModel,
                False,
            ),
            ([{"uid": 1}, {"uid": 2}], DummyModel, True),  # Missing 'name' field
        ],
    )
    def test_handle_decompressed_json_payload_content(
        self, instance, data, model, expect_error
    ):
        func = lambda: list(  # noqa: E731
            instance._handle_decompressed_json_payload_content(
                msgpack.packb(data), model
            )
        )
        if expect_error:
            with pytest.raises(pydantic.ValidationError):
                func()
        else:
            result = func()
            if model:
                assert all(
                    isinstance(r, model) for r in result
                ), "Expected all results to be wrapped in the provided model"
            else:
                assert result == data, "Expected raw data without model wrapping"