"""Unit tests for scikit-learn autolog functionality."""

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


class TestSklearnAutolog:
    """Test autolog for scikit-learn framework."""

    def test_patch_sklearn_predict(self, mocker):
        """Test that sklearn predict() is instrumented."""

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

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

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

        # Mock MCAClient singleton
        mock_client = MagicMock()
        mock_client.trace = MagicMock(
            return_value=MagicMock(__enter__=MagicMock(), __exit__=MagicMock())
        )
        mock_client.record_prediction = MagicMock()

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

            # Apply autolog
            autolog(frameworks=["sklearn"])

            # Verify that BaseEstimator.predict was patched
            assert hasattr(MockBaseEstimator, "predict")

            # Verify instrumentation works
            estimator = MockBaseEstimator()
            estimator.predict([[1, 2]])
            mock_client.record_prediction.assert_called()

    def test_sklearn_instrumentation_records_metrics(self, mocker):
        """Test that sklearn predict() calls record metrics."""

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

        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 MCAClient
        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 sklearn
            patch_sklearn(capture_input=False, capture_output=False)

            # Make prediction
            estimator = MockEstimator()
            result = estimator.predict([[1, 2, 3]])

            # Verify instrumentation was called
            assert result == [1, 0, 1]
            mock_client.trace.assert_called()
            mock_client.record_prediction.assert_called()

    def test_sklearn_capture_input_output(self, mocker):
        """Test that sklearn autolog captures input/output shapes when requested."""

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

        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 numpy array with shape
        mock_input = MagicMock()
        mock_input.shape = (100, 10)

        # Mock MCAClient
        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 capture enabled
            patch_sklearn(capture_input=True, capture_output=True)

            # Make prediction
            estimator = MockEstimator()
            estimator.predict(mock_input)

            # Story 8.5: Verify record_prediction was called with raw input_data
            # Shape extraction now happens inside record_prediction()
            mock_client.record_prediction.assert_called()
            call_kwargs = mock_client.record_prediction.call_args[1]
            assert "input_data" in call_kwargs
            # input_data should be the raw mock object (not pre-serialized with shape)
            assert call_kwargs["input_data"] == mock_input

    def test_sklearn_graceful_fallback_no_client(self, mocker):
        """Test that sklearn autolog gracefully falls back if client not initialized."""

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

        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 MCAClient.get_instance() to return None
        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=None):
            from mca_sdk.autolog._sklearn import patch_sklearn

            # Patch sklearn
            patch_sklearn(capture_input=False, capture_output=False)

            # Make prediction (should not raise error)
            estimator = MockEstimator()
            result = estimator.predict([[1, 2, 3]])

            # Should still work without instrumentation
            assert result == [1, 0, 1]

    def test_sklearn_exception_handling(self, mocker):
        """Test that sklearn autolog handles exceptions and re-raises them."""

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

        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 MCAClient
        mock_client = MagicMock()
        mock_context = MagicMock()
        mock_client.trace = MagicMock(return_value=mock_context)
        mock_context.__enter__ = MagicMock(return_value=mock_context)
        mock_context.__exit__ = MagicMock(return_value=False)  # Don't suppress exception

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

            # Patch sklearn
            patch_sklearn(capture_input=False, capture_output=False)

            # Make prediction (should raise error AND record metrics)
            estimator = MockEstimator()
            with pytest.raises(ValueError, match="Model error"):
                estimator.predict([[1, 2, 3]])

            # Verify metrics were still recorded with error info
            mock_client.record_prediction.assert_called()
            call_kwargs = mock_client.record_prediction.call_args[1]
            assert "error" in call_kwargs

    def test_no_double_patching(self, mocker):
        """Test that calling autolog twice doesn't cause double patching."""

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

        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 MCAClient
        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 twice
            patch_sklearn(capture_input=False, capture_output=False)
            patch_sklearn(capture_input=False, capture_output=False)

            # Make prediction
            estimator = MockEstimator()
            result = estimator.predict([[1, 2, 3]])

            # Should only be instrumented once
            assert result == [1, 0, 1]
            # Note: We can't easily verify single instrumentation without
            # inspecting the method wrapper, but no exception is good enough


class TestAutologEntry:
    """Test the main autolog() entry point."""

    def test_autolog_requires_client(self):
        """Test that autolog raises error if MCAClient not initialized."""
        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=None):
            from mca_sdk.autolog import autolog

            with pytest.raises(RuntimeError, match="MCAClient must be initialized"):
                autolog()

    def test_autolog_detects_frameworks(self, mocker):
        """Test that autolog auto-detects installed frameworks."""

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

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

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

        # Mock MCAClient
        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 auto-detect and patch sklearn
            autolog()

            # Test that it works
            estimator = MockBaseEstimator()
            estimator.predict([[1, 2]])
            mock_client.record_prediction.assert_called()

    def test_autolog_exclude_framework(self, mocker):
        """Test that autolog respects exclude parameter."""

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

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

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

        # Mock xgboost
        xgboost_module = MagicMock()

        class MockBooster:
            def predict(self, dmatrix):
                return [0.8]

        xgboost_module.Booster = MockBooster
        sys.modules["xgboost"] = xgboost_module

        # Mock MCAClient
        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

            # Exclude xgboost
            autolog(exclude=["xgboost"])

            # sklearn should work
            estimator = MockBaseEstimator()
            estimator.predict([[1, 2]])
            assert mock_client.record_prediction.call_count >= 1

            # xgboost should NOT be instrumented
            booster = MockBooster()
            booster.predict(None)
            # Call count should not increase (xgboost not instrumented)
            assert mock_client.record_prediction.call_count == 1

    def test_autolog_specific_frameworks(self, mocker):
        """Test that autolog respects frameworks parameter."""

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

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

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

        # Mock MCAClient
        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

            # Request only sklearn
            autolog(frameworks=["sklearn"])

            # Verify it works
            estimator = MockBaseEstimator()
            estimator.predict([[1, 2]])
            mock_client.record_prediction.assert_called_once()
