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

from pytest import fixture
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'},
        },
    }


@patch('laborious.activities.api.PIWebAPIClient')
def test_get_pi_web_api_core_labels_without_operation_type(mock_pi_web_api_client):
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring

    api_instance = API(
        base_url='https://test-pi-server.com',
        auth_type='bearer',
        auth_token='test_token',
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=MagicMock(),
    )
    with patch.object(
        SientiaMonitoring,
        'get_core_labels',
        return_value={
            'pod_id': 'test_pod',
            'runtime': 'local',
            'model_name': 'test_model',
            'workflow_name': 'test_workflow',
            'operation_type': 'write_pi_web_api_data',
        },
    ):
        labels = api_instance.get_pi_web_api_core_labels(metadata=metadata['metadata'])
    assert labels['operation_type'] == 'write_pi_web_api_data'
    assert labels == {
        'pod_id': 'test_pod',
        'runtime': 'local',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
        'operation_type': 'write_pi_web_api_data',
    }


@patch('laborious.activities.api.PIWebAPIClient')
def test_get_pi_web_api_core_labels_with_operation_type(mock_pi_web_api_client):
    from sientia_do.observability.sientia_monitoring import SientiaMonitoring

    api_instance = API(
        base_url='https://test-pi-server.com',
        auth_type='bearer',
        auth_token='test_token',
        logger=MagicMock(),
        notification_handler=MagicMock(),
        metrics_controller=MagicMock(),
    )
    with patch.object(
        SientiaMonitoring,
        'get_core_labels',
        return_value={
            'pod_id': 'test_pod',
            'runtime': 'k8s',
            'model_name': 'test_model',
            'workflow_name': 'test_workflow',
            'operation_type': 'write',
        },
    ):
        labels = api_instance.get_pi_web_api_core_labels(
            metadata=metadata['metadata'], operation_type='write'
        )
    assert labels['operation_type'] == 'write'
    assert labels['runtime'] == 'k8s'


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=MagicMock(),
    )

    assert api.pi_web_api_client is not None


@fixture
@patch('laborious.activities.api.PIWebAPIClient')
def api(mock_pi_web_api_client):
    mock_client = MagicMock()
    mock_client.write_value = MagicMock()
    mock_client.close = MagicMock()
    mock_client.base_url = 'https://test-pi-server.com'
    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=MagicMock(),
    )
    api_instance.send_notification = MagicMock()
    api_instance.info = MagicMock()
    api_instance.error = MagicMock()
    api_instance.emit_metric_sync = MagicMock()
    api_instance.get_core_labels = MagicMock(
        return_value={
            'pod_id': 'test_pod',
            'runtime': 'local',
            'model_name': 'test_model',
            'workflow_name': 'test_workflow',
        }
    )
    return api_instance


@patch('laborious.activities.api.DataFrame')
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()

    # Mock successful responses
    api.pi_web_api_client.write_value.side_effect = [
        [{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': []}],
        [{'WebId': 'web_id_3', 'Errors': []}, {'WebId': 'web_id_4', 'Errors': []}],
    ]

    result = api.write_pi_web_api_data(input_data)

    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,
                },
                metadata=metadata['metadata'],
            ),
            call(
                web_ids=['web_id_3', 'web_id_4'],
                value={
                    'Timestamp': '2024-01-01T00:00:00+00:00',
                    'Value': 0.95,
                },
                metadata=metadata['metadata'],
            ),
        ]
    )

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


@patch('laborious.activities.api.DataFrame')
def test_write_pi_web_api_data_partial_prediction_mismatch_downgrades_confidence(
    mock_dataframe, api, base_input_data
):
    """A partial write (some tags report Errors) doesn't raise — it's caught by the tag-count
    mismatch check, which still must downgrade confidence and attach the mismatch message."""
    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'},
        },
    }
    mock_df = _create_mock_dataframe()
    mock_dataframe.return_value = mock_df

    api.pi_web_api_client.write_value.side_effect = [
        [{'WebId': 'web_id_1', 'Errors': []}, {'WebId': 'web_id_2', 'Errors': ['bad tag']}],
        [{'WebId': 'web_id_3', 'Errors': []}],
    ]

    api.write_pi_web_api_data(input_data)

    mock_df.__setitem__.assert_any_call(
        'prediction_confidence', PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
    )
    mock_df.__setitem__.assert_any_call('comments', ANY)
    api.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='WRITE_PI_WEB_API_PREDICTION_ERROR',
        message=ANY,
        block='write_pi_web_api_data',
        level=NotificationLevel.ERROR,
    )


@patch('laborious.activities.api.DataFrame')
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 = api.write_pi_web_api_data(base_input_data)

    api.send_notification.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


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

    # First call succeeds, second fails
    api.pi_web_api_client.write_value.side_effect = [
        [{'WebId': 'web_id_1', 'Errors': []}],
        Exception('Confidence write failed'),
    ]

    result = api.write_pi_web_api_data(base_input_data)

    api.send_notification.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


@patch('laborious.activities.api.DataFrame')
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()

    # Mock empty responses
    api.pi_web_api_client.write_value.side_effect = [
        [],
        [],
    ]

    result = 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,
                },
                metadata=metadata['metadata'],
            ),
            call(
                web_ids=[],
                value={
                    'Timestamp': '2024-01-01T00:00:00+00:00',
                    'Value': 0.95,
                },
                metadata=metadata['metadata'],
            ),
        ]
    )

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


@patch('laborious.activities.api.DataFrame')
def test_write_pi_web_api_data_updates_confidence_and_comments(
    mock_dataframe, api, base_input_data
):
    mock_dataframe.return_value = _create_mock_dataframe()
    api.pi_web_api_client.write_value.side_effect = [
        [{'WebId': 'web_id_1', 'Errors': []}],
        [{'WebId': 'web_id_2', 'Errors': []}],
    ]
    with patch.object(
        api,
        'process_pi_web_api_response',
        new=MagicMock(side_effect=[(0.33, 'PI warning'), (0, '')]),
    ) as process_mock:
        result = api.write_pi_web_api_data(base_input_data)

    assert process_mock.call_count == 2
    assert result is not None


@patch('laborious.activities.api.SientiaMonitoring.shutdown')
def test_close(mock_shutdown, api):
    api.close()

    api.pi_web_api_client.close.assert_called_once()
    mock_shutdown.assert_called_once_with(api)


def test_process_pi_web_api_response_success(api):
    """Test successful processing of PI Web API response with all tags written."""
    response_data = [
        {'WebId': 'web_id_1', 'Errors': []},
        {'WebId': 'web_id_2', 'Errors': []},
    ]
    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
    core_labels = {
        'pod_id': 'test_pod',
        'runtime': 'local',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
    }

    confidence, message = api.process_pi_web_api_response(
        response_data=response_data,
        tags=tags,
        core_labels=core_labels,
        metadata=metadata['metadata'],
    )

    assert confidence == 0
    assert message == ''
    assert api.emit_metric_sync.call_count == 2
    # Verify that emit_metric_sync was called with correct tags structure
    call_args_list = api.emit_metric_sync.call_args_list
    assert len(call_args_list) == 2
    # Check that all calls include core_labels and tag_name
    for call_args in call_args_list:
        assert 'tag_name' in call_args.kwargs['tags']
        assert call_args.kwargs['tags']['tag_name'] in ['tag1', 'tag2']


def test_process_pi_web_api_response_with_errors(api):
    """Test processing response with errors in some tags."""
    response_data = [
        {'WebId': 'web_id_1', 'Errors': ['Error writing tag']},
        {'WebId': 'web_id_2', 'Errors': []},
    ]
    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
    core_labels = {
        'pod_id': 'test_pod',
        'runtime': 'local',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
    }

    confidence, message = api.process_pi_web_api_response(
        response_data=response_data,
        tags=tags,
        core_labels=core_labels,
        metadata=metadata['metadata'],
    )

    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
    assert (
        message
        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
    )
    assert api.emit_metric_sync.call_count == 2


def test_process_pi_web_api_response_missing_tags(api):
    """Test processing response when number of written tags doesn't match expected."""
    response_data = [
        {'WebId': 'web_id_1', 'Errors': []},
    ]
    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
    core_labels = {
        'pod_id': 'test_pod',
        'runtime': 'local',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
    }

    confidence, message = api.process_pi_web_api_response(
        response_data=response_data,
        tags=tags,
        core_labels=core_labels,
        metadata=metadata['metadata'],
    )

    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
    assert (
        message
        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag1'] tags were written."
    )
    api.send_notification.assert_called_once()
    call_args = api.send_notification.call_args
    assert call_args.kwargs['notification_id'] == 'WRITE_PI_WEB_API_PREDICTION_ERROR'
    assert call_args.kwargs['level'] == NotificationLevel.ERROR


def test_process_pi_web_api_response_missing_webid(api):
    """Test processing response when WebId is missing in response item."""
    response_data = [
        {'Errors': []},
        {'WebId': 'web_id_2', 'Errors': []},
    ]
    tags = {'tag1': 'web_id_1', 'tag2': 'web_id_2'}
    core_labels = {
        'pod_id': 'test_pod',
        'runtime': 'local',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
    }

    confidence, message = api.process_pi_web_api_response(
        response_data=response_data,
        tags=tags,
        core_labels=core_labels,
        metadata=metadata['metadata'],
    )

    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
    assert (
        message
        == "The number of written tags does not match the number of tag names: Expected ['tag1', 'tag2'] tags, but ['tag2'] tags were written."
    )
    api.error.assert_any_call('The response did not contain some WebIds', metadata['metadata'])


def test_process_pi_web_api_response_missing_tag_name(api):
    """Test processing response when tag name is not found for WebId."""
    response_data = [
        {'WebId': 'unknown_web_id', 'Errors': []},
    ]
    tags = {'tag1': 'web_id_1'}
    core_labels = {
        'pod_id': 'test_pod',
        'runtime': 'local',
        'model_name': 'test_model',
        'workflow_name': 'test_workflow',
    }

    confidence, message = api.process_pi_web_api_response(
        response_data=response_data,
        tags=tags,
        core_labels=core_labels,
        metadata=metadata['metadata'],
    )

    assert confidence == PI_WEB_API_PREDICTION_ERROR_CONFIDENCE
    assert (
        message
        == "The number of written tags does not match the number of tag names: Expected ['tag1'] tags, but [] tags were written."
    )
    api.error.assert_any_call(
        'The response did not contain the tag name for WebId unknown_web_id', metadata['metadata']
    )