from unittest.mock import ANY, MagicMock, patch
from pytest import fixture
from pytest import mark
from sientia_do.notifications.models import NotificationLevel

from scouter.activities.postgres import Postgres


@fixture
@patch("scouter.activities.postgres.create_engine")
@patch("scouter.activities.postgres.sessionmaker")
def postgres_client(mock_sessionmaker, mock_engine):
    # Create a mock session
    mock_session = MagicMock()
    mock_session.commit = MagicMock()
    mock_session.close = MagicMock()

    # Configure the session to work with context management
    mock_session.__enter__ = MagicMock(return_value=mock_session)
    mock_session.__exit__ = MagicMock(return_value=None)

    # Configure the sessionmaker to return our mock session
    mock_sessionmaker.return_value = mock_session

    # Configure the engine to return our mock sessionmaker
    mock_engine.return_value = MagicMock()
    mock_engine.return_value.dispose = MagicMock()

    # Create the Postgres client
    client = Postgres(
        host="localhost",
        port=5432,
        user="postgres",
        password="postgres",
        dbname="postgres",
        min_connections=1,
        max_connections=10,
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    # Set up the session factory
    client.session_factory = mock_sessionmaker

    return client


@mark.asyncio
@patch("scouter.activities.postgres.DataFrame")
async def test_export_data_to_postgres_success(mock_dataframe, postgres_client):
    data = {"schema": "test", "table_name": "test",
            "data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
    await postgres_client.export_data_to_postgres(data)

    # Verify notification handler wasn't called
    postgres_client.notification_handler.build_and_send_notification.assert_not_called()

    # Verify session handling
    mock_dataframe.assert_called_once_with(data["data"])
    mock_dataframe.return_value.to_sql.assert_called_once_with(
        data["table_name"],
        postgres_client.engine,
        schema=data["schema"],
        if_exists="append",
        index=False
    )
    postgres_client.session_factory.return_value.commit.assert_called_once()
    postgres_client.session_factory.return_value.close.assert_called_once()


@mark.asyncio
@patch("scouter.activities.postgres.DataFrame", return_value=MagicMock(
    to_sql=MagicMock(side_effect=Exception("Error exporting data to postgres"))
))
async def test_export_data_to_postgres_error(_mock_dataframe, postgres_client):
    data = {"schema": "test", "table_name": "test",
            "data": {"a": [1, 2, 3], "b": [4, 5, 6]}}
    await postgres_client.export_data_to_postgres(data)

    # Verify error notification was sent
    postgres_client.notification_handler.build_and_send_notification.assert_called_once_with(
        notification_id="ERROR_EXPORTING_DATA_TO_POSTGRES",
        message="Error exporting data to postgres: Error exporting data to postgres",
        block="export_data_to_postgres",
        level=NotificationLevel.ERROR,
        attachment_content=ANY
    )

    # Verify session handling
    postgres_client.session_factory.return_value.close.assert_called_once()