from unittest.mock import ANY, MagicMock, patch

import numpy as np
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel
from sientia_do.temporal.constants import DATETIME_FORMAT, DATETIME_FORMAT_WITH_TZ

from model_manager.activities.mlflow import MLFlow


@patch('model_manager.activities.mlflow.MLFlowRepository')
def test___init__(mock_mlflow_repository):
    mlflow = MLFlow(
        mlflow_host='http://localhost',
        mlflow_port=5000,
        mlflow_username='admin',
        mlflow_password='admin',
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    assert mlflow.mlflow_host == 'http://localhost'
    assert mlflow.mlflow_port == 5000
    assert mlflow.mlflow_username == 'admin'
    assert mlflow.mlflow_password == 'admin'

    mock_mlflow_repository.assert_called_once_with('http://localhost:5000', 'admin', 'admin', ANY)


@fixture
@patch('model_manager.activities.mlflow.MLFlowRepository')
def mlflow(mock_mlflow_repository):
    mlflow = MLFlow(
        mlflow_host='http://localhost:5000',
        mlflow_port=5000,
        mlflow_username='admin',
        mlflow_password='admin',
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )

    mlflow.send_notification = MagicMock()

    return mlflow


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


@mark.asyncio
@patch('model_manager.activities.mlflow.DataFrame')
@patch('model_manager.activities.mlflow.max')
async def test_request_transform_success(mock_max, mock_dataframe, mlflow):
    mock_max.return_value = '2024-01-02'
    # Mock input data
    input_data = {
        **metadata,
        'data': [
            {
                'timestamp': '2024-01-01',
                'variable': 'var1',
                'value': 1.0,
                'created_at': '2024-01-01 12:00:00',
            },
            {
                'timestamp': '2024-01-01',
                'variable': 'var2',
                'value': 2.0,
                'created_at': '2024-01-01 12:00:00',
            },
            {
                'timestamp': '2024-01-02',
                'variable': 'var1',
                'value': 3.0,
                'created_at': '2024-01-02 12:00:00',
            },
            {
                'timestamp': '2024-01-02',
                'variable': 'var2',
                'value': 4.0,
                'created_at': '2024-01-02 12:00:00',
            },
            {
                'timestamp': '2024-01-02',
                'variable': 'var1',
                'value': 1.0,
                'created_at': '2024-01-01 12:00:00',
            },
            {
                'timestamp': '2024-01-02',
                'variable': 'var2',
                'value': 1.0,
                'created_at': '2024-01-01 12:00:00',
            },
        ],
        'model_name': 'test_model',
        'model_config': {},
    }

    # Mock the transform response
    expected_response = {'prediction': [0.5, 0.6], 'timestamp': ['2024-01-01', '2024-01-02']}
    mlflow.model_monitoring_repository.transform.return_value = expected_response

    mock_dataframe.return_value.sort_values.return_value = mock_dataframe.return_value
    mock_dataframe.return_value.drop_duplicates.return_value = mock_dataframe.return_value

    # Call the method
    response_data = await mlflow.request_transform(input_data)

    # Verify the data was correctly transformed
    mock_dataframe.assert_called_once_with(input_data['data'])
    mock_dataframe.return_value.pivot.assert_called_once_with(
        index='timestamp', columns='variable', values='value'
    )
    mock_dataframe = mock_dataframe.return_value.pivot.return_value
    mock_dataframe.fillna.assert_called_once_with(np.nan, inplace=True)
    # mock_dataframe.reset_index.assert_called_once()
    mock_dataframe.columns.name = None

    # Verify the response
    assert response_data == expected_response

    # Verify the repository was called with correct arguments
    mlflow.model_monitoring_repository.transform.assert_called_once_with(
        'test_model', mock_dataframe, {}, metadata['metadata']
    )


@mark.asyncio
@patch('model_manager.activities.mlflow.DataFrame')
@patch('model_manager.activities.mlflow.to_datetime')
@patch('model_manager.activities.mlflow.max')
async def test_request_predict(mock_max, mock_to_datetime, mock_dataframe, mlflow):
    mock_max.return_value = '2024-01-02'
    # Mock input data
    input_data = {
        **metadata,
        'data': {
            'variable': {
                '2024-01-01': 'var1',
                '2024-01-02': 'var2',
                '2024-01-03': 'var1',
                '2024-01-04': 'var2',
            },
            'value': {'2024-01-01': 1.0, '2024-01-02': 2.0, '2024-01-03': 3.0, '2024-01-04': 4.0},
        },
        'model_name': 'test_model',
        'model_config': {},
    }

    # Mock the predict response
    expected_response = {'prediction': [0.5, 0.6]}
    mlflow.model_monitoring_repository.predict.return_value = expected_response

    # Call the method
    response_data = await mlflow.request_predict(input_data)

    mock_dataframe.assert_called_once_with(input_data['data'])
    mock_dataframe.return_value.replace.assert_called_once_with(np.nan, None, inplace=True)
    mock_dataframe.return_value.__setitem__.assert_any_call(
        'timestamp', mock_to_datetime.return_value.dt.strftime.return_value
    )

    mock_to_datetime.assert_called_once_with(
        mock_dataframe.return_value.__getitem__.return_value, format=DATETIME_FORMAT_WITH_TZ
    )
    mock_to_datetime.return_value.dt.strftime.assert_called_once_with(DATETIME_FORMAT)

    # Verify the response
    assert response_data == expected_response

    # Verify the repository was called with correct arguments
    mlflow.model_monitoring_repository.predict.assert_called_once_with(
        'test_model', mock_dataframe.return_value, {}, metadata['metadata']
    )


@mark.asyncio
async def test_retrain_model(mlflow):
    data = {
        'model_id': [4, 5, 6, 7],
        'created_at': [1, 2, 3, 4],
        'timestamp': [1, 1, 2, 2],
        'variable': ['var1', 'var2', 'var1', 'var2'],
        'value': [1, 2, 3, 4],
    }

    mlflow.model_monitoring_repository.retrain_model.return_value = (
        'Model retrained successfully',
        'test',
    )

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

    mlflow.model_monitoring_repository.retrain_model.assert_called_once()

    assert response == {
        'status': 'Model retrained successfully',
        'timestamp': 2,
        'experiment': 'test',
    }


@mark.asyncio
async def test_retrain_model_error(mlflow):
    mlflow.model_monitoring_repository.retrain_model.side_effect = Exception(
        'Error retraining model'
    )

    data = {
        'model_id': [4, 5, 6, 7],
        'created_at': [1, 2, 3, 4],
        'timestamp': [1, 1, 2, 2],
        'variable': ['var1', 'var2', 'var1', 'var2'],
        'value': [1, 2, 3, 4],
    }

    try:
        await mlflow.retrain_model({**metadata, 'data': data, 'model_name': 'test_model'})
    except Exception as e:  # noqa: BLE001
        assert str(e) == 'Error retraining model'
        mlflow.send_notification.assert_called_once_with(
            metadata=metadata['metadata'],
            notification_id='RETRAIN_MODEL_ERROR',
            message='Error retraining model test_model: Error retraining model',
            block='retrain_model',
            level=NotificationLevel.ERROR,
            attachment_content=ANY,
        )
    else:
        raise AssertionError('No exception raised')


@mark.asyncio
async def test_update_production_model(mlflow):
    mlflow.model_monitoring_repository.update_production_model.return_value = {
        'data1': 1,
        'data2': 2,
    }

    input_data = {
        **metadata,
        'model_name': 'test_model',
        'model_id': 1,
        'experiment': 'test',
        'timestamp': 2,
        'status': 'success',
    }

    response = await mlflow.update_production_model(input_data)

    mlflow.model_monitoring_repository.update_production_model.assert_called_once_with(
        experiment='test', model_name='test_model'
    )

    assert response == {
        'data1': {0: 1},
        'data2': {0: 2},
        'model_id': {0: 1},
        'model_name': {0: 'test_model'},
        'timestamp': {0: 2},
        'status': {0: 'success'},
    }


@mark.asyncio
async def test_update_production_model_error(mlflow):
    mlflow.model_monitoring_repository.update_production_model.side_effect = Exception(
        'Error updating production model'
    )

    input_data = {
        **metadata,
        'model_name': 'test_model',
        'model_id': 1,
        'experiment': 'test',
        'timestamp': 2,
        'status': 'success',
    }

    try:
        await mlflow.update_production_model(input_data)
    except Exception as e:  # noqa: BLE001
        assert str(e) == 'Error updating production model'
        mlflow.send_notification.assert_called_once_with(
            metadata=metadata['metadata'],
            notification_id='UPDATE_PRODUCTION_MODEL_ERROR',
            message='Error updating production model test_model: Error updating production model',
            block='update_production_model',
            level=NotificationLevel.ERROR,
            attachment_content=ANY,
        )
    else:
        raise AssertionError('No exception raised')


@mark.asyncio
async def test_save_model_success(mlflow):
    """Test save_model successfully saves model and artifacts to MLflow."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    # Mock train result
    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params
    train_result.run_name = None  # Will be set by get_next_run_name

    # Mock repository methods
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
    mlflow.model_monitoring_repository.save_run.return_value = None

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify repository methods were called
    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('test_experiment')
    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once_with(train_result)
    mlflow.model_monitoring_repository.save_run.assert_called_once_with(train_result)

    # Verify response
    assert response['success'] is True
    assert response['result'] == train_result
    assert response['error_message'] is None
    assert train_result.run_name == 'test_experiment-1'


@mark.asyncio
async def test_save_model_get_next_run_name_error(mlflow):
    """Test save_model handles error during get_next_run_name."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock error in get_next_run_name
    mlflow.model_monitoring_repository.get_next_run_name.side_effect = Exception(
        'MLflow connection error'
    )

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify error handling
    assert response['success'] is False
    assert response['result'] is None
    assert 'MLflow connection error' in response['error_message']

    # Verify notification was sent
    mlflow.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='SAVE_MODEL_ERROR',
        message=ANY,
        block='save_model',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_save_model_generate_artifacts_error(mlflow):
    """Test save_model handles error during generate_artifacts."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock successful get_next_run_name but error in generate_artifacts
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.side_effect = FileNotFoundError(
        'Reports directory does not exist'
    )

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify error handling
    assert response['success'] is False
    assert response['result'] is None
    assert 'Reports directory does not exist' in response['error_message']

    # Verify notification was sent
    mlflow.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='SAVE_MODEL_ERROR',
        message=ANY,
        block='save_model',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_save_model_save_run_error(mlflow):
    """Test save_model handles error during save_run."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock successful get_next_run_name and generate_artifacts but error in save_run
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
    mlflow.model_monitoring_repository.save_run.side_effect = ValueError(
        'One or more metrics (MSE, R2, MAE) are None'
    )

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify error handling
    assert response['success'] is False
    assert response['result'] is None
    assert 'One or more metrics (MSE, R2, MAE) are None' in response['error_message']

    # Verify notification was sent
    mlflow.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='SAVE_MODEL_ERROR',
        message=ANY,
        block='save_model',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_save_model_missing_metadata(mlflow):
    """Test save_model handles missing metadata gracefully."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'test_experiment'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params

    # Mock repository methods
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'test_experiment-1'
    mlflow.model_monitoring_repository.generate_artifacts.return_value = train_result
    mlflow.model_monitoring_repository.save_run.return_value = None

    # Input data without metadata
    input_data = {
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify it still works (metadata defaults to {})
    assert response['success'] is True
    assert response['result'] == train_result
    assert response['error_message'] is None


@mark.asyncio
async def test_save_model_complete_flow(mlflow):
    """Test save_model complete flow with all steps."""
    from model_manager.utils.models.train_model_params import TrainModelParams
    from model_manager.utils.models.train_model_result import TrainModelResult

    params = MagicMock(spec=TrainModelParams)
    params.experiment_name = 'production_model'

    train_result = MagicMock(spec=TrainModelResult)
    train_result.params = params
    train_result.run_name = None
    train_result.run_dir = None
    train_result.report_path = None

    # Mock complete flow
    mlflow.model_monitoring_repository.get_next_run_name.return_value = 'production_model-5'

    # After generate_artifacts, paths should be set
    updated_result = MagicMock(spec=TrainModelResult)
    updated_result.params = params
    updated_result.run_name = 'production_model-5'
    updated_result.run_dir = '/reports/production_model-5_20231010'
    updated_result.report_path = '/reports/production_model-5_20231010/report.html'
    updated_result.train_data_path = '/reports/production_model-5_20231010/train_data.csv'
    updated_result.test_data_path = '/reports/production_model-5_20231010/test_data.csv'

    mlflow.model_monitoring_repository.generate_artifacts.return_value = updated_result
    mlflow.model_monitoring_repository.save_run.return_value = None

    input_data = {
        **metadata,
        'train_result': train_result,
    }

    # Call the method
    response = await mlflow.save_model(input_data)

    # Verify complete flow
    mlflow.model_monitoring_repository.get_next_run_name.assert_called_once_with('production_model')
    mlflow.model_monitoring_repository.generate_artifacts.assert_called_once()
    mlflow.model_monitoring_repository.save_run.assert_called_once_with(updated_result)

    # Verify response
    assert response['success'] is True
    assert response['result'] == updated_result
    assert response['error_message'] is None
    assert updated_result.run_name == 'production_model-5'
    assert updated_result.run_dir is not None
    assert updated_result.report_path is not None