from unittest.mock import ANY, AsyncMock, call, patch

from pytest import fixture, mark

from model_manager.activities.activities import Activities
from model_manager.workflows.minimal_retrain import MinimalRetrain


@fixture
def minimal_retrain() -> MinimalRetrain:
    return MinimalRetrain()


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


@mark.asyncio
@patch('model_manager.workflows.minimal_retrain.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, minimal_retrain: MinimalRetrain):
    input_data = {
        'model_id': 'test_model_id',
        'model_name': 'test_model',
        'workflow_name': 'minimal_retrain',
        'schedule_name': 'test_schedule',
        'query': 'test_query',
        'schema': 'test_schema',
        'table_name': 'test_table',
    }

    workflow_mock.execute_activity_method = AsyncMock(
        return_value={
            'data1': '1',
            'data2': '2',
        }
    )

    await minimal_retrain.run(input_data)

    workflow_mock.execute_local_activity_method.assert_has_calls(
        [
            call(
                Activities.load_custom_query,
                {
                    **metadata,
                    'query': input_data['query'],
                    'datetime_columns': input_data.get('datetime_columns', []),
                },
                retry_policy=ANY,
                start_to_close_timeout=ANY,
            )
        ]
    )

    workflow_mock.execute_activity_method.assert_has_calls(
        [
            call(
                Activities.retrain_model,
                {
                    **metadata,
                    'data': workflow_mock.execute_local_activity_method.return_value,
                    'model_name': input_data['model_name'],
                },
                retry_policy=ANY,
                start_to_close_timeout=ANY,
            )
        ]
    )

    workflow_mock.execute_activity_method.assert_has_calls(
        [
            call(
                Activities.update_production_model,
                {
                    **metadata,
                    'model_name': input_data['model_name'],
                    'model_id': input_data['model_id'],
                    **workflow_mock.execute_activity_method.return_value,
                },
                retry_policy=ANY,
                start_to_close_timeout=ANY,
            )
        ]
    )

    workflow_mock.execute_activity_method.assert_has_calls(
        [
            call(
                Activities.export_data_to_postgres,
                {
                    **metadata,
                    'data': workflow_mock.execute_activity_method.return_value,
                    'schema': input_data['schema'],
                    'table_name': input_data['table_name'],
                },
                retry_policy=ANY,
                start_to_close_timeout=ANY,
            )
        ]
    )