import datetime
from unittest.mock import ANY, AsyncMock, MagicMock, patch

from pytest import fixture, mark, raises
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.activities.postgres import Postgres

from laborious.activities.storage import Storage

metadata = {
    'metadata': {
        'model_id': 'test_model_id',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
        'schedule_name': 'test_schedule',
    }
}


@fixture
@patch('laborious.activities.storage.MinioRepository')
def storage(mock_minio_repository):
    return Storage(
        host='localhost',
        port=5432,
        user='postgres',
        password='postgres',
        dbname='postgres',
        min_connections=1,
        max_connections=10,
        minio_config={
            'endpoint_url': 'localhost:9000',
            'access_key': 'minio',
            'secret_key': 'minio123',
            'region_name': 'us-east-1',
            'default_bucket': 'test',
        },
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
    )


@patch('laborious.activities.storage.MinioRepository')
def test___init___not_hasattr(mock_minio_repository):
    logger = MagicMock()
    notification_handler = MagicMock()
    metrics_controller = AsyncMock()
    storage = Storage(
        host='localhost',
        port=5432,
        user='postgres',
        password='postgres',
        dbname='postgres',
        min_connections=1,
        max_connections=10,
        minio_config={
            'endpoint_url': 'localhost:9000',
            'access_key': 'minio',
            'secret_key': 'minio123',
            'region_name': 'us-east-1',
            'default_bucket': 'test',
        },
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
    )
    assert isinstance(storage, Postgres)

    mock_minio_repository.assert_called_once_with(
        logger=logger,
        notification_handler=notification_handler,
        minio_endpoint_url='localhost:9000',
        minio_access_key='minio',
        minio_secret_key='minio123',
        minio_region_name='us-east-1',
        minio_default_bucket='test',
        metrics_controller=metrics_controller,
    )


@patch('laborious.activities.storage.MinioRepository')
def test___init___none_minio_repository(mock_minio_repository, storage):
    storage.minio_repository = None
    logger = MagicMock()
    notification_handler = MagicMock()
    metrics_controller = AsyncMock()
    storage.__init__(
        host='localhost',
        port=5432,
        user='postgres',
        password='postgres',
        dbname='postgres',
        min_connections=1,
        max_connections=10,
        minio_config={
            'endpoint_url': 'localhost:9000',
            'access_key': 'minio',
            'secret_key': 'minio123',
            'region_name': 'us-east-1',
            'default_bucket': 'test',
        },
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
    )

    mock_minio_repository.assert_called_once_with(
        logger=logger,
        notification_handler=notification_handler,
        minio_endpoint_url='localhost:9000',
        minio_access_key='minio',
        minio_secret_key='minio123',
        minio_region_name='us-east-1',
        minio_default_bucket='test',
        metrics_controller=metrics_controller,
    )


@patch('laborious.activities.storage.MinioRepository')
def test___init___done_repository(mock_minio_repository, storage):
    storage.__init__(
        host='localhost',
        port=5432,
        user='postgres',
        password='postgres',
        dbname='postgres',
        min_connections=1,
        max_connections=10,
        minio_config={
            'endpoint_url': 'localhost:9000',
            'access_key': 'minio',
            'secret_key': 'minio123',
            'region_name': 'us-east-1',
            'default_bucket': 'test',
        },
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
    )
    mock_minio_repository.assert_not_called()
    assert storage.minio_repository is not None


@mark.asyncio
async def test_query_to_minio_minio_repository_not_initialized(storage):
    storage.minio_repository = None

    with raises(ValueError) as e:
        await storage.query_to_minio({})

        assert str(e.value) == 'Minio repository not initialized'


@mark.asyncio
async def test_query_to_minio_not_data(storage):
    storage.load_custom_query = AsyncMock(return_value=None)
    result = await storage.query_to_minio({})

    storage.load_custom_query.assert_called_once_with({})
    assert result['success'] is False
    assert result['message'] == 'No data returned from query'


@mark.asyncio
@patch('laborious.activities.storage.pd.DataFrame')
@patch('laborious.activities.storage.now')
async def test_query_to_minio_success(now, dataframe, storage):
    data = [{'a': 1}, {'a': 2}, {'a': 3}]
    storage.load_custom_query = AsyncMock(return_value=data)
    now.return_value = datetime.datetime(2024, 1, 1, 0, 0, 0)
    storage.minio_repository.store_dataframe_as_parquet = AsyncMock()
    storage.minio_repository.minio_bucket = 'test'

    result = await storage.query_to_minio({'object_prefix': 'test', **metadata})

    dataframe.assert_called_once_with(data)

    storage.minio_repository.store_dataframe_as_parquet.assert_called_once_with(
        dataframe=dataframe.return_value,
        uri='s3://test/test_2024-01-01_00-00-00.parquet',
        object_name='test_2024-01-01_00-00-00.parquet',
        metadata=metadata['metadata'],
    )

    assert result['success'] is True
    assert result['object_key'] == 'test_2024-01-01_00-00-00.parquet'
    assert result['uri'] == 's3://test/test_2024-01-01_00-00-00.parquet'


@mark.asyncio
async def test_query_to_minio_error(storage):
    storage.send_notification = MagicMock()
    storage.send_notification_async = AsyncMock()
    storage.minio_repository.store_dataframe_as_parquet = AsyncMock()

    storage.load_custom_query = AsyncMock(side_effect=Exception('test'))
    result = await storage.query_to_minio({**metadata, 'object_prefix': 'test'})
    assert result['success'] is False
    assert result['message'] == 'test'
    storage.send_notification_async.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='ERROR_STORING_QUERY_TO_MINIO',
        message='Error storing query to MinIO: test',
        block='query_to_minio',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


def test_close(storage):
    storage.minio_repository = MagicMock()

    storage.close()

    assert storage.minio_repository is None


def test___del__(storage):
    storage.close = MagicMock()

    storage.__del__()

    storage.close.assert_called_once()