"""Unit tests for XGBoost autolog functionality."""

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


class TestXGBoostAutolog:
    """Test autolog for XGBoost framework."""

    def test_patch_xgboost_booster(self, mocker):
        """Test that XGBoost Booster.predict() is instrumented."""
        # Mock xgboost import
        xgboost_module = MagicMock()

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

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

        # 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._xgboost import patch_xgboost

            # Apply patch
            patch_xgboost(capture_input=False, capture_output=False)

            # Make prediction
            booster = MockBooster()
            result = booster.predict(None)

            # Verify instrumentation was called
            assert result == [0.8, 0.2, 0.9]
            mock_client.trace.assert_called()
            mock_client.record_prediction.assert_called()

    def test_patch_xgboost_sklearn_api(self, mocker):
        """Test that XGBoost sklearn API (XGBClassifier) is instrumented."""
        # Mock xgboost import
        xgboost_module = MagicMock()
        xgboost_sklearn = MagicMock()

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

        xgboost_sklearn.XGBClassifier = MockXGBClassifier
        xgboost_sklearn.XGBRegressor = MagicMock
        xgboost_module.sklearn = xgboost_sklearn
        xgboost_module.Booster = MagicMock
        sys.modules["xgboost"] = xgboost_module
        sys.modules["xgboost.sklearn"] = xgboost_sklearn

        # 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._xgboost import patch_xgboost

            # Apply patch
            patch_xgboost(capture_input=False, capture_output=False)

            # Make prediction
            classifier = MockXGBClassifier()
            result = classifier.predict([[1, 2, 3]])

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

            # Verify attributes include sklearn API type
            call_kwargs = mock_client.record_prediction.call_args[1]
            assert call_kwargs.get("api_type") == "sklearn"
            assert call_kwargs.get("model_type") == "classifier"

    def test_xgboost_graceful_fallback_no_client(self, mocker):
        """Test that xgboost autolog gracefully falls back if client not initialized."""
        # Mock xgboost
        xgboost_module = MagicMock()

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

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

        # Mock MCAClient.get_instance() to return None
        with patch("mca_sdk.core.client.MCAClient.get_instance", return_value=None):
            from mca_sdk.autolog._xgboost import patch_xgboost

            # Patch xgboost
            patch_xgboost(capture_input=False, capture_output=False)

            # Make prediction (should not raise error)
            booster = MockBooster()
            result = booster.predict(None)

            # Should still work without instrumentation
            assert result == [0.8, 0.2, 0.9]

    def test_xgboost_exception_handling(self, mocker):
        """Test that xgboost autolog handles exceptions and re-raises them."""
        # Mock xgboost
        xgboost_module = MagicMock()

        class MockBooster:
            def predict(self, dmatrix):
                raise RuntimeError("Prediction failed")

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

        # 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._xgboost import patch_xgboost

            # Patch xgboost
            patch_xgboost(capture_input=False, capture_output=False)

            # Make prediction (should raise error AND record metrics)
            booster = MockBooster()
            with pytest.raises(RuntimeError, match="Prediction failed"):
                booster.predict(None)

            # 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
            assert "error_type" in call_kwargs

    def test_xgboost_multiple_methods_patched(self, mocker):
        """Test that multiple XGBoost methods are patched correctly."""
        # Mock xgboost with both Booster and sklearn API
        xgboost_module = MagicMock()
        xgboost_sklearn = MagicMock()

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

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

            def predict_proba(self, X):
                return [[0.2, 0.8]]

        xgboost_module.Booster = MockBooster
        xgboost_sklearn.XGBClassifier = MockXGBClassifier
        xgboost_sklearn.XGBRegressor = MagicMock
        xgboost_module.sklearn = xgboost_sklearn
        sys.modules["xgboost"] = xgboost_module
        sys.modules["xgboost.sklearn"] = xgboost_sklearn

        # Mock MCAClient
        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._xgboost import patch_xgboost

            # Apply patch
            patch_xgboost(capture_input=False, capture_output=False)

            # Test Booster.predict
            booster = MockBooster()
            booster.predict(None)
            assert mock_client.record_prediction.call_count >= 1

            # Test XGBClassifier.predict
            classifier = MockXGBClassifier()
            classifier.predict([[1, 2]])
            assert mock_client.record_prediction.call_count >= 2

            # Test XGBClassifier.predict_proba
            classifier.predict_proba([[1, 2]])
            assert mock_client.record_prediction.call_count >= 3


class TestAutologIntegration:
    """Integration tests for autolog with multiple frameworks."""

    def test_autolog_mixed_frameworks(self, mocker):
        """Test that autolog can handle multiple frameworks simultaneously."""

        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

            autolog(frameworks=["sklearn", "xgboost"])

            # Test both frameworks work
            estimator = MockBaseEstimator()
            estimator.predict([[1, 2]])
            assert mock_client.record_prediction.call_count >= 1

            booster = MockBooster()
            booster.predict(None)
            assert mock_client.record_prediction.call_count >= 2
