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

from pytest import fixture, mark

from model_manager.activities.activities import Activities
from model_manager.workflows.predictions_batch import PredictionsBatch


@fixture
def predictions_batch() -> PredictionsBatch:
    return PredictionsBatch()


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


@mark.asyncio
@patch('model_manager.workflows.predictions_batch.workflow', new_callable=AsyncMock)
async def test_run(workflow_mock: AsyncMock, predictions_batch: PredictionsBatch):
    workflow_mock.execute_local_activity_method.return_value = {'data': 'test_data'}
    input_data = {
        'schedule_name': 'test_schedule',
        'model_name': 'test_model',
        'model_id': 'test_model_id',
        'query': 'SELECT * FROM test',
        'schema': 'test_schema',
        'table_name': 'test_table',
        'datetime_columns': ['timestamp', 'created_at'],
        'prediction_store_policy': 'erl:1',
        'model_config': {'retention': '30'},
    }

    await predictions_batch.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,
            )
        ]
    )
    prediction_input = {
        'metadata': metadata,
        'data': {'data': 'test_data'},
        'schema': input_data['schema'],
        'table_name': input_data['table_name'],
        'model_id': input_data['model_id'],
        'model_name': input_data['model_name'],
        'input_filters': input_data.get('input_filters', {'EMPTY_DATA': {'POLICY': 'STOP'}}),
        'mlflow_transform_filters': input_data.get(
            'mlflow_transform_filters', {'API_ERROR': {'POLICY': 'STOP'}}
        ),
        'mlflow_predict_filters': input_data.get(
            'mlflow_predict_filters', {'API_ERROR': {'POLICY': 'STOP'}}
        ),
        'model_config': input_data.get('model_config', {}),
        'path_priority': input_data.get('path_priority', ['STOP', 'CONTINUE', 'REPEAT']),
        'prediction_store_policy': input_data.get('prediction_store_policy', 'erl:1'),
    }

    workflow_mock.execute_child_workflow.assert_has_calls(
        [call('prediction_process', prediction_input)]
    )