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

import pytest_asyncio
from pytest import fixture, mark
from sientia_do.notifications.models import NotificationLevel

from laborious.activities.api import API, PI_WEB_API_PREDICTION_ERROR_CONFIDENCE

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


def _create_mock_dataframe(to_dict_return=None):
    """Helper function to create a mocked DataFrame for testing."""
    mock_df = MagicMock()
    mock_head = MagicMock()

    def get_column_values(key):
        if key == 'prediction':
            return MagicMock(values=[0.75])
        elif key == 'prediction_confidence':
            return MagicMock(values=[0.95])
        else:
            return MagicMock(values=['2024-01-01T00:00:00+00:00'])

    mock_head.__getitem__.side_effect = get_column_values
    mock_df.head.return_value = mock_head

    if to_dict_return is None:
        to_dict_return = {
            'prediction': [0.75],
            'prediction_confidence': [0.95],
            'timestamp': ['2024-01-01T00:00:00+00:00'],
        }
    mock_df.to_dict.return_value = to_dict_return

    return mock_df


@fixture
def base_input_data():
    """Base input data for PI Web API tests."""
    return {
        **metadata,
        'data': {
            'prediction': [0.75],
            'prediction_confidence': [0.95],
            'timestamp': ['2024-01-01T00:00:00+00:00'],
        },
        'pi_web_api_output_config': {
            'endpoint': 'https://test-pi-server.com/piwebapi',
            'prediction_tags': {'tag1': 'web_id_1'},
            'confidence_tags': {'tag2': 'web_id_2'},
        },
    }


def test__init__():
    api = API(
        base_url='https://test-pi-server.com',
        auth_type='bearer',
        auth_token='test_token',
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
    )

    assert api.pi_web_api_client is not None


@pytest_asyncio.fixture
@patch('laborious.activities.api.PIWebAPIClient')
def api(mock_pi_web_api_client):
    mock_client = MagicMock()
    mock_client.write_value = AsyncMock()
    mock_client.close = MagicMock()
    mock_pi_web_api_client.return_value = mock_client

    api_instance = API(
        base_url='https://test-pi-server.com',
        auth_type='bearer',
        auth_token='test_token',
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=AsyncMock(),
    )
    api_instance.send_notification_async = AsyncMock()
    api_instance.info = MagicMock()
    return api_instance


@mark.asyncio
@patch('laborious.activities.api.DataFrame')
async def test_write_pi_web_api_data_success(mock_dataframe, api, base_input_data):
    input_data = {
        **base_input_data,
        'pi_web_api_output_config': {
            'endpoint': 'https://test-pi-server.com/piwebapi',
            'prediction_tags': {'tag1': 'web_id_1', 'tag2': 'web_id_2'},
            'confidence_tags': {'tag3': 'web_id_3', 'tag4': 'web_id_4'},
        },
    }

    mock_dataframe.return_value = _create_mock_dataframe()

    result = await api.write_pi_web_api_data(input_data)

    api.info.assert_called_once_with('Writing data to PI Web API...', metadata['metadata'])

    api.pi_web_api_client.write_value.assert_has_calls(
        [
            call(
                web_ids=['web_id_1', 'web_id_2'],
                value={
                    'Timestamp': '2024-01-01T00:00:00+00:00',
                    'Value': 0.75,
                },
                endpoint='https://test-pi-server.com/piwebapi',
                metadata=metadata['metadata'],
            ),
            call(
                web_ids=['web_id_3', 'web_id_4'],
                value={
                    'Timestamp': '2024-01-01T00:00:00+00:00',
                    'Value': 0.95,
                },
                endpoint='https://test-pi-server.com/piwebapi',
                metadata=metadata['metadata'],
            ),
        ]
    )

    assert result == {
        'prediction': [0.75],
        'prediction_confidence': [0.95],
        'timestamp': ['2024-01-01T00:00:00+00:00'],
    }


@mark.asyncio
@patch('laborious.activities.api.DataFrame')
async def test_write_pi_web_api_data_prediction_error(mock_dataframe, api, base_input_data):
    mock_dataframe.return_value = _create_mock_dataframe(
        {
            'prediction': [0.75],
            'prediction_confidence': [PI_WEB_API_PREDICTION_ERROR_CONFIDENCE],
            'timestamp': ['2024-01-01T00:00:00+00:00'],
        }
    )

    api.pi_web_api_client.write_value.side_effect = Exception('Prediction write failed')

    result = await api.write_pi_web_api_data(base_input_data)

    api.send_notification_async.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
        message="Error writing prediction data to PI Web API: Prediction write failed\n Tags: {'tag1': 'web_id_1'}",
        block='write_pi_web_api_data',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )

    assert result['prediction_confidence'][0] == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
    assert api.pi_web_api_client.write_value.call_count == 1


@mark.asyncio
@patch('laborious.activities.api.DataFrame')
async def test_write_pi_web_api_data_confidence_error(mock_dataframe, api, base_input_data):
    mock_dataframe.return_value = _create_mock_dataframe()

    api.pi_web_api_client.write_value.side_effect = [
        None,
        Exception('Confidence write failed'),
    ]

    result = await api.write_pi_web_api_data(base_input_data)

    api.send_notification_async.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='WRITE_PI_WEB_API_CONFIDENCE_ERROR',
        message="Error writing confidence data to PI Web API: Confidence write failed\n Tags: {'tag2': 'web_id_2'}",
        block='write_pi_web_api_data',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )

    assert result == {
        'prediction': [0.75],
        'prediction_confidence': [0.95],
        'timestamp': ['2024-01-01T00:00:00+00:00'],
    }
    assert api.pi_web_api_client.write_value.call_count == 2


@mark.asyncio
@patch('laborious.activities.api.DataFrame')
async def test_write_pi_web_api_data_empty_tags(mock_dataframe, api, base_input_data):
    input_data = {
        **base_input_data,
        'pi_web_api_output_config': {
            'endpoint': 'https://test-pi-server.com/piwebapi',
            'prediction_tags': {},
            'confidence_tags': {},
        },
    }

    mock_dataframe.return_value = _create_mock_dataframe()

    result = await api.write_pi_web_api_data(input_data)

    api.pi_web_api_client.write_value.assert_has_calls(
        [
            call(
                web_ids=[],
                value={
                    'Timestamp': '2024-01-01T00:00:00+00:00',
                    'Value': 0.75,
                },
                endpoint='https://test-pi-server.com/piwebapi',
                metadata=metadata['metadata'],
            ),
            call(
                web_ids=[],
                value={
                    'Timestamp': '2024-01-01T00:00:00+00:00',
                    'Value': 0.95,
                },
                endpoint='https://test-pi-server.com/piwebapi',
                metadata=metadata['metadata'],
            ),
        ]
    )

    assert result == {
        'prediction': [0.75],
        'prediction_confidence': [0.95],
        'timestamp': ['2024-01-01T00:00:00+00:00'],
    }


@mark.asyncio
async def test_close(api):
    api.close()

    api.pi_web_api_client.close.assert_called_once()