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

import numpy as np
import pandas as pd
from pytest import fixture, raises
from sientia_do.notifications.models import NotificationLevel
from sientia_do.observability.metrics_controller import MetricsController
from sientia_do.temporal.constants import DATETIME_FORMAT_WITH_TZ

from laborious.activities.mlflow import MLFlow


@fixture(autouse=True)
def _passthrough_from_dict():
    with patch(
        'laborious.activities.mlflow.MinioDataFramePayload.from_dict', side_effect=lambda x: x
    ):
        yield


@patch('laborious.activities.mlflow.MinioRepository')
def test___init__(mock_minio_repository):
    logger = MagicMock()
    notification_handler = MagicMock()
    metrics_controller = MagicMock(spec=MetricsController)
    mlflow_repo = MagicMock()
    plugin_store = MagicMock()

    minio_repo = mock_minio_repository(
        endpoint='localhost:9000',
        access_key='minio',
        secret_key='minio123',
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
        bucket='test',
    )
    mlflow = MLFlow(
        mlflow_repository=mlflow_repo,
        plugin_store=plugin_store,
        minio_repository=minio_repo,
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
    )

    assert mlflow.mlflow_repository is mlflow_repo
    assert mlflow.plugin_store is plugin_store

    mock_minio_repository.assert_called_once_with(
        endpoint='localhost:9000',
        access_key='minio',
        secret_key='minio123',
        logger=ANY,
        notification_handler=ANY,
        metrics_controller=ANY,
        bucket='test',
    )


@fixture
@patch('laborious.activities.mlflow.MinioRepository')
def mlflow(mock_minio_repository):
    logger = MagicMock()
    notification_handler = MagicMock()
    metrics_controller = MagicMock(spec=MetricsController)
    mlflow_repo = MagicMock()
    plugin_store = MagicMock()

    minio_repo = mock_minio_repository(
        endpoint='localhost:9000',
        access_key='minio',
        secret_key='minio123',
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
        bucket='test',
    )
    mlflow = MLFlow(
        mlflow_repository=mlflow_repo,
        plugin_store=plugin_store,
        minio_repository=minio_repo,
        logger=logger,
        notification_handler=notification_handler,
        metrics_controller=metrics_controller,
    )

    mlflow.minio_repository = MagicMock()

    mlflow.send_notification = MagicMock()
    mlflow.emit_metric = MagicMock()
    mlflow.error = MagicMock()
    mlflow.debug = MagicMock()
    mlflow.info = MagicMock()
    mlflow.warning = MagicMock()
    mlflow.critical = MagicMock()

    return mlflow


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


def test_detect_and_parse_datetime_index_empty(mlflow):
    df = pd.DataFrame()
    out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
    assert out.empty


def test_detect_and_parse_datetime_index_mixed_types_error(mlflow):
    idx = pd.Index([pd.Timestamp('2020-01-01', tz='UTC'), 'x'])
    df = pd.DataFrame({'a': [1, 2]}, index=idx)
    with raises(ValueError):
        mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])


def test_detect_and_parse_datetime_index_invalid_string_error(mlflow):
    idx = pd.Index(['bad-format'])
    df = pd.DataFrame({'a': [1]}, index=idx)
    with raises(ValueError):
        mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])


def test_detect_and_parse_datetime_index_unsupported_type_error(mlflow):
    idx = pd.Index([pd.Period('2020-01', freq='M')])
    df = pd.DataFrame({'a': [1]}, index=idx)
    with raises(ValueError):
        mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])


def test_detect_and_parse_datetime_index_datetime_success(mlflow):
    idx = pd.Index([datetime(2020, 1, 1, 0, 0, 0)])
    df = pd.DataFrame({'a': [1]}, index=idx)
    out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
    assert out.index[0].endswith('+0000')


def test_detect_and_parse_datetime_index_timestamp_with_tz_success(mlflow):
    idx = pd.DatetimeIndex([pd.Timestamp('2020-01-01 00:00:00', tz='UTC')])
    df = pd.DataFrame({'a': [1]}, index=idx)
    out = mlflow._detect_and_parse_datetime_index(df, metadata['metadata'])
    assert out.index[0].endswith('+0000')


@patch(
    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
    new_callable=MagicMock,
)
def test_request_transform_success(mock_from_dataframe, mlflow):
    ts = pd.Timestamp('2020-01-01', tz='UTC')
    raw = pd.DataFrame(
        {
            'variable': ['v1', 'v1'],
            'timestamp': [ts, ts],
            'value': [1.0, 2.0],
            'created_at': [ts, ts],
        }
    )
    pivoted = raw.sort_values('created_at', ascending=False).drop_duplicates(
        subset=['variable', 'timestamp'], keep='first'
    )
    pivoted = pivoted.pivot(index='timestamp', columns='variable', values='value')
    pivoted = pivoted.fillna(np.nan)
    pivoted.columns.name = None
    pivoted.index.name = None
    pivoted['timestamp'] = pivoted.index

    out_idx = pd.Index([ts.strftime(DATETIME_FORMAT_WITH_TZ)], name=None)
    out_df = pd.DataFrame({'v1': [1.0]}, index=out_idx)
    wrapper = MagicMock()
    wrapper.transform.return_value = (out_df, {'meta': True})
    mlflow.mlflow_repository.get_cached_model.return_value = wrapper

    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=raw)

    input_data = {
        **metadata,
        'data': payload,
        'model_name': 'test_model',
        'model_config': {},
    }

    response_data = mlflow.request_transform(input_data)

    mlflow.mlflow_repository.get_cached_model.assert_called_once_with(
        model_name='test_model',
        alias='production',
        retention_minutes=0,
        metadata=metadata['metadata'],
    )
    mock_from_dataframe.assert_called_once()
    assert response_data == mock_from_dataframe.return_value


@patch(
    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
    new_callable=MagicMock,
)
def test_request_transform_success_without_transform_meta(mock_from_dataframe, mlflow):
    ts = pd.Timestamp('2020-01-01', tz='UTC')
    raw = pd.DataFrame(
        {
            'variable': ['v1'],
            'timestamp': [ts],
            'value': [1.0],
            'created_at': [ts],
        }
    )
    out_idx = pd.Index([ts.strftime(DATETIME_FORMAT_WITH_TZ)], name=None)
    out_df = pd.DataFrame({'v1': [1.0]}, index=out_idx)
    wrapper = MagicMock()
    wrapper.transform.return_value = (out_df, {})
    mlflow.mlflow_repository.get_cached_model.return_value = wrapper

    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=raw)

    input_data = {
        **metadata,
        'data': payload,
        'model_name': 'test_model',
        'model_config': {},
    }

    mlflow.request_transform(input_data)
    mock_from_dataframe.assert_called_once()


@patch(
    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
    new_callable=MagicMock,
)
def test_request_transform_failure(mock_from_dataframe, mlflow):
    mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('boom')

    data_mock = MagicMock()
    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=data_mock)

    input_data = {
        **metadata,
        'data': payload,
        'model_name': 'test_model',
        'model_config': {},
    }

    data_mock.sort_values.return_value = data_mock
    data_mock.drop_duplicates.return_value = data_mock
    data_mock.pivot.return_value = data_mock

    mlflow.request_transform(input_data)

    mock_from_dataframe.assert_called_once_with(
        dataframe=None,
        minio_repo=mlflow.minio_repository,
        model_name='test_model',
        operation='transform',
        status={'success': False, 'content': ANY},
        workflow_metadata=metadata['metadata'],
        last_timestamp=payload.last_timestamp,
        logger=mlflow.logger,
    )


@patch(
    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
    new_callable=MagicMock,
)
@patch('laborious.activities.mlflow.to_datetime')
def test_request_predict(mock_to_datetime, mock_from_dataframe, mlflow):
    wrapper = MagicMock()
    pred_df = MagicMock()
    wrapper.predict.return_value = (pred_df, {})
    mlflow.mlflow_repository.get_cached_model.return_value = wrapper

    data_mock = MagicMock()
    data_mock.index = pd.DatetimeIndex([pd.Timestamp('2020-01-01', tz='UTC')])
    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=data_mock)

    input_data = {
        **metadata,
        'data': payload,
        'model_name': 'test_model',
        'model_config': {},
    }

    pred_df.columns = MagicMock()
    pred_df.__setitem__ = MagicMock()

    response_data = mlflow.request_predict(input_data)

    data_mock.replace.assert_called_once_with(np.nan, None, inplace=True)
    mock_to_datetime.assert_called()
    mlflow.mlflow_repository.get_cached_model.assert_called_once()
    mock_from_dataframe.assert_called_once()
    assert response_data == mock_from_dataframe.return_value


@patch(
    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
    new_callable=MagicMock,
)
@patch('laborious.activities.mlflow.to_datetime')
def test_request_predict_success_dataframe_and_meta(mock_to_datetime, mock_from_dataframe, mlflow):
    mock_to_datetime.side_effect = lambda x, **kwargs: x
    wrapper = MagicMock()
    pred_df = pd.DataFrame({'raw': [0.3]})
    wrapper.predict.return_value = (pred_df, {'m': 1})
    mlflow.mlflow_repository.get_cached_model.return_value = wrapper

    data_mock = MagicMock()
    data_mock.index = pd.DatetimeIndex([pd.Timestamp('2020-01-01', tz='UTC')])
    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=data_mock)

    input_data = {
        **metadata,
        'data': payload,
        'model_name': 'test_model',
        'model_config': {},
    }

    mlflow.request_predict(input_data)

    assert list(pred_df.columns) == ['prediction', 'response_time']
    mlflow.info.assert_any_call("Wrapper predict metadata: {'m': 1}", metadata['metadata'])
    mock_from_dataframe.assert_called_once()


@patch(
    'laborious.activities.mlflow.MinioDataFramePayload.from_dataframe',
    new_callable=MagicMock,
)
@patch('laborious.activities.mlflow.to_datetime')
def test_request_predict_failure(mock_to_datetime, mock_from_dataframe, mlflow):
    mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('predict boom')

    data_mock = MagicMock()
    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=data_mock)

    input_data = {
        **metadata,
        'data': payload,
        'model_name': 'test_model',
        'model_config': {},
    }

    mlflow.request_predict(input_data)

    mock_from_dataframe.assert_called_once_with(
        dataframe=None,
        minio_repo=mlflow.minio_repository,
        model_name='test_model',
        operation='predict',
        status={'success': False, 'content': ANY},
        workflow_metadata=metadata['metadata'],
        last_timestamp=payload.last_timestamp,
        logger=mlflow.logger,
    )


@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_success_data_success_retrain(
    mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow, tmp_path
):
    mock_mkdtemp.return_value = str(tmp_path)
    mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)

    mv_alias = MagicMock()
    mv_alias.run_id = 'source-run'
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias

    wrapper = MagicMock()
    mlflow.mlflow_repository.get_cached_model.return_value = wrapper

    mock_cm = MagicMock()
    mock_cm.__enter__.return_value = MagicMock(run_id='new-run', experiment_id='exp-1')
    mock_cm.__exit__.return_value = False
    mlflow.mlflow_repository.start_run.return_value = mock_cm

    ts = pd.Timestamp('2020-01-01', tz='UTC')
    raw_data = pd.DataFrame(
        {
            'variable': ['target', 'f1'],
            'timestamp': [ts, ts],
            'value': [1.0, 2.0],
        }
    )
    pivoted_index = pd.DatetimeIndex([ts])
    wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0])

    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=raw_data)

    response = mlflow.retrain_model(
        {
            **metadata,
            'data': payload,
            'model_name': 'test_model',
            'model_config': {
                'target': 'target',
            },
        }
    )

    wrapper.retrain.assert_called_once()
    wrapper.store_model.assert_called_once_with(name='test_model')
    assert response['success'] is True
    assert response['experiment']['run_id'] == 'new-run'


@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_success_with_payload_data(
    mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow, tmp_path
):
    mock_mkdtemp.return_value = str(tmp_path)
    mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
    mv_alias = MagicMock(run_id='src')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
    wrapper = MagicMock()
    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
    mock_cm = MagicMock()
    mock_cm.__enter__.return_value = MagicMock(run_id='r', experiment_id='e')
    mock_cm.__exit__.return_value = False
    mlflow.mlflow_repository.start_run.return_value = mock_cm

    ts = pd.Timestamp('2020-01-01', tz='UTC')
    raw_data = pd.DataFrame(
        {
            'variable': ['target', 'f1', 'target', 'f1'],
            'timestamp': [ts, ts, ts + pd.Timedelta(hours=1), ts + pd.Timedelta(hours=1)],
            'value': [1.0, 2.0, 3.0, 4.0],
            'created_at': [ts, ts, ts + pd.Timedelta(hours=1), ts + pd.Timedelta(hours=1)],
        }
    )
    pivoted_index = pd.DatetimeIndex([ts, ts + pd.Timedelta(hours=1)])
    wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0, 3.0])

    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=raw_data)

    response = mlflow.retrain_model(
        {
            **metadata,
            'data': payload,
            'model_name': 'test_model',
            'model_config': {'target': 'target'},
        }
    )

    assert response['success'] is True


@patch('laborious.activities.mlflow.mlflow.log_artifact')
@patch('laborious.activities.mlflow.tempfile.mkdtemp')
@patch('laborious.activities.mlflow.rmtree')
@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_always_uses_retrain_even_with_full_retrain_flag(
    mock_to_datetime, mock_rmtree, mock_mkdtemp, mock_log_artifact, mlflow, tmp_path
):
    mock_mkdtemp.return_value = str(tmp_path)
    mock_to_datetime.side_effect = lambda idx, **kwargs: pd.DatetimeIndex(idx)
    mv_alias = MagicMock(run_id='src')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
    wrapper = MagicMock()
    mlflow.mlflow_repository.get_cached_model.return_value = wrapper
    mock_cm = MagicMock()
    mock_cm.__enter__.return_value = MagicMock(run_id='r', experiment_id='e')
    mock_cm.__exit__.return_value = False
    mlflow.mlflow_repository.start_run.return_value = mock_cm

    ts = pd.Timestamp('2020-01-01', tz='UTC')
    ts_next = ts + pd.Timedelta(days=1)
    raw_data = pd.DataFrame(
        {
            'variable': ['target', 'f1', 'target', 'f1'],
            'timestamp': [ts, ts, ts_next, ts_next],
            'value': [1.0, 2.0, 3.0, 4.0],
        }
    )
    pivoted_index = pd.DatetimeIndex([ts, ts_next])
    wrapper.retrain.return_value = _retrain_prediction_frame(pivoted_index, 'target', [1.0, 3.0])

    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=raw_data)

    response = mlflow.retrain_model(
        {
            **metadata,
            'data': payload,
            'model_name': 'test_model',
            'model_config': {'target': 'target', 'full_retrain': True, 'validation_fraction': 0.5},
        }
    )

    wrapper.retrain.assert_called_once()
    wrapper.train.assert_not_called()
    assert response['success'] is True


@patch('laborious.activities.mlflow.to_datetime')
def test_retrain_model_success_data_fail_retrain(mock_to_datetime, mlflow):
    mv_alias = MagicMock(run_id='src')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv_alias
    mlflow.mlflow_repository.get_cached_model.side_effect = RuntimeError('retrain failed')

    raw_data = MagicMock(columns=['variable', 'timestamp', 'value', 'created_at'])
    raw_data.__getitem__.return_value.max.return_value = 'tsmax'
    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=raw_data)

    pivoted = MagicMock()
    raw_data.sort_values.return_value = raw_data
    raw_data.drop_duplicates.return_value = raw_data
    raw_data.pivot.return_value = pivoted
    pivoted.fillna = MagicMock()
    pivoted.columns.name = None
    pivoted.index = MagicMock()
    pivoted.__setitem__ = MagicMock()

    response = mlflow.retrain_model(
        {
            **metadata,
            'data': payload,
            'model_name': 'test_model',
            'model_config': {'target': 'target'},
        }
    )

    assert response['success'] is False
    assert 'retrain failed' in response['message']


def test_retrain_model_data_error(mlflow):
    response = mlflow.retrain_model(
        {
            **metadata,
            'model_name': 'test_model',
            'model_config': {
                'target': 'target',
            },
        }
    )

    assert response['success'] is False
    assert 'data' in response['message'].lower() or 'loading' in response['message'].lower()


def test_retrain_model_missing_target(mlflow):
    ts = pd.Timestamp('2020-01-01', tz='UTC')
    raw_data = pd.DataFrame(
        {
            'variable': ['f1', 'f2'],
            'timestamp': [ts, ts],
            'value': [1.0, 2.0],
        }
    )
    payload = MagicMock()
    payload.retrieve = MagicMock(return_value=raw_data)

    response = mlflow.retrain_model(
        {
            **metadata,
            'data': payload,
            'model_name': 'test_model',
            'model_config': {},
        }
    )

    assert response['success'] is False
    assert 'target' in response['message']


def test_retrain_model_data_error_no_minio_repository(mlflow):
    mlflow.minio_repository = None

    with raises(ValueError) as e:
        mlflow.retrain_model(
            {
                **metadata,
                'object_key': 'test_object_key',
                'model_name': 'test_model',
                'model_config': {
                    'target': 'target',
                },
            }
        )

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


def test_update_production_model(mlflow):
    mlflow.mlflow_repository._client.search_model_versions.return_value = [
        MagicMock(version='3', run_id='run-x'),
        MagicMock(version='2', run_id='run-x'),
    ]

    input_data = {
        **metadata,
        'model_name': 'test_model',
        'model_id': 1,
        'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
        'timestamp': 2,
        'status': 'success',
    }

    response = mlflow.update_production_model(input_data)

    mlflow.mlflow_repository.promote_to_alias.assert_called_once_with(
        model_name='test_model',
        version='3',
        alias='production',
        metadata=metadata['metadata'],
    )

    assert response['model_name'] == 'test_model'
    assert response['version'] == '3'


def test_update_production_model_error(mlflow):
    mlflow.mlflow_repository._client.search_model_versions.return_value = []

    input_data = {
        **metadata,
        'model_name': 'test_model',
        'model_id': 1,
        'experiment': {'run_id': 'run-x', 'experiment_id': 'e1'},
        'timestamp': 2,
        'status': 'success',
    }

    try:
        mlflow.update_production_model(input_data)
    except Exception:
        mlflow.send_notification.assert_called_once_with(
            metadata=metadata['metadata'],
            notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
            message=ANY,
            block='update_production_model',
            level=NotificationLevel.ERROR,
            attachment_content=ANY,
        )
    else:
        raise AssertionError('Expected exception')


def _artifact_file_info(path: str) -> MagicMock:
    file_info = MagicMock()
    file_info.path = path
    return file_info


def _retrain_prediction_frame(
    index: pd.DatetimeIndex, target: str, values: list[float]
) -> pd.DataFrame:
    return pd.DataFrame({target: values}, index=index)


@patch('laborious.activities.mlflow.to_datetime')
def test_get_reference_data_success(mock_to_datetime, mlflow):
    input_data = {
        **metadata,
        'model_name': 'test_model',
    }

    mv = MagicMock(run_id='run1')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
    mlflow.mlflow_repository._client.list_artifacts.return_value = [
        _artifact_file_info('evaluation_data.csv'),
    ]

    mock_reference_data = MagicMock()
    mock_reference_data.__getitem__.return_value = MagicMock()
    mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
    mock_reference_data.to_dict.return_value = [
        {'timestamp': '2023-05-26 11:12:27', 'value': 1.0},
    ]

    with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
        with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
            with patch('laborious.activities.mlflow.rmtree'):
                with patch.object(
                    mlflow,
                    '_find_downloaded_csv',
                    return_value=MagicMock(),
                ):
                    result = mlflow.get_reference_data(input_data)

    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
        run_id='run1',
        artifact_path='evaluation_data.csv',
        dst_path='/t',
        metadata=metadata['metadata'],
    )
    mock_reference_data.to_dict.assert_called_once_with(orient='records')
    assert result == mock_reference_data.to_dict.return_value


def test_get_reference_data_not_found(mlflow):
    input_data = {
        **metadata,
        'model_name': 'test_model',
    }

    mlflow.mlflow_repository._client.get_model_version_by_alias.side_effect = Exception('missing')

    result = mlflow.get_reference_data(input_data)

    mlflow.warning.assert_called()
    assert result is None


def test_get_reference_data_only_test_data_csv(mlflow):
    input_data = {
        **metadata,
        'model_name': 'test_model',
    }
    mv = MagicMock(run_id='run1')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
    mlflow.mlflow_repository._client.list_artifacts.return_value = [
        _artifact_file_info('test_data.csv'),
    ]

    mock_reference_data = MagicMock()
    mock_reference_data.__getitem__.return_value = MagicMock()
    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]

    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
                with patch('laborious.activities.mlflow.rmtree'):
                    with patch.object(
                        mlflow,
                        '_find_downloaded_csv',
                        return_value=MagicMock(),
                    ):
                        result = mlflow.get_reference_data(input_data)

    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
        run_id='run1',
        artifact_path='test_data.csv',
        dst_path='/t',
        metadata=metadata['metadata'],
    )
    assert result == mock_reference_data.to_dict.return_value


def test_get_reference_data_prefers_evaluation_data_when_both_listed(mlflow):
    input_data = {
        **metadata,
        'model_name': 'test_model',
    }
    mv = MagicMock(run_id='run1')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
    mlflow.mlflow_repository._client.list_artifacts.return_value = [
        _artifact_file_info('test_data.csv'),
        _artifact_file_info('evaluation_data.csv'),
    ]

    mock_reference_data = MagicMock()
    mock_reference_data.__getitem__.return_value = MagicMock()
    mock_reference_data.to_dict.return_value = [{'timestamp': '2023-05-26 11:12:27', 'value': 1.0}]

    with patch('laborious.activities.mlflow.to_datetime') as mock_to_datetime:
        mock_to_datetime.return_value.dt.strftime.return_value = MagicMock()
        with patch('laborious.activities.mlflow.pd.read_csv', return_value=mock_reference_data):
            with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='/t'):
                with patch('laborious.activities.mlflow.rmtree'):
                    with patch.object(
                        mlflow,
                        '_find_downloaded_csv',
                        return_value=MagicMock(),
                    ):
                        result = mlflow.get_reference_data(input_data)

    mlflow.mlflow_repository.download_artifacts.assert_called_once_with(
        run_id='run1',
        artifact_path='evaluation_data.csv',
        dst_path='/t',
        metadata=metadata['metadata'],
    )
    assert result == mock_reference_data.to_dict.return_value


def test_get_reference_data_no_candidate_artifacts_returns_none(mlflow):
    input_data = {
        **metadata,
        'model_name': 'test_model',
    }
    mv = MagicMock(run_id='run1')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
    mlflow.mlflow_repository._client.list_artifacts.return_value = [
        _artifact_file_info('other_artifact.csv'),
    ]

    result = mlflow.get_reference_data(input_data)

    mlflow.mlflow_repository.download_artifacts.assert_not_called()
    mlflow.warning.assert_called()
    assert result is None


def test_get_reference_data_missing_csv_file_returns_none(mlflow):
    input_data = {
        **metadata,
        'model_name': 'test_model',
    }
    mv = MagicMock(run_id='run1')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
    mlflow.mlflow_repository._client.list_artifacts.return_value = [
        _artifact_file_info('evaluation_data.csv'),
    ]

    with patch('laborious.activities.mlflow.tempfile.mkdtemp', return_value='tmp'):
        with patch('laborious.activities.mlflow.rmtree'):
            with patch.object(mlflow, '_find_downloaded_csv', return_value=None):
                result = mlflow.get_reference_data(input_data)

    mlflow.mlflow_repository.download_artifacts.assert_called_once()
    mlflow.warning.assert_called()
    assert result is None


def test_get_reference_data_exception(mlflow):
    input_data = {
        **metadata,
        'model_name': 'test_model',
    }

    mv = MagicMock(run_id='run1')
    mlflow.mlflow_repository._client.get_model_version_by_alias.return_value = mv
    mlflow.mlflow_repository._client.list_artifacts.return_value = [
        _artifact_file_info('evaluation_data.csv'),
    ]
    mlflow.mlflow_repository.download_artifacts.side_effect = Exception('dl fail')

    result = mlflow.get_reference_data(input_data)

    assert result is None


def test_find_downloaded_csv_prefers_the_direct_path(mlflow, tmp_path):
    direct = tmp_path / 'evaluation_data.csv'
    direct.write_text('a\n1\n')

    assert mlflow._find_downloaded_csv(str(tmp_path), 'evaluation_data.csv') == direct


def test_find_downloaded_csv_searches_nested_directories(mlflow, tmp_path):
    nested = tmp_path / 'artifacts' / 'data' / 'evaluation_data.csv'
    nested.parent.mkdir(parents=True)
    nested.write_text('a\n1\n')

    assert mlflow._find_downloaded_csv(str(tmp_path), 'evaluation_data.csv') == nested


def test_find_downloaded_csv_returns_none_when_missing(mlflow, tmp_path):
    assert mlflow._find_downloaded_csv(str(tmp_path), 'evaluation_data.csv') is None