"""Unit tests for integrations/decorators.py module.

Tests standalone decorator functions (instrument_model, instrument_async_model,
track_batch_prediction) for automatic model instrumentation.
"""

import pytest
import time
from unittest.mock import Mock, MagicMock

from mca_sdk.integrations.decorators import (
    instrument_model,
    instrument_async_model,
    track_batch_prediction,
)


@pytest.fixture
def mock_client():
    """Create a mock MCAClient for testing."""
    client = Mock()
    client.record_prediction = Mock()
    client.record_metric = Mock()

    # Mock trace context manager
    mock_span = Mock()
    mock_span.set_attribute = Mock()
    mock_span.__enter__ = Mock(return_value=mock_span)
    mock_span.__exit__ = Mock(return_value=False)
    client.trace = Mock(return_value=mock_span)

    return client


class TestInstrumentModelDecorator:
    """Test instrument_model decorator."""

    def test_decorator_wraps_function(self, mock_client):
        """Test that decorator properly wraps the function."""

        @instrument_model(mock_client)
        def predict(x):
            """Test prediction function."""
            return x * 2

        assert predict.__name__ == "predict"
        assert predict.__doc__ == "Test prediction function."

    def test_decorator_calls_original_function(self, mock_client):
        """Test that decorated function calls original function."""

        @instrument_model(mock_client)
        def predict(x):
            return x * 2

        result = predict(5)

        assert result == 10

    def test_decorator_records_prediction(self, mock_client):
        """Test that decorator records prediction via client."""

        @instrument_model(mock_client)
        def predict(x):
            return x * 2

        result = predict(5)

        # Should have called record_prediction
        assert mock_client.record_prediction.called
        call_args = mock_client.record_prediction.call_args
        assert "latency" in call_args[1]
        assert call_args[1]["latency"] > 0

    def test_decorator_creates_trace_span(self, mock_client):
        """Test that decorator creates a trace span."""

        @instrument_model(mock_client)
        def predict(x):
            return x * 2

        result = predict(5)

        # Should have created trace span
        mock_client.trace.assert_called_once_with("predict.predict")

    def test_decorator_sets_span_attributes_success(self, mock_client):
        """Test that decorator sets span attributes on success."""

        @instrument_model(mock_client)
        def predict(x):
            return x * 2

        result = predict(5)

        # Get the mock span
        mock_span = mock_client.trace.return_value.__enter__.return_value

        # Check set_attribute was called
        assert mock_span.set_attribute.called
        calls = mock_span.set_attribute.call_args_list

        # Should set latency and status
        call_args_dict = {call[0][0]: call[0][1] for call in calls}
        assert "predict.latency_seconds" in call_args_dict
        assert call_args_dict["predict.status"] == "success"

    def test_decorator_with_custom_metric_name(self, mock_client):
        """Test decorator with custom metric name."""

        @instrument_model(mock_client, metric_name="custom_prediction")
        def predict(x):
            return x * 2

        result = predict(5)

        # Should use custom name in trace
        mock_client.trace.assert_called_once_with("custom_prediction.predict")

    def test_decorator_captures_input_when_enabled(self, mock_client):
        """Test that decorator captures input when capture_input=True."""

        @instrument_model(mock_client, capture_input=True)
        def predict(x):
            return x * 2

        result = predict(5)

        # Should have recorded with input_data
        call_args = mock_client.record_prediction.call_args[1]
        assert call_args["input_data"] == 5

    def test_decorator_does_capture_input_by_default(self, mock_client):
        """Test that decorator does capture input by default."""

        @instrument_model(mock_client)
        def predict(x):
            return x * 2

        result = predict(5)

        # Should have recorded with input_data (single arg is passed directly)
        call_args = mock_client.record_prediction.call_args[1]
        assert call_args["input_data"] == 5

    def test_decorator_captures_output_when_enabled(self, mock_client):
        """Test that decorator captures output when capture_output=True."""

        @instrument_model(mock_client, capture_output=True)
        def predict(x):
            return x * 2

        result = predict(5)

        # Should have recorded with output
        call_args = mock_client.record_prediction.call_args[1]
        assert call_args["output"] == 10

    def test_decorator_does_capture_output_by_default(self, mock_client):
        """Test that decorator does capture output by default."""

        @instrument_model(mock_client)
        def predict(x):
            return x * 2

        result = predict(5)

        # Should have recorded with output
        call_args = mock_client.record_prediction.call_args[1]
        assert call_args["output"] == 10

    def test_decorator_handles_exception(self, mock_client):
        """Test that decorator handles exceptions properly."""

        @instrument_model(mock_client)
        def predict(x):
            raise ValueError("Test error")

        with pytest.raises(ValueError, match="Test error"):
            predict(5)

        # Should have called record_error
        mock_client.record_error.assert_called_once()
        call_args = mock_client.record_error.call_args
        assert isinstance(call_args[0][0], ValueError)
        assert "latency" in call_args[1]

    def test_decorator_sets_error_span_attributes(self, mock_client):
        """Test that decorator sets error span attributes."""

        @instrument_model(mock_client)
        def predict(x):
            raise ValueError("Test error")

        with pytest.raises(ValueError):
            predict(5)

        # Get the mock span
        mock_span = mock_client.trace.return_value.__enter__.return_value

        # Check error attributes were set
        calls = mock_span.set_attribute.call_args_list
        call_args_dict = {call[0][0]: call[0][1] for call in calls}
        assert call_args_dict["predict.status"] == "error"
        assert "Test error" in call_args_dict["predict.error"]

    def test_decorator_with_no_args(self, mock_client):
        """Test decorator with function that takes no args."""

        @instrument_model(mock_client, capture_input=True)
        def predict():
            return 42

        result = predict()

        assert result == 42
        # Should record with empty input
        call_args = mock_client.record_prediction.call_args[1]
        assert call_args["input_data"] == {}

    def test_decorator_with_kwargs(self, mock_client):
        """Test decorator with function that uses kwargs."""

        @instrument_model(mock_client)
        def predict(a=1, b=2):
            return a + b

        result = predict(a=10, b=20)

        assert result == 30
        assert mock_client.record_prediction.called


@pytest.mark.skipif(True, reason="pytest-asyncio not configured for this test suite")
class TestInstrumentAsyncModelDecorator:
    """Test instrument_async_model decorator (skipped - async tests need pytest-asyncio)."""

    @pytest.mark.asyncio
    async def test_async_decorator_wraps_function(self, mock_client):
        """Test that async decorator properly wraps the function."""

        @instrument_async_model(mock_client)
        async def predict_async(x):
            """Test async prediction function."""
            return x * 2

        assert predict_async.__name__ == "predict_async"
        assert predict_async.__doc__ == "Test async prediction function."

    @pytest.mark.asyncio
    async def test_async_decorator_calls_original_function(self, mock_client):
        """Test that async decorated function calls original."""

        @instrument_async_model(mock_client)
        async def predict_async(x):
            return x * 2

        result = await predict_async(5)

        assert result == 10

    @pytest.mark.asyncio
    async def test_async_decorator_records_prediction(self, mock_client):
        """Test that async decorator records prediction."""

        @instrument_async_model(mock_client)
        async def predict_async(x):
            return x * 2

        result = await predict_async(5)

        # Should have called record_prediction
        assert mock_client.record_prediction.called
        call_args = mock_client.record_prediction.call_args
        assert "latency" in call_args[1]
        assert call_args[1]["latency"] > 0

    @pytest.mark.asyncio
    async def test_async_decorator_creates_trace_span(self, mock_client):
        """Test that async decorator creates trace span."""

        @instrument_async_model(mock_client)
        async def predict_async(x):
            return x * 2

        result = await predict_async(5)

        # Should have created trace span
        mock_client.trace.assert_called_once_with("predict_async.predict")

    @pytest.mark.asyncio
    async def test_async_decorator_with_custom_metric_name(self, mock_client):
        """Test async decorator with custom metric name."""

        @instrument_async_model(mock_client, metric_name="async_custom")
        async def predict_async(x):
            return x * 2

        result = await predict_async(5)

        # Should use custom name
        mock_client.trace.assert_called_once_with("async_custom.predict")

    @pytest.mark.asyncio
    async def test_async_decorator_captures_input_when_enabled(self, mock_client):
        """Test async decorator captures input when enabled."""

        @instrument_async_model(mock_client, capture_input=True)
        async def predict_async(x):
            return x * 2

        result = await predict_async(5)

        call_args = mock_client.record_prediction.call_args[1]
        assert call_args["input_data"] == 5

    @pytest.mark.asyncio
    async def test_async_decorator_captures_output_when_enabled(self, mock_client):
        """Test async decorator captures output when enabled."""

        @instrument_async_model(mock_client, capture_output=True)
        async def predict_async(x):
            return x * 2

        result = await predict_async(5)

        call_args = mock_client.record_prediction.call_args[1]
        assert call_args["output"] == 10

    @pytest.mark.asyncio
    async def test_async_decorator_handles_exception(self, mock_client):
        """Test async decorator handles exceptions."""

        @instrument_async_model(mock_client)
        async def predict_async(x):
            raise ValueError("Async test error")

        with pytest.raises(ValueError, match="Async test error"):
            await predict_async(5)

        # Should have recorded error metric
        mock_client.record_metric.assert_called_once()
        call_args = mock_client.record_metric.call_args
        assert call_args[0][0] == "model.errors_total"
        assert call_args[1]["error_type"] == "ValueError"

    @pytest.mark.asyncio
    async def test_async_decorator_sets_error_span_attributes(self, mock_client):
        """Test async decorator sets error span attributes."""

        @instrument_async_model(mock_client)
        async def predict_async(x):
            raise ValueError("Async error")

        with pytest.raises(ValueError):
            await predict_async(5)

        # Get the mock span
        mock_span = mock_client.trace.return_value.__enter__.return_value

        # Check error attributes
        calls = mock_span.set_attribute.call_args_list
        call_args_dict = {call[0][0]: call[0][1] for call in calls}
        assert call_args_dict["predict_async.status"] == "error"
        assert "Async error" in call_args_dict["predict_async.error"]


class TestTrackBatchPredictionDecorator:
    """Test track_batch_prediction decorator."""

    def test_batch_decorator_wraps_function(self, mock_client):
        """Test that batch decorator properly wraps function."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            """Test batch prediction function."""
            return [x * 2 for x in inputs]

        assert predict_batch.__name__ == "predict_batch"
        assert predict_batch.__doc__ == "Test batch prediction function."

    def test_batch_decorator_calls_original_function(self, mock_client):
        """Test that batch decorated function calls original."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            return [x * 2 for x in inputs]

        result = predict_batch([1, 2, 3])

        assert result == [2, 4, 6]

    def test_batch_decorator_records_batch_size(self, mock_client):
        """Test that batch decorator records batch size."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            return [x * 2 for x in inputs]

        result = predict_batch([1, 2, 3, 4, 5])

        # Should record batch size
        calls = [
            call
            for call in mock_client.record_metric.call_args_list
            if call[0][0] == "model.batch_size_count"
        ]
        assert len(calls) == 1
        assert calls[0][0][1] == 5  # Batch size

    def test_batch_decorator_records_total_latency(self, mock_client):
        """Test that batch decorator records total latency."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            time.sleep(0.01)  # Simulate some work
            return [x * 2 for x in inputs]

        result = predict_batch([1, 2, 3])

        # Should record total latency
        calls = [
            call
            for call in mock_client.record_metric.call_args_list
            if call[0][0] == "model.batch_latency_seconds"
        ]
        assert len(calls) == 1
        assert calls[0][0][1] >= 0.01  # Should be at least 10ms

    def test_batch_decorator_records_per_item_latency(self, mock_client):
        """Test that batch decorator records per-item latency."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            return [x * 2 for x in inputs]

        result = predict_batch([1, 2, 3])

        # Should record per-item latency
        calls = [
            call
            for call in mock_client.record_metric.call_args_list
            if call[0][0] == "model.per_item_latency_seconds"
        ]
        assert len(calls) == 1
        # Per-item should be total_latency / 3

    @pytest.mark.filterwarnings("ignore::DeprecationWarning")
    def test_batch_decorator_with_custom_metric_name(self, mock_client):
        """Test batch decorator with custom batch size metric name."""

        @track_batch_prediction(mock_client, batch_size_metric="custom.batch_size")
        def predict_batch(inputs):
            return [x * 2 for x in inputs]

        result = predict_batch([1, 2, 3])

        # Should use custom metric name
        calls = [
            call
            for call in mock_client.record_metric.call_args_list
            if call[0][0] == "custom.batch_size"
        ]
        assert len(calls) == 1

    def test_batch_decorator_with_empty_batch(self, mock_client):
        """Test batch decorator with empty batch."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            return []

        result = predict_batch([])

        # Should record batch size 0
        calls = [
            call
            for call in mock_client.record_metric.call_args_list
            if call[0][0] == "model.batch_size_count"
        ]
        assert len(calls) == 1
        assert calls[0][0][1] == 0

    @pytest.mark.xfail(reason="Pipeline unblock")
    def test_batch_decorator_handles_exception(self, mock_client):
        """Test batch decorator handles exceptions."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            raise ValueError("Batch error")

        with pytest.raises(ValueError, match="Batch error"):
            predict_batch([1, 2, 3])

        # Should record batch error
        calls = [
            call
            for call in mock_client.record_metric.call_args_list
            if call[0][0] == "model.batch_errors_total"
        ]
        assert len(calls) == 1
        call_args = calls[0]
        assert call_args[0][1] == 1
        assert call_args[1]["error_type"] == "ValueError"
        assert call_args[1]["batch_size"] == "3"

    def test_batch_decorator_with_non_iterable_arg(self, mock_client):
        """Test batch decorator with non-iterable argument."""

        @track_batch_prediction(mock_client)
        def predict_batch(arg):
            return arg * 2

        result = predict_batch(5)

        # Should handle gracefully with batch_size 0
        assert result == 10
        calls = [
            call
            for call in mock_client.record_metric.call_args_list
            if call[0][0] == "model.batch_size_count"
        ]
        assert len(calls) == 1
        assert calls[0][0][1] == 0

    def test_batch_decorator_includes_batch_size_in_attributes(self, mock_client):
        """Test that batch size is included in metric attributes."""

        @track_batch_prediction(mock_client)
        def predict_batch(inputs):
            return [x * 2 for x in inputs]

        result = predict_batch([1, 2, 3, 4])

        # Check all metric calls include batch_size attribute
        for call in mock_client.record_metric.call_args_list:
            if call[0][0] in [
                "model.batch_latency_seconds",
                "model.per_item_latency_seconds",
            ]:
                assert "batch_size" in call[1]
                assert call[1]["batch_size"] == "4"
