from unittest.mock import MagicMock, patch
from pytest import fixture, mark
import pandas as pd
from laborious.activities.postgres import Postgres


@fixture
@patch("laborious.activities.postgres.create_engine")
def postgres_activity(_mock_create_engine):
    return Postgres(
        host="localhost",
        port=5432,
        user="test_user",
        password="test_password",
        dbname="test_db",
        min_connections=1,
        max_connections=5,
        logger=MagicMock(),
        notification_handler=MagicMock()
    )


@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_none_data(mock_read_sql_query, postgres_activity):
    query = "SELECT * FROM test_table LIMIT 1"
    mock_read_sql_query.return_value = None

    result = await postgres_activity.load_custom_query(query)

    assert isinstance(result, dict)
    assert len(result) == 0


@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_date_converted(mock_read_sql_query, postgres_activity):
    query = "SELECT * FROM test_table LIMIT 1"
    mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})
    mock_data['date'] = pd.to_datetime('2022-01-01')

    mock_read_sql_query.return_value = mock_data

    result = await postgres_activity.load_custom_query(query)

    assert isinstance(result, dict)
    assert len(result) == 3
    assert "column1" in result
    assert "column2" in result
    assert "date" in result
    assert result['date'] == {0: '2022-01-01 00:00:00'}


@mark.asyncio
@patch("laborious.activities.postgres.read_sql_query")
async def test_load_custom_query_success(mock_read_sql_query, postgres_activity):
    query = "SELECT * FROM test_table LIMIT 1"
    mock_data = pd.DataFrame({"column1": [1], "column2": ["test"]})

    mock_read_sql_query.return_value = mock_data

    result = await postgres_activity.load_custom_query(query)

    assert isinstance(result, dict)
    assert len(result) == 2
    assert "column1" in result
    assert "column2" in result
    postgres_activity.logger.info.assert_called()


@mark.asyncio
async def test_load_custom_query_error(postgres_activity):
    query = "SELECT * FROM non_existent_table"
    error_msg = "Table not found"

    with patch("laborious.activities.postgres.read_sql_query", side_effect=ValueError(error_msg)):
        result = await postgres_activity.load_custom_query(query)

    assert isinstance(result, dict)
    assert len(result) == 0
    postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
    postgres_activity.logger.error.assert_called()


@mark.asyncio
async def test_repeat_last_prediction_success(postgres_activity):
    query_items = {
        "schema": "public",
        "table_name": "predictions",
        "model": 1
    }

    with patch("sqlalchemy.orm.session.Session.execute") as mock_execute:
        await postgres_activity.repeat_last_prediction(query_items)

    mock_execute.assert_called_once()
    postgres_activity.logger.info.assert_called()


@mark.asyncio
async def test_repeat_last_prediction_error(postgres_activity):
    query_items = {
        "schema": "public",
        "table_name": "predictions",
        "model": 1
    }
    error_msg = "Database error"

    with patch("sqlalchemy.orm.session.Session.execute", side_effect=ValueError(error_msg)):
        await postgres_activity.repeat_last_prediction(query_items)

    postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
    postgres_activity.logger.error.assert_called()


@mark.asyncio
async def test_export_data_to_postgres_success(postgres_activity):
    input_data = {
        "schema": "public",
        "table_name": "test_table",
        "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
    }

    with patch("laborious.activities.postgres.DataFrame.to_sql") as mock_to_sql:
        await postgres_activity.export_data_to_postgres(input_data)

    mock_to_sql.assert_called_once()
    postgres_activity.logger.debug.assert_called()


@mark.asyncio
async def test_export_data_to_postgres_error(postgres_activity):
    input_data = {
        "schema": "public",
        "table_name": "test_table",
        "data": pd.DataFrame({"column1": [1, 2], "column2": ["a", "b"]})
    }
    error_msg = "Export failed"

    with patch("laborious.activities.postgres.DataFrame.to_sql", side_effect=ValueError(error_msg)):
        await postgres_activity.export_data_to_postgres(input_data)

    postgres_activity.notification_handler.build_and_send_notification.assert_called_once()
    postgres_activity.logger.error.assert_called()


@mark.asyncio
async def test_close(postgres_activity):
    postgres_activity.close()

    postgres_activity.engine.dispose.assert_called_once()


@mark.asyncio
async def test_del(postgres_activity):
    postgres_activity.close = MagicMock()
    postgres_activity.__del__()

    postgres_activity.close.assert_called_once()