import json
from itertools import chain
from unittest.mock import AsyncMock, MagicMock

import pytest
import pytest_asyncio

from src.backend.connectors.db.components import audit_rows

sample_add_args = {
    "audit_id": 1,
    audit_rows.DBColumns.ROW_IDX: 1,
    "row_type": "row_type",
    "row_data": {"key": "value"},
}


class TestAuditRows:
    _class = audit_rows.AuditRows

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

    class TestAdd:
        @pytest.mark.asyncio
        async def test_add(self, instance):
            args = sample_add_args
            await instance.add(**args)
            assert len(instance.rows) == 1
            first_row = instance.rows[0]
            assert first_row.audit == args["audit_id"]
            assert first_row.idx == args[audit_rows.DBColumns.ROW_IDX]
            assert first_row.type == args["row_type"]
            assert first_row.data == args["row_data"]

        @pytest.mark.asyncio
        async def test_add_uses_asyncio_lock(self, instance):
            """Ensure the asyncio.Lock is used to avoid concurrent access to the
            rows list (i.e. simultaneous addition and saving of rows)."""
            instance._rows_lock = AsyncMock()
            await instance.add(
                "audit_id", audit_rows.DBColumns.ROW_IDX, "row_type", {"key": "value"}
            )
            assert instance._rows_lock.__aenter__.called

    class TestGetByAuditGroupId:
        expected_data = {"key": "value"}

        @pytest_asyncio.fixture
        async def mocker_execute_query_fetchall(self, mocker, instance):
            sample_row = {
                **sample_add_args,
                "data": json.dumps(self.expected_data),
            }

            mock_client = mocker.AsyncMock()
            mock_client.db.execute_query_fetchall.return_value = [sample_row]
            instance.client = mock_client

        @pytest_asyncio.fixture
        async def get_by_audit_group_id(self, mocker_execute_query_fetchall, instance):
            return await instance.get_by_audit_group_id(123)

        @pytest.mark.asyncio
        async def test_get_by_audit_group_id_returns(self, get_by_audit_group_id):
            assert isinstance(get_by_audit_group_id, list)
            assert len(get_by_audit_group_id) == 1

        @pytest.mark.asyncio
        async def test_get_by_audit_group_id_row_content(self, get_by_audit_group_id):
            row = get_by_audit_group_id[0]
            assert row["audit_id"] == sample_add_args["audit_id"]
            assert (
                row[audit_rows.DBColumns.ROW_IDX]
                == sample_add_args[audit_rows.DBColumns.ROW_IDX]
            )
            assert row["row_type"] == sample_add_args["row_type"]
            assert row["row_data"] == self.expected_data

    class TestSave:
        @pytest_asyncio.fixture
        async def mocker_client(self, instance):
            mock_conn = MagicMock()

            mock_conn.__enter__.return_value = mock_conn
            mock_conn.__exit__.return_value = None
            mock_conn.commit = MagicMock()
            mock_conn.rollback = MagicMock()

            mock_client = MagicMock()
            mock_client.db.get_connection = AsyncMock(return_value=mock_conn)

            instance.client = mock_client
            return mock_client

        @pytest_asyncio.fixture
        async def save(self, mocker_client, instance):
            async def _(error: bool = False):
                if error:
                    mocker_client.db.get_connection.return_value.commit.side_effect = (
                        Exception
                    )
                return await instance.save()

            return _

        @pytest.mark.asyncio
        async def test_save_commit(self, save, mocker_client):
            await save()
            assert mocker_client.db.get_connection.return_value.commit.called

        @pytest.mark.asyncio
        async def test_save_rollback_if_error(self, save, mocker_client):
            with pytest.raises(Exception):
                await save(error=True)
            assert mocker_client.db.get_connection.return_value.rollback.called

    class TestPrepareRowsForSave:
        sample_types = {"type A", "type B", "type C"}

        @property
        def sample_rows(self):
            return [
                audit_rows.Row(
                    audit=1,
                    idx=1,
                    type=type_,
                    data={"key": "value"},
                )
                for type_ in self.sample_types
            ]

        @pytest.mark.asyncio
        async def test_prepare_rows_for_save_grouped_by_type(self, instance):
            row_count_per_type = 3
            instance.rows.extend(self.sample_rows * row_count_per_type)
            result = await instance._prepare_rows_for_save()
            assert result.keys() == set(self.sample_types)
            for type_ in self.sample_types:
                assert len(result[type_]) == row_count_per_type, (
                    f"Expected {row_count_per_type} rows for type {type_}, "
                    f"got {len(result[type_])}"
                )

        @pytest.mark.asyncio
        @pytest.mark.parametrize("row_idx_present", [True, False])
        async def test_prepare_rows_for_save_row_idx_removed_from_data_if_present(
            self, instance, row_idx_present
        ):
            sample_rows = self.sample_rows
            if row_idx_present:
                for row in sample_rows:
                    row.data[audit_rows.DBColumns.ROW_IDX] = 1
            instance.rows.extend(sample_rows)
            result = await instance._prepare_rows_for_save()
            individual_result_rows = list(chain.from_iterable(result.values()))
            assert all(
                audit_rows.DBColumns.ROW_IDX not in row.data
                for row in individual_result_rows
            ), "ROW_IDX should be removed from row data if present"

        @pytest.mark.asyncio
        async def test_prepare_rows_for_save_rows_cleared(self, instance):
            instance.rows = self.sample_rows
            await instance._prepare_rows_for_save()
            assert instance.rows == []

        @pytest.mark.asyncio
        async def test_prepare_rows_for_save_uses_asyncio_lock(self, instance):
            """Ensure the asyncio.Lock is used to avoid concurrent access to the
            rows list (i.e. simultaneous addition and saving of rows)."""
            instance._rows_lock = AsyncMock()
            await instance._prepare_rows_for_save()
            assert instance._rows_lock.__aenter__.called


@pytest.mark.parametrize(
    "value,expected",
    [
        ({"key": "value"}, '{"key": "value"}'),
        ({"key": ["value"]}, '{"key": ["value"]}'),
        ({"key": {"key": "value"}}, '{"key": {"key": "value"}}'),
        ({"key": (1, 2, 3)}, '{"key": [1, 2, 3]}'),
        ({"key": {1, 2, 3}}, '{"key": [1, 2, 3]}'),
        ({"key": float("nan")}, '{"key": null}'),  # Ensure NaN is converted to null
    ],
)
def test_prepare_data(value, expected):
    """Test the _prepare_data function."""
    assert audit_rows._prepare_data(value) == expected