"""Tests for connection failure recovery and resilience."""

import pytest
from unittest.mock import AsyncMock, Mock, PropertyMock, patch

from metapg.pool import cursor, init_pool, transaction
from metapg.pool.pool import _async_connection, _sync_connection


class TestConnectionResilience:
    """Test connection failure recovery."""

    @pytest.mark.asyncio
    async def test_async_cursor_recovers_from_closed_connection(self):
        """Test that async cursor detects and recovers from closed connections."""
        # Initialize pool
        init_pool("postgresql://localhost/test", db_name="test_resilience")

        # Simulate first successful connection
        async with cursor("test_resilience") as cur:
            # This creates and stores a connection in context var
            pass

        # Now simulate the connection being closed (e.g., server timeout)
        # by directly accessing and "closing" the connection in the context var
        from metapg.pool.pool import _get_async_context_vars

        conn_var, cursor_var = _get_async_context_vars("test_resilience")

        # Mock a closed connection scenario
        mock_conn = AsyncMock()
        type(mock_conn).closed = PropertyMock(return_value=True)
        conn_var.set(mock_conn)

        # Mock cursor from closed connection
        mock_cursor = AsyncMock()
        type(mock_cursor).connection = PropertyMock(return_value=mock_conn)
        cursor_var.set(mock_cursor)

        # Now try to use cursor again - it should detect closed connection
        # and get a fresh one from the pool
        with patch("metapg.pool.pool._get_async_pool") as mock_get_pool:
            mock_pool = AsyncMock()
            mock_pool.closed = False
            mock_new_conn = AsyncMock()
            type(mock_new_conn).closed = PropertyMock(return_value=False)
            mock_new_cursor = AsyncMock()

            mock_pool.connection.return_value.__aenter__.return_value = mock_new_conn
            mock_new_conn.cursor.return_value.__aenter__.return_value = mock_new_cursor
            mock_get_pool.return_value = mock_pool

            async with cursor("test_resilience") as cur:
                # Should get new cursor from pool, not the stale one
                assert cur == mock_new_cursor

    def test_sync_cursor_recovers_from_closed_connection(self):
        """Test that sync cursor detects and recovers from closed connections."""
        # Initialize pool
        init_pool("postgresql://localhost/test", db_name="test_sync_resilience")

        # Simulate first successful connection
        with cursor("test_sync_resilience") as cur:
            pass

        # Mock a closed connection scenario
        from metapg.pool.pool import _get_sync_context_vars

        conn_var, cursor_var = _get_sync_context_vars("test_sync_resilience")

        mock_conn = Mock()
        type(mock_conn).closed = PropertyMock(return_value=True)
        conn_var.set(mock_conn)

        mock_cursor = Mock()
        type(mock_cursor).connection = PropertyMock(return_value=mock_conn)
        cursor_var.set(mock_cursor)

        # Try to use cursor - should get fresh connection
        with patch("metapg.pool.pool._get_sync_pool") as mock_get_pool:
            mock_pool = Mock()
            mock_new_conn = Mock()
            type(mock_new_conn).closed = PropertyMock(return_value=False)
            mock_new_cursor = Mock()

            mock_pool.connection.return_value.__enter__.return_value = mock_new_conn
            mock_new_conn.cursor.return_value.__enter__.return_value = mock_new_cursor
            mock_get_pool.return_value = mock_pool

            with cursor("test_sync_resilience") as cur:
                assert cur == mock_new_cursor

    @pytest.mark.asyncio
    async def test_async_connection_clears_stale_context(self):
        """Test that stale connections are cleared from context vars."""
        from metapg.pool.pool import _get_async_context_vars

        conn_var, _ = _get_async_context_vars("test_stale")

        # Set a closed connection in context
        mock_conn = AsyncMock()
        type(mock_conn).closed = PropertyMock(return_value=True)
        conn_var.set(mock_conn)

        # Mock pool to return new connection
        with patch("metapg.pool.pool._get_async_pool") as mock_get_pool:
            mock_pool = AsyncMock()
            mock_pool.closed = False
            mock_new_conn = AsyncMock()
            type(mock_new_conn).closed = PropertyMock(return_value=False)

            mock_pool.connection.return_value.__aenter__.return_value = mock_new_conn
            mock_get_pool.return_value = mock_pool

            async with _async_connection("test_stale") as conn:
                # Should get new connection
                assert conn == mock_new_conn
                # Stale connection should be cleared
                # (context var will be set to new connection)

    def test_sync_connection_clears_stale_context(self):
        """Test that stale sync connections are cleared from context vars."""
        from metapg.pool.pool import _get_sync_context_vars

        conn_var, _ = _get_sync_context_vars("test_sync_stale")

        # Set a closed connection in context
        mock_conn = Mock()
        type(mock_conn).closed = PropertyMock(return_value=True)
        conn_var.set(mock_conn)

        # Mock pool to return new connection
        with patch("metapg.pool.pool._get_sync_pool") as mock_get_pool:
            mock_pool = Mock()
            mock_new_conn = Mock()
            type(mock_new_conn).closed = PropertyMock(return_value=False)

            mock_pool.connection.return_value.__enter__.return_value = mock_new_conn
            mock_get_pool.return_value = mock_pool

            with _sync_connection("test_sync_stale") as conn:
                # Should get new connection
                assert conn == mock_new_conn

    @pytest.mark.asyncio
    async def test_async_transaction_with_closed_connection(self):
        """Test that transactions recover from closed connections."""
        init_pool("postgresql://localhost/test", db_name="test_tx")

        # Set up a closed connection in context
        from metapg.pool.pool import _get_async_context_vars

        conn_var, _ = _get_async_context_vars("test_tx")
        mock_conn = AsyncMock()
        type(mock_conn).closed = PropertyMock(return_value=True)
        conn_var.set(mock_conn)

        # Mock pool to return new connection
        with patch("metapg.pool.pool._get_async_pool") as mock_get_pool:
            mock_pool = AsyncMock()
            mock_pool.closed = False
            mock_new_conn = AsyncMock()
            type(mock_new_conn).closed = PropertyMock(return_value=False)
            mock_tx = AsyncMock()

            mock_pool.connection.return_value.__aenter__.return_value = mock_new_conn
            mock_new_conn.transaction.return_value.__aenter__.return_value = mock_tx
            mock_get_pool.return_value = mock_pool

            # Should successfully create transaction with new connection
            async with transaction("test_tx"):
                pass

            # Verify transaction was created with new connection
            mock_new_conn.transaction.assert_called_once()
