"""Integration tests for autolog with real ML framework patching.

Tests that autolog correctly patches real scikit-learn, XGBoost, and LightGBM
models and verifies telemetry emission using in-memory OTel exporters.

These tests require ML frameworks to be installed:
    pip install scikit-learn xgboost lightgbm
"""

import logging
import numpy as np
import pytest
from unittest.mock import Mock

from mca_sdk import MCAClient
from mca_sdk.autolog import autolog


class TestAutologSklearnIntegration:
    """Integration tests for sklearn autolog with real models."""

    def test_sklearn_random_forest_patching(
        self, mca_client_factory, in_memory_span_exporter, in_memory_metric_reader
    ):
        """Test autolog patches real sklearn RandomForestClassifier and emits telemetry."""
        # Import sklearn
        from sklearn.ensemble import RandomForestClassifier
        from sklearn.datasets import make_classification

        # Create MCAClient
        client = mca_client_factory(
            service_name="sklearn-test", model_id="rf-001", team_name="test-team"
        )

        # Enable autolog
        autolog(frameworks=["sklearn"], capture_input=True, capture_output=True)

        # Create and train a real model
        X, y = make_classification(n_samples=100, n_features=10, random_state=42)
        model = RandomForestClassifier(n_estimators=10, random_state=42)
        model.fit(X, y)

        # Make predictions (should be instrumented)
        predictions = model.predict(X[:5])

        # Verify predictions worked
        assert predictions is not None
        assert len(predictions) == 5

        # Force flush to ensure spans are exported
        client._tracer_provider.force_flush()
        client._meter_provider.force_flush()

        # Verify spans were emitted
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) > 0, "Expected spans to be emitted for sklearn predict"

        # Find the predict span
        predict_spans = [s for s in spans if "predict" in s.name.lower()]
        assert len(predict_spans) > 0, "Expected at least one predict span"

        # Verify span attributes contain framework info (proper attribute checks)
        predict_span = predict_spans[0]
        assert predict_span.attributes is not None
        # Check for specific framework attribute key/value
        assert "framework" in predict_span.attributes
        assert predict_span.attributes["framework"] == "sklearn"
        assert "method" in predict_span.attributes
        assert predict_span.attributes["method"] == "predict"

        # Verify metrics were recorded
        metrics = in_memory_metric_reader.get_metrics_data()
        assert metrics is not None

    def test_sklearn_predict_proba_instrumentation(
        self, mca_client_factory, in_memory_span_exporter
    ):
        """Test autolog instruments predict_proba method."""
        from sklearn.linear_model import LogisticRegression
        from sklearn.datasets import make_classification

        # Create client and enable autolog
        client = mca_client_factory(
            service_name="sklearn-proba-test",
            model_id="lr-001",
            team_name="test-team",
        )
        autolog(frameworks=["sklearn"], capture_output=True)

        # Create and train model
        X, y = make_classification(n_samples=50, n_features=5, random_state=42)
        model = LogisticRegression(random_state=42)
        model.fit(X, y)

        # Call predict_proba
        probabilities = model.predict_proba(X[:3])

        # Verify predictions
        assert probabilities is not None
        assert probabilities.shape == (3, 2)

        # Force flush and verify spans
        client._tracer_provider.force_flush()
        spans = in_memory_span_exporter.get_finished_spans()

        # Should have spans for both fit and predict_proba
        assert len(spans) >= 1
        span_names = [s.name for s in spans]
        assert any("predict" in name.lower() for name in span_names)


class TestAutologXGBoostIntegration:
    """Integration tests for XGBoost autolog with real models."""

    def test_xgboost_classifier_patching(
        self, mca_client_factory, in_memory_span_exporter, in_memory_metric_reader
    ):
        """Test autolog patches real XGBoost XGBClassifier and emits telemetry."""
        # Import xgboost
        import xgboost as xgb
        from sklearn.datasets import make_classification

        # Create MCAClient
        client = mca_client_factory(
            service_name="xgboost-test", model_id="xgb-001", team_name="test-team"
        )

        # Enable autolog
        autolog(frameworks=["xgboost"], capture_input=True, capture_output=True)

        # Create and train a real XGBoost model
        X, y = make_classification(n_samples=100, n_features=10, random_state=42)
        model = xgb.XGBClassifier(n_estimators=10, random_state=42, verbosity=0)
        model.fit(X, y)

        # Make predictions
        predictions = model.predict(X[:5])

        # Verify predictions worked
        assert predictions is not None
        assert len(predictions) == 5

        # Force flush
        client._tracer_provider.force_flush()
        client._meter_provider.force_flush()

        # Verify spans were emitted
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) > 0, "Expected spans to be emitted for xgboost predict"

        # Find the predict span
        predict_spans = [s for s in spans if "predict" in s.name.lower()]
        assert len(predict_spans) > 0, "Expected at least one predict span"

        # Verify span attributes
        predict_span = predict_spans[0]
        assert predict_span.attributes is not None

    def test_xgboost_booster_api(self, mca_client_factory, in_memory_span_exporter):
        """Test autolog works with native XGBoost Booster API."""
        import xgboost as xgb
        from sklearn.datasets import make_classification

        # Create client and enable autolog
        client = mca_client_factory(
            service_name="xgb-booster-test", model_id="xgb-002", team_name="test-team"
        )
        autolog(frameworks=["xgboost"])

        # Create DMatrix and train with native API
        X, y = make_classification(n_samples=50, n_features=5, random_state=42)
        dtrain = xgb.DMatrix(X, label=y)
        params = {"objective": "binary:logistic", "eval_metric": "logloss"}
        bst = xgb.train(params, dtrain, num_boost_round=5, verbose_eval=False)

        # Make predictions using Booster.predict
        dtest = xgb.DMatrix(X[:3])
        predictions = bst.predict(dtest)

        # Verify predictions
        assert predictions is not None
        assert len(predictions) == 3

        # Force flush and verify telemetry
        client._tracer_provider.force_flush()
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) >= 1


class TestAutologLightGBMIntegration:
    """Integration tests for LightGBM autolog with real models."""

    def test_lightgbm_classifier_patching(
        self, mca_client_factory, in_memory_span_exporter, in_memory_metric_reader
    ):
        """Test autolog patches real LightGBM LGBMClassifier and emits telemetry."""
        # Import lightgbm
        import lightgbm as lgb
        from sklearn.datasets import make_classification

        # Create MCAClient
        client = mca_client_factory(
            service_name="lightgbm-test", model_id="lgbm-001", team_name="test-team"
        )

        # Enable autolog
        autolog(frameworks=["lightgbm"], capture_input=True, capture_output=True)

        # Create and train a real LightGBM model
        X, y = make_classification(n_samples=100, n_features=10, random_state=42)
        model = lgb.LGBMClassifier(n_estimators=10, random_state=42, verbosity=-1)
        model.fit(X, y)

        # Make predictions
        predictions = model.predict(X[:5])

        # Verify predictions worked
        assert predictions is not None
        assert len(predictions) == 5

        # Force flush
        client._tracer_provider.force_flush()
        client._meter_provider.force_flush()

        # Verify spans were emitted
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) > 0, "Expected spans to be emitted for lightgbm predict"

        # Find the predict span
        predict_spans = [s for s in spans if "predict" in s.name.lower()]
        assert len(predict_spans) > 0, "Expected at least one predict span"

    def test_lightgbm_booster_api(self, mca_client_factory, in_memory_span_exporter):
        """Test autolog works with native LightGBM Booster API."""
        import lightgbm as lgb
        from sklearn.datasets import make_classification

        # Create client and enable autolog
        client = mca_client_factory(
            service_name="lgbm-booster-test",
            model_id="lgbm-002",
            team_name="test-team",
        )
        autolog(frameworks=["lightgbm"])

        # Create Dataset and train with native API
        X, y = make_classification(n_samples=50, n_features=5, random_state=42)
        train_data = lgb.Dataset(X, label=y)
        params = {"objective": "binary", "metric": "binary_logloss", "verbosity": -1}
        bst = lgb.train(params, train_data, num_boost_round=5, verbose_eval=False)

        # Make predictions using Booster.predict
        predictions = bst.predict(X[:3])

        # Verify predictions
        assert predictions is not None
        assert len(predictions) == 3

        # Force flush and verify telemetry
        client._tracer_provider.force_flush()
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) >= 1


class TestAutologExcludeParameter:
    """Integration tests for exclude parameter."""

    def test_exclude_xgboost(
        self, mca_client_factory, in_memory_span_exporter, in_memory_metric_reader
    ):
        """Test that excluded frameworks are not patched."""
        from sklearn.ensemble import RandomForestClassifier
        from sklearn.datasets import make_classification

        # Create client
        client = mca_client_factory(
            service_name="exclude-test", model_id="excl-001", team_name="test-team"
        )

        # Enable autolog but exclude xgboost
        autolog(exclude=["xgboost"])

        # Verify sklearn is patched (should work)
        X, y = make_classification(n_samples=50, n_features=5, random_state=42)
        sklearn_model = RandomForestClassifier(n_estimators=5, random_state=42)
        sklearn_model.fit(X, y)
        predictions = sklearn_model.predict(X[:3])
        assert predictions is not None

        # Force flush
        client._tracer_provider.force_flush()

        # Verify spans were emitted for sklearn
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) > 0

    def test_exclude_all_frameworks(self, mca_client_factory, caplog):
        """Test that excluding all frameworks logs a warning."""
        # Create client
        client = mca_client_factory(
            service_name="exclude-all-test",
            model_id="excl-002",
            team_name="test-team",
        )

        # Enable autolog but exclude all frameworks
        with caplog.at_level(logging.WARNING):
            autolog(exclude=["sklearn", "xgboost", "lightgbm"])

        # Verify warning was logged
        assert any("No ML frameworks detected" in record.message for record in caplog.records)


class TestAutologGracefulFallback:
    """Integration tests for graceful fallback behavior."""

    def test_autolog_without_initialized_client(self):
        """Test that autolog raises RuntimeError if MCAClient not initialized."""
        # Clear any existing client instance
        MCAClient._instance = None

        # Attempt to call autolog without initializing client
        with pytest.raises(
            RuntimeError, match="MCAClient must be initialized before calling autolog"
        ):
            autolog()


class TestAutologDataCapture:
    """Integration tests for input/output data capture."""

    def test_capture_input_output_with_shapes(self, mca_client_factory, in_memory_span_exporter):
        """Test that capture_input/capture_output records data shapes."""
        from sklearn.ensemble import RandomForestClassifier
        from sklearn.datasets import make_classification

        # Create client
        client = mca_client_factory(
            service_name="capture-test", model_id="cap-001", team_name="test-team"
        )

        # Enable autolog with data capture
        autolog(frameworks=["sklearn"], capture_input=True, capture_output=True)

        # Create model and make predictions
        X, y = make_classification(n_samples=50, n_features=8, random_state=42)
        model = RandomForestClassifier(n_estimators=5, random_state=42)
        model.fit(X, y)
        predictions = model.predict(X[:10])

        # Force flush
        client._tracer_provider.force_flush()

        # Verify spans contain shape information
        spans = in_memory_span_exporter.get_finished_spans()
        predict_spans = [s for s in spans if "predict" in s.name.lower()]

        # At least one predict span should exist
        assert len(predict_spans) > 0

        # Check if any span has input/output shape attributes
        # Note: Actual attribute names depend on implementation
        predict_span = predict_spans[0]
        assert predict_span.attributes is not None

    def test_payload_size_limit_truncation(self, mca_client_factory, in_memory_span_exporter):
        """Test that large payloads are truncated per max_payload_size."""
        from sklearn.ensemble import RandomForestClassifier
        from sklearn.datasets import make_classification

        # Create client
        client = mca_client_factory(
            service_name="truncate-test", model_id="trunc-001", team_name="test-team"
        )

        # Enable autolog with small payload limit
        autolog(
            frameworks=["sklearn"],
            capture_input=True,
            capture_output=True,
            max_payload_size=100,  # Very small limit
        )

        # Create model with large feature space
        X, y = make_classification(n_samples=1000, n_features=50, random_state=42)  # Large dataset
        model = RandomForestClassifier(n_estimators=5, random_state=42)
        model.fit(X, y)
        predictions = model.predict(X[:100])  # Predict on 100 samples

        # Verify predictions still work despite truncation
        assert predictions is not None
        assert len(predictions) == 100

        # Force flush
        client._tracer_provider.force_flush()

        # Verify spans were still emitted (truncation shouldn't break instrumentation)
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) > 0


class TestAutologMultipleFrameworks:
    """Integration tests for using multiple frameworks simultaneously."""

    def test_autolog_multiple_frameworks_simultaneously(
        self, mca_client_factory, in_memory_span_exporter
    ):
        """Test that autolog can patch multiple frameworks at once."""
        from sklearn.ensemble import RandomForestClassifier
        import xgboost as xgb
        from sklearn.datasets import make_classification

        # Create client
        client = mca_client_factory(
            service_name="multi-framework-test",
            model_id="multi-001",
            team_name="test-team",
        )

        # Enable autolog for multiple frameworks
        autolog(frameworks=["sklearn", "xgboost"])

        # Train and predict with sklearn
        X, y = make_classification(n_samples=50, n_features=5, random_state=42)
        sklearn_model = RandomForestClassifier(n_estimators=3, random_state=42)
        sklearn_model.fit(X, y)
        sklearn_preds = sklearn_model.predict(X[:3])

        # Train and predict with xgboost
        xgb_model = xgb.XGBClassifier(n_estimators=3, random_state=42, verbosity=0)
        xgb_model.fit(X, y)
        xgb_preds = xgb_model.predict(X[:3])

        # Verify both predictions worked
        assert sklearn_preds is not None
        assert xgb_preds is not None

        # Force flush
        client._tracer_provider.force_flush()

        # Verify spans from both frameworks
        spans = in_memory_span_exporter.get_finished_spans()
        assert len(spans) > 0

        # Should have spans from both frameworks
        predict_spans = [s for s in spans if "predict" in s.name.lower()]
        assert len(predict_spans) >= 2  # At least one from each framework
