"""Unit tests for autolog validation and robustness improvements."""

import pytest
from unittest.mock import MagicMock, patch
import sys


class TestAutologValidation:
    """Test autolog parameter validation."""

    def test_invalid_framework_name(self):
        """Test that invalid framework names raise ValueError."""
        mock_client = MagicMock()
        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=mock_client):
            from mca_sdk import autolog

            with pytest.raises(ValueError, match="Invalid framework names"):
                autolog(frameworks=["pytorch", "tensorflow"])

    def test_invalid_exclude_name(self):
        """Test that invalid exclude names raise ValueError."""
        mock_client = MagicMock()
        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=mock_client):
            from mca_sdk import autolog

            with pytest.raises(ValueError, match="Invalid framework names in exclude"):
                autolog(exclude=["pytorch"])

    def test_valid_framework_names(self):
        """Test that valid framework names are accepted."""
        # Mock sklearn
        sklearn_base = MagicMock()

        class MockBaseEstimator:
            def predict(self, X):
                return [1, 0]

        sklearn_base.BaseEstimator = MockBaseEstimator
        sys.modules["sklearn"] = MagicMock()
        sys.modules["sklearn.base"] = sklearn_base

        mock_client = MagicMock()
        mock_client.trace = MagicMock(
            return_value=MagicMock(__enter__=MagicMock(), __exit__=MagicMock())
        )

        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=mock_client):
            from mca_sdk import autolog

            # Should not raise
            autolog(frameworks=["sklearn"])

    def test_max_payload_size_parameter(self):
        """Test that max_payload_size parameter is accepted and used."""
        sklearn_base = MagicMock()

        class MockBaseEstimator:
            def predict(self, X):
                return [1, 0]

        sklearn_base.BaseEstimator = MockBaseEstimator
        sys.modules["sklearn"] = MagicMock()
        sys.modules["sklearn.base"] = sklearn_base

        mock_client = MagicMock()
        mock_client.trace = MagicMock(
            return_value=MagicMock(__enter__=MagicMock(), __exit__=MagicMock())
        )

        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=mock_client):
            from mca_sdk import autolog

            # Should not raise and should accept custom max_payload_size
            autolog(frameworks=["sklearn"], max_payload_size=5000)


class TestPayloadSizeLimiting:
    """Test payload size limiting functionality."""

    def test_payload_truncation_for_large_output(self, mocker):
        """Test that large outputs are truncated to prevent memory bloat."""

        # Create a model that returns large output
        class MockEstimator:
            def predict(self, X):
                # Return a very large string (>10KB)
                return ["x" * 20000]

        sklearn_mock = MagicMock()
        sklearn_mock.ensemble.RandomForestClassifier = MockEstimator
        sys.modules["sklearn"] = sklearn_mock

        sklearn_base_mock = MagicMock()
        sklearn_base_mock.BaseEstimator = MockEstimator
        sys.modules["sklearn.base"] = sklearn_base_mock

        mock_client = MagicMock()
        mock_client.trace = MagicMock(
            return_value=MagicMock(__enter__=MagicMock(), __exit__=MagicMock())
        )

        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=mock_client):
            from mca_sdk.autolog._sklearn import patch_sklearn

            # Patch with small payload limit
            patch_sklearn(capture_input=False, capture_output=True, max_payload_size=100)

            estimator = MockEstimator()
            result = estimator.predict([[1, 2]])

            # Verify prediction still works
            assert result == ["x" * 20000]

            # Verify metric was recorded
            mock_client.record_prediction.assert_called()

            # Verify output was truncated (check call kwargs)
            call_kwargs = mock_client.record_prediction.call_args[1]
            if "output_sample" in call_kwargs:
                # Should be truncated
                assert "truncated" in call_kwargs["output_sample"] or call_kwargs.get(
                    "output_truncated"
                )


class TestGracefulClientHandling:
    """Test graceful handling when MCAClient is not initialized."""

    def test_prediction_works_without_client_after_patch(self, mocker):
        """Test that predictions work even if client becomes unavailable after patching."""

        class MockEstimator:
            def predict(self, X):
                return [1, 0]

        sklearn_mock = MagicMock()
        sklearn_mock.ensemble.RandomForestClassifier = MockEstimator
        sys.modules["sklearn"] = sklearn_mock

        sklearn_base_mock = MagicMock()
        sklearn_base_mock.BaseEstimator = MockEstimator
        sys.modules["sklearn.base"] = sklearn_base_mock

        # First patch with client available
        mock_client = MagicMock()
        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=mock_client):
            from mca_sdk.autolog._sklearn import patch_sklearn

            patch_sklearn()

        # Now client becomes unavailable
        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=None):
            # Prediction should still work without crashing
            estimator = MockEstimator()
            result = estimator.predict([[1, 2]])
            assert result == [1, 0]


class TestErrorHandling:
    """Test proper error handling and span marking."""

    def test_exception_is_recorded_and_reraised(self, mocker):
        """Test that exceptions are recorded in telemetry and re-raised properly."""

        class MockEstimator:
            def predict(self, X):
                raise ValueError("Model failed")

        sklearn_mock = MagicMock()
        sklearn_mock.ensemble.RandomForestClassifier = MockEstimator
        sys.modules["sklearn"] = sklearn_mock

        sklearn_base_mock = MagicMock()
        sklearn_base_mock.BaseEstimator = MockEstimator
        sys.modules["sklearn.base"] = sklearn_base_mock

        mock_client = MagicMock()
        mock_span = MagicMock()
        mock_client.trace = MagicMock(return_value=mock_span)
        mock_span.__enter__ = MagicMock(return_value=mock_span)
        mock_span.__exit__ = MagicMock(return_value=False)

        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=mock_client):
            from mca_sdk.autolog._sklearn import patch_sklearn

            patch_sklearn()

            estimator = MockEstimator()

            # Exception should be raised
            with pytest.raises(ValueError, match="Model failed"):
                estimator.predict([[1, 2]])

            # Verify error was recorded in metrics
            mock_client.record_prediction.assert_called()
            call_kwargs = mock_client.record_prediction.call_args[1]
            assert "error" in call_kwargs or "failed" in call_kwargs
