from unittest.mock import ANY, MagicMock, patch

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

from model_manager.activities.gates import Gates


@fixture
def gates_activity():
    gates = Gates(
        logger=MagicMock(),
        notification_handler=MagicMock(),
    )
    gates.error = MagicMock()
    gates.debug = MagicMock()
    gates.info = MagicMock()
    gates.warning = MagicMock()
    gates.critical = MagicMock()
    gates.send_notification = MagicMock()
    return gates


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


@mark.asyncio
async def test_input_gate_invalid_filter(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
        'data': {'value': [1, 2, 3]},
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.input_gate(input_data)

    # Assert
    assert result == (None, 0, '')
    gates_activity.error.assert_called_once_with(
        'Filter INVALID_FILTER not found', metadata['metadata']
    )


@mark.asyncio
@patch('model_manager.activities.gates.input_filter_functions')
async def test_input_gate_filter_exception(mock_input_filter_functions, gates_activity):
    # Arrange
    mock_input_filter_functions.__contains__.return_value = True
    mock_input_filter_functions.__getitem__.return_value = MagicMock(
        side_effect=Exception('Test error')
    )
    input_data = {
        **metadata,
        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
        'data': {'value': []},
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.input_gate(input_data)

    # Assert
    assert result == (None, 0, '')
    gates_activity.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='INTPUT_GATE_ERROR__EMPTY_DATA',
        message="Error in filter EMPTY_DATA:{'policy': 'STOP', 'config': {}}: \n Test error",
        block='input_gate',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_input_gate_no_filters(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {},
        'data': {'value': [1, 2, 3]},
        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
    }

    # Act
    result = await gates_activity.input_gate(input_data)

    # Assert
    assert result == (None, 0, '')
    gates_activity.debug.assert_called()


@mark.asyncio
async def test_input_gate_with_filter(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
        'data': {'value': []},
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.input_gate(input_data)

    # Assert
    assert result == ('STOP', -1, 'Input data with bad quality')
    gates_activity.debug.assert_called()


@mark.asyncio
async def test_input_gate_filter_returns_false(gates_activity):
    """Test to cover line 129 branch when filter returns False (filter passes)."""
    # Arrange - Use data that will NOT trigger EMPTY_DATA filter (has data)
    input_data = {
        **metadata,
        'filters': {'EMPTY_DATA': {'policy': 'STOP', 'config': {}}},
        'data': {'value': [1, 2, 3, 4, 5]},  # Has data, filter returns False
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.input_gate(input_data)

    # Assert - Filter returns False, so no policy is added to filter_output
    assert result == (None, 0, '')  # No filter triggered
    gates_activity.debug.assert_called()


@mark.asyncio
async def test_mlflow_response_gate_invalid_filter(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
        'data': {'content': {'message': 'success'}},
        'type': 'test',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_response_gate(input_data)

    # Assert
    assert result == (None, 0, '')


@mark.asyncio
@patch('model_manager.activities.gates.mlflow_response_filter_functions')
async def test_mlflow_response_gate_filter_exception(
    mock_mlflow_response_filter_functions, gates_activity
):
    # Arrange
    mock_mlflow_response_filter_functions.__contains__.return_value = True
    mock_mlflow_response_filter_functions.__getitem__.return_value = MagicMock(
        side_effect=Exception('Test error')
    )
    input_data = {
        **metadata,
        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
        'data': {'content': {'message': 'success'}},
        'type': 'test',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_response_gate(input_data)

    # Assert
    assert result == (None, 0, '')
    gates_activity.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='MLFLOW_GATE_RESPONSE_FILTER__INVALID_FILTER',
        message="Error in filter INVALID_FILTER:{'POLICY': 'STOP'}: \n Test error",
        block='mlflow_gate',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_mlflow_response_gate_no_filters(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {},
        'data': {'content': {'message': 'success'}},
        'type': 'test',
        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_response_gate(input_data)

    # Assert
    assert result == (None, 0, '')
    gates_activity.debug.assert_called()


@mark.asyncio
async def test_mlflow_response_gate_with_filter(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {'API_ERROR': {'policy': 'STOP'}},
        'data': {
            'success': False,
            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
        },
        'type': 'test',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_response_gate(input_data)

    # Assert
    assert result == ('STOP', -1, 'API error occurred')
    gates_activity.debug.assert_called()
    gates_activity.send_notification.assert_called()


@mark.asyncio
async def test_mlflow_response_gate_filter_returns_false(gates_activity):
    """Test to cover line 208 branch when filter returns False (no API error)."""
    # Arrange - Use data that will NOT trigger API_ERROR filter (success=True)
    input_data = {
        **metadata,
        'filters': {'API_ERROR': {'policy': 'STOP'}},
        'data': {
            'success': True,  # Success=True, filter returns False
            'content': {'message': 'Operation successful', 'result': 'data'},
        },
        'type': 'transform',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_response_gate(input_data)

    # Assert - Filter returns False, so no policy is added to filter_output
    assert result == (None, 0, '')  # No filter triggered
    gates_activity.debug.assert_called()


@mark.asyncio
async def test_mlflow_content_gate_invalid_filter(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {'INVALID_FILTER': {'POLICY': 'STOP'}},
        'data': {'value': [1, 2, 3]},
        'type': 'test',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_content_gate(input_data)

    # Assert
    assert result == (None, 0, '')


@mark.asyncio
@patch('model_manager.activities.gates.mlflow_content_filter_functions')
async def test_mlflow_content_gate_filter_exception(
    mock_mlflow_content_filter_functions, gates_activity
):
    # Arrange
    mock_mlflow_content_filter_functions.__contains__.return_value = True
    mock_mlflow_content_filter_functions.__getitem__.return_value = MagicMock(
        side_effect=Exception('Test error')
    )
    input_data = {
        **metadata,
        'filters': {'API_ERROR': {'POLICY': 'STOP'}},
        'data': {
            'success': False,
            'content': {'message': 'API error occurred', 'traceback': 'error trace'},
        },
        'type': 'test',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_content_gate(input_data)

    # Assert
    assert result == (None, 0, '')
    gates_activity.debug.assert_called()
    gates_activity.send_notification.assert_called_once_with(
        metadata=metadata['metadata'],
        notification_id='MLFLOW_GATE_CONTENT_FILTER__API_ERROR',
        message="Error in filter API_ERROR:{'POLICY': 'STOP'}: \n Test error",
        block='mlflow_gate',
        level=NotificationLevel.ERROR,
        attachment_content=ANY,
    )


@mark.asyncio
async def test_mlflow_content_gate_no_filters(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {},
        'data': {'value': [1, 2, 3]},
        'type': 'test',
        'path_priority': ['CONTINUE', 'STOP', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_content_gate(input_data)

    # Assert
    assert result == (None, 0, '')
    gates_activity.debug.assert_called()


@mark.asyncio
async def test_mlflow_content_gate_with_filter(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
        'data': {'value': [None, None, None]},
        'type': 'test',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_content_gate(input_data)

    # Assert
    assert result == ('STOP', -1, 'Transformed data not passed the content filter')
    gates_activity.debug.assert_called()
    gates_activity.send_notification.assert_called()


@mark.asyncio
async def test_mlflow_content_gate_filter_returns_false(gates_activity):
    """Test to cover line 293 branch when filter returns False (no NaN values)."""
    # Arrange - Use data that will NOT trigger NAN_VALUES filter (no NaN)
    input_data = {
        **metadata,
        'filters': {'NAN_VALUES': {'policy': 'STOP', 'config': {}}},
        'data': {'value': [1.0, 2.0, 3.0, 4.0, 5.0]},  # All valid numbers, no NaN
        'type': 'predict',
        'path_priority': ['STOP', 'CONTINUE', 'REPEAT'],
    }

    # Act
    result = await gates_activity.mlflow_content_gate(input_data)

    # Assert - Filter returns False, so no policy is added to filter_output
    assert result == (None, 0, '')  # No filter triggered
    gates_activity.debug.assert_called()


def test_get_prediction_store_policy_invalid_policy(gates_activity):
    # Arrange
    prediction_store_policy = 'INVALID_POLICY'

    # Act
    policy_type, policy_value = gates_activity.get_prediction_store_policy(
        prediction_store_policy, metadata
    )

    # Assert
    assert policy_type == 'lts'
    assert policy_value == 1


def test_get_prediction_store_policy_invalid_policy_value(gates_activity):
    # Arrange
    prediction_store_policy = 'abc:INVALID_VALUE'

    # Act
    policy_type, policy_value = gates_activity.get_prediction_store_policy(
        prediction_store_policy, metadata
    )

    # Assert
    assert policy_type == 'lts'
    assert policy_value == 1


def test_get_prediction_store_policy_valid_policy_type(gates_activity):
    # Arrange
    prediction_store_policy = 'abc:1'

    # Act
    policy_type, policy_value = gates_activity.get_prediction_store_policy(
        prediction_store_policy, metadata
    )

    # Assert
    assert policy_type == 'lts'
    assert policy_value == 1


def test_get_prediction_store_policy_valid_policy(gates_activity):
    # Arrange
    prediction_store_policy = 'erl:1'

    # Act
    policy_type, policy_value = gates_activity.get_prediction_store_policy(
        prediction_store_policy, metadata
    )

    # Assert
    assert policy_type == 'erl'
    assert policy_value == 1


@mark.asyncio
async def test_format_prediction_no_timestamp(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'data': {
            'prediction': {'2023-05-26 11:12:27': 1},
            'response_time': {'2023-05-26 11:12:27': 0.1},
        },
        'model_id': 'test_model',
        'prediction_confidence': 0.9,
        'prediction_store_policy': 'lts:1',
    }

    # Act
    result = await gates_activity.format_prediction(input_data)

    # Assert
    assert result['prediction'] == {0: 1}
    assert result['response_time'] == {0: ANY}
    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
    assert result['model_id'] == {0: 'test_model'}
    assert result['prediction_confidence'] == {0: 0.9}
    assert result['prediction_status'] == {0: 'Good'}
    assert result['comments'] == {0: ''}


@mark.asyncio
async def test_format_prediction_with_timestamp_erl(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'data': {
            'prediction': {
                '2023-05-26 11:12:27': 1,
                '2023-05-26 11:12:28': 2,
                '2023-05-26 11:12:29': 3,
            },
            'response_time': {
                '2023-05-26 11:12:27': 0.1,
                '2023-05-26 11:12:28': 0.2,
                '2023-05-26 11:12:29': 0.3,
            },
        },
        'model_id': 'test_model',
        'prediction_confidence': 0.9,
        'prediction_store_policy': 'erl:2',
    }

    # Act
    result = await gates_activity.format_prediction(input_data)

    # Assert
    assert result['prediction'] == {0: 2, 1: 1}
    assert result['response_time'] == {0: 0.2, 1: 0.1}
    assert result['timestamp'] == {0: '2023-05-26 11:12:28', 1: '2023-05-26 11:12:27'}
    assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
    assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
    assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
    assert result['comments'] == {0: '', 1: ''}


@mark.asyncio
async def test_format_prediction_with_timestamp_lts(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'data': {
            'prediction': {
                '2023-05-26 11:12:27': 1,
                '2023-05-26 11:12:28': 2,
                '2023-05-26 11:12:29': 3,
            },
            'response_time': {
                '2023-05-26 11:12:27': 0.1,
                '2023-05-26 11:12:28': 0.2,
                '2023-05-26 11:12:29': 0.3,
            },
        },
        'model_id': 'test_model',
        'prediction_confidence': 0.9,
        'prediction_store_policy': 'lts:2',
    }

    # Act
    result = await gates_activity.format_prediction(input_data)

    # Assert
    assert result['prediction'] == {0: 3, 1: 2}
    assert result['response_time'] == {0: 0.3, 1: 0.2}
    assert result['timestamp'] == {0: '2023-05-26 11:12:29', 1: '2023-05-26 11:12:28'}
    assert result['model_id'] == {0: 'test_model', 1: 'test_model'}
    assert result['prediction_confidence'] == {0: 0.9, 1: 0.9}
    assert result['prediction_status'] == {0: 'Good', 1: 'Good'}
    assert result['comments'] == {0: '', 1: ''}


@mark.asyncio
async def test_format_prediction_with_timestamp_invalid_policy(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'data': {
            'prediction': [1, 2, 3],
            'response_time': [0.1, 0.2, 0.3],
            'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28', '2023-05-26 11:12:29'],
        },
        'model_id': 'test_model',
        'prediction_confidence': 0.9,
        'prediction_store_policy': 'lts:2',
    }
    gates_activity.get_prediction_store_policy = MagicMock(return_value=('invalid', 1))

    try:
        await gates_activity.format_prediction(input_data)
    except ValueError as e:
        assert str(e) == 'Invalid policy type: invalid'
    else:
        raise AssertionError('Expected ValueError')


@mark.asyncio
async def test_format_default_prediction(gates_activity):
    # Arrange
    input_data = {
        **metadata,
        'timestamp': '2023-05-26 11:12:27',
        'model_id': 'test_model',
        'prediction_confidence': 0.1,
        'comment': 'Test comment',
    }

    # Act
    result = await gates_activity.format_default_prediction(input_data)

    # Assert
    assert result['prediction'] == {0: 0}
    assert result['response_time'] == {0: 0}
    assert result['timestamp'] == {0: '2023-05-26 11:12:27'}
    assert result['model_id'] == {0: 'test_model'}
    assert result['prediction_confidence'] == {0: 0.1}
    assert result['prediction_status'] == {0: 'Bad'}
    assert result['comments'] == {0: 'Test comment'}
    gates_activity.debug.assert_called()


@mark.asyncio
async def test_get_last_timestamp_with_data(gates_activity):
    # Arrange
    input_data = {**metadata, 'data': {'timestamp': ['2023-05-26 11:12:27', '2023-05-26 11:12:28']}}

    # Act
    result = await gates_activity.get_last_timestamp(input_data)

    # Assert
    assert result == '2023-05-26 11:12:28'


@mark.asyncio
async def test_get_last_timestamp_no_data(gates_activity):
    # Arrange
    input_data = {'data': {}, **metadata}

    # Act
    result = await gates_activity.get_last_timestamp(input_data)

    # Assert
    assert isinstance(result, str)  # Should be a timestamp string
    assert len(result) > 0


@mark.asyncio
@patch('model_manager.activities.gates.metrics')
async def test_write_metrics(mock_metrics, gates_activity):
    """Test write_metrics method."""
    input_data = {
        **metadata,
        'prediction': {
            'prediction': [1, 2, 3],
            'prediction_confidence': [0.9, 0.8, 0.7],
            'response_time': [0.1, 0.2, 0.3],
        },
    }
    await gates_activity.write_metrics(input_data)
    mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.assert_called_once_with(
        pod_id=gates_activity.pod_id,
        model_name=metadata['metadata']['model_name'],
        pipeline_name=metadata['metadata']['workflow_name'],
    )
    mock_metrics.PREDICTIONS_WRITTEN_COUNT.labels.return_value.inc.assert_called_once_with()

    mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.assert_called_once_with(
        pod_id=gates_activity.pod_id,
        model_name=metadata['metadata']['model_name'],
        pipeline_name=metadata['metadata']['workflow_name'],
    )
    mock_metrics.PREDICTION_CONFIDENCE_MONITOR.labels.return_value.set.assert_called_once_with(0.9)

    mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.assert_called_once_with(
        pod_id=gates_activity.pod_id,
        model_name=metadata['metadata']['model_name'],
        pipeline_name=metadata['metadata']['workflow_name'],
    )
    mock_metrics.PREDICTION_RESPONSE_TIME_MONITOR.labels.return_value.observe.assert_called_once_with(
        0.1
    )