"""Integration tests for @predict() decorator with MCAClient.

Tests real-world usage scenarios including ML model predictions,
multi-threaded environments, and async functions.
"""

import asyncio
import concurrent.futures
import time
import pytest
from unittest.mock import Mock

from mca_sdk import MCAClient


class MockMLModel:
    """Mock machine learning model for integration testing."""

    def __init__(self, model_name: str = "test_model"):
        self.model_name = model_name
        self.prediction_count = 0

    def predict(self, features: dict) -> dict:
        """Simulate ML model prediction."""
        self.prediction_count += 1
        # Simulate some processing time
        time.sleep(0.01)

        # Simple mock prediction logic
        score = sum(features.values()) * 0.1
        return {
            "prediction": "positive" if score > 0.5 else "negative",
            "confidence": min(score, 1.0),
            "model": self.model_name,
        }


class TestDecoratorWithMLModel:
    """Integration tests with ML model-like functions."""

    def test_decorator_with_ml_model_prediction(self):
        """Test decorator works with ML model prediction method."""
        client = MCAClient(
            service_name="ml-service", model_id="test-ml-model", team_name="test-team"
        )
        model = MockMLModel("classifier-v1")

        # Decorate the model's predict method
        decorated_predict = client.predict(
            span_name="ml_model_prediction", capture_input=True, capture_output=True
        )(model.predict)

        # Make predictions
        features = {"feature1": 2.5, "feature2": 3.8, "feature3": 1.2}
        result = decorated_predict(features)

        # Verify prediction executed successfully
        assert result["prediction"] in ["positive", "negative"]
        assert 0.0 <= result["confidence"] <= 1.0
        assert result["model"] == "classifier-v1"
        assert model.prediction_count == 1

    def test_decorator_with_multiple_predictions(self):
        """Test decorator tracks multiple predictions correctly."""
        client = MCAClient(
            service_name="ml-service", model_id="test-ml-model", team_name="test-team"
        )
        model = MockMLModel()

        @client.predict(capture_input=True, capture_output=True)
        def predict_batch(batch_features: list) -> list:
            """Process batch of predictions."""
            return [model.predict(features) for features in batch_features]

        # Make batch prediction
        batch = [
            {"feature1": 1.0, "feature2": 2.0},
            {"feature1": 5.0, "feature2": 6.0},
            {"feature1": 0.5, "feature2": 0.3},
        ]
        results = predict_batch(batch)

        # Verify all predictions completed
        assert len(results) == 3
        assert all("prediction" in r for r in results)


class TestDecoratorMultiThreaded:
    """Integration tests for multi-threaded scenarios."""

    def test_decorator_in_threadpool_executor(self):
        """Test decorator works correctly with ThreadPoolExecutor."""
        client = MCAClient(
            service_name="threaded-service",
            model_id="thread-model",
            team_name="test-team",
        )

        results = []
        prediction_ids_captured = []

        @client.predict(capture_input=True, capture_output=True)
        def threaded_predict(value: int) -> dict:
            """Function to be called from multiple threads."""
            time.sleep(0.001)  # Small delay to simulate work
            return {
                "input": value,
                "output": value * 2,
                "thread_id": (id(threading.current_thread()) if "threading" in dir() else 0),
            }

        # Execute predictions in thread pool
        with concurrent.futures.ThreadPoolExecutor(max_workers=10) as executor:
            futures = [executor.submit(threaded_predict, i) for i in range(50)]
            results = [f.result() for f in concurrent.futures.as_completed(futures)]

        # Verify all predictions completed
        assert len(results) == 50

        # Verify each input was processed correctly
        outputs = {r["input"]: r["output"] for r in results}
        for i in range(50):
            assert outputs[i] == i * 2

    def test_no_race_conditions_concurrent_predictions(self):
        """Test decorator has no race conditions with concurrent calls."""
        client = MCAClient(
            service_name="concurrent-service",
            model_id="concurrent-model",
            team_name="test-team",
        )

        shared_state = {"counter": 0}

        @client.predict()
        def increment_and_return(increment_by: int) -> int:
            """Function that modifies shared state (intentional race condition test)."""
            current = shared_state["counter"]
            time.sleep(0.0001)  # Small delay to expose race conditions
            shared_state["counter"] = current + increment_by
            return shared_state["counter"]

        # Run concurrent predictions
        with concurrent.futures.ThreadPoolExecutor(max_workers=5) as executor:
            futures = [executor.submit(increment_and_return, 1) for _ in range(20)]
            results = [f.result() for f in futures]

        # Decorator should not interfere with execution
        # (Note: shared_state will have race conditions, but decorator shouldn't add to them)
        assert len(results) == 20


class TestDecoratorAsync:
    """Integration tests for async functions."""

    @pytest.mark.asyncio
    async def test_decorator_with_async_ml_prediction(self):
        """Test decorator works with async ML prediction functions."""
        client = MCAClient(
            service_name="async-ml-service",
            model_id="async-model",
            team_name="test-team",
        )

        @client.predict(capture_input=True, capture_output=True)
        async def async_predict(features: dict) -> dict:
            """Async ML prediction function."""
            await asyncio.sleep(0.01)  # Simulate async I/O
            score = sum(features.values()) * 0.1
            return {
                "prediction": "positive" if score > 0.5 else "negative",
                "confidence": score,
            }

        # Make async prediction
        features = {"feature1": 5.0, "feature2": 8.0}
        result = await async_predict(features)

        # Verify prediction succeeded
        assert result["prediction"] == "positive"
        assert result["confidence"] > 0.5

    @pytest.mark.asyncio
    async def test_decorator_with_concurrent_async_predictions(self):
        """Test decorator with multiple concurrent async predictions."""
        client = MCAClient(
            service_name="async-concurrent-service",
            model_id="async-concurrent-model",
            team_name="test-team",
        )

        @client.predict(capture_input=True, capture_output=True)
        async def async_batch_predict(value: int) -> dict:
            """Async prediction for concurrent testing."""
            await asyncio.sleep(0.001)
            return {"input": value, "output": value**2}

        # Run 100 concurrent async predictions
        tasks = [async_batch_predict(i) for i in range(100)]
        results = await asyncio.gather(*tasks)

        # Verify all completed correctly
        assert len(results) == 100
        assert results[10]["output"] == 100  # 10^2
        assert results[99]["output"] == 9801  # 99^2


class TestDecoratorErrorHandling:
    """Integration tests for error handling."""

    def test_decorator_handles_prediction_errors(self):
        """Test decorator properly handles and re-raises prediction errors."""
        client = MCAClient(
            service_name="error-service", model_id="error-model", team_name="test-team"
        )

        @client.predict(capture_input=True)
        def failing_predict(data: dict) -> dict:
            """Prediction that fails."""
            if data.get("trigger_error"):
                raise ValueError("Invalid input data")
            return {"result": "ok"}

        # Successful prediction
        result = failing_predict({"trigger_error": False})
        assert result["result"] == "ok"

        # Failed prediction should re-raise
        with pytest.raises(ValueError, match="Invalid input data"):
            failing_predict({"trigger_error": True})

    def test_decorator_records_error_and_latency_on_failure(self):
        """Test decorator records metrics even when prediction fails."""
        client = MCAClient(
            service_name="error-metrics-service",
            model_id="error-metrics-model",
            team_name="test-team",
        )

        @client.predict()
        def slow_failing_predict():
            """Prediction that fails after some time."""
            time.sleep(0.05)  # 50ms delay
            raise RuntimeError("Model execution failed")

        # Capture exception
        start = time.perf_counter()
        with pytest.raises(RuntimeError):
            slow_failing_predict()
        elapsed = time.perf_counter() - start

        # Verify execution took expected time (decorator didn't short-circuit)
        assert elapsed >= 0.05


class TestDecoratorWithComplexInputs:
    """Integration tests with complex input types."""

    def test_decorator_with_numpy_like_arrays(self):
        """Test decorator handles array-like inputs."""
        client = MCAClient(
            service_name="array-service", model_id="array-model", team_name="test-team"
        )

        @client.predict(capture_input=True, capture_output=True)
        def predict_with_arrays(features: list, weights: list) -> float:
            """Prediction using list inputs (simulating numpy arrays)."""
            return sum(f * w for f, w in zip(features, weights))

        features = [1.0, 2.0, 3.0, 4.0]
        weights = [0.1, 0.2, 0.3, 0.4]

        result = predict_with_arrays(features, weights)

        # Verify prediction calculated correctly
        assert result == 3.0  # 1*0.1 + 2*0.2 + 3*0.3 + 4*0.4

    def test_decorator_with_multiple_args_and_kwargs(self):
        """Test decorator captures multiple args and kwargs correctly."""
        client = MCAClient(
            service_name="multi-arg-service",
            model_id="multi-arg-model",
            team_name="test-team",
        )

        @client.predict(capture_input=True, capture_output=True)
        def predict_multi_param(
            model_input: dict, threshold: float = 0.5, normalize: bool = True
        ) -> dict:
            """Prediction with multiple parameters."""
            score = sum(model_input.values())
            if normalize:
                score = score / len(model_input)

            return {
                "prediction": "positive" if score > threshold else "negative",
                "score": score,
                "threshold": threshold,
            }

        # Call with positional and keyword arguments
        result = predict_multi_param({"feat1": 0.8, "feat2": 0.9}, threshold=0.7, normalize=True)

        # Verify execution
        expected_score = (0.8 + 0.9) / 2  # 0.85
        assert result["score"] == expected_score
        assert result["threshold"] == 0.7
        assert expected_score > 0.7, "Score should be greater than threshold"
        assert (
            result["prediction"] == "positive"
        ), f"Expected 'positive' but got '{result['prediction']}' with score {result['score']} and threshold {result['threshold']}"


# Import threading for thread ID capture
import threading
