from unittest.mock import MagicMock, patch

import numpy as np
from pytest import fixture, mark
from laborious.activities.mlflow import MLFlow


@patch("laborious.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"
    )


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


@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
async def test_request_transform(mock_max, mock_dataframe, mlflow):
    mock_max.return_value = '2024-01-02'
    # Mock input data
    input_data = {
        'data': [
            {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
            {'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
            {'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
            {'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
        ],
        'model_name': 'test_model',
        'model_retention': 30
    }

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

    # 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, 30
    )


@mark.asyncio
@patch("laborious.activities.mlflow.DataFrame")
@patch("laborious.activities.mlflow.max")
async def test_request_predict(mock_max, mock_dataframe, mlflow):
    mock_max.return_value = '2024-01-02'
    # Mock input data
    input_data = {
        'data': [
            {'timestamp': '2024-01-01', 'variable': 'var1', 'value': 1.0},
            {'timestamp': '2024-01-01', 'variable': 'var2', 'value': 2.0},
            {'timestamp': '2024-01-02', 'variable': 'var1', 'value': 3.0},
            {'timestamp': '2024-01-02', 'variable': 'var2', 'value': 4.0}
        ],
        'model_name': 'test_model',
        'model_retention': 30
    }

    # 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
    )

    # 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, 30
    )