"""Unit tests for GenAI integration metric naming.

Tests that genai.py helper functions use the correct metric prefix
based on the client's model type (generative vs internal).
"""

import pytest
from datetime import datetime, timedelta
from unittest.mock import Mock, patch
from opentelemetry.sdk.metrics import MeterProvider
from opentelemetry.sdk.metrics.export import InMemoryMetricReader
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.resources import Resource

from mca_sdk import MCAClient
from mca_sdk.integrations.genai import track_llm_completion, track_llm_embedding
from mca_sdk.integrations.litellm_callback import MCALiteLLMCallback


class TestGenAIMetricNaming:
    """Test that GenAI helpers use correct metric prefix."""

    def test_track_llm_completion_uses_genai_prefix_for_generative_models(self):
        """Verify track_llm_completion uses 'genai' prefix for generative models."""
        # Create client with model_type='generative'
        client = MCAClient(
            service_name="test-genai",
            model_id="mdl-genai-001",
            team_name="test",
            model_type="generative",
            llm_provider="openai",
            llm_model="gpt-4",
        )

        # Track what metrics were recorded
        recorded_metrics = []
        original_record_metric = client.record_metric

        def capture_record_metric(name, value, **attrs):
            recorded_metrics.append(name)
            return original_record_metric(name, value, **attrs)

        client.record_metric = capture_record_metric

        # Call track_llm_completion
        track_llm_completion(
            client=client,
            model="gpt-4",
            prompt_tokens=100,
            completion_tokens=50,
            latency=1.5,
            status="success",
        )

        # Verify all metrics use 'genai.' prefix, not 'llm.'
        for metric_name in recorded_metrics:
            assert metric_name.startswith(
                "genai."
            ), f"Metric '{metric_name}' should use 'genai.' prefix for generative models"
            assert not metric_name.startswith(
                "llm."
            ), f"Metric '{metric_name}' should not use 'llm.' prefix"

        # Verify expected metrics were recorded
        assert "genai.requests_total" in recorded_metrics
        assert "genai.tokens_total" in recorded_metrics
        assert "genai.latency_seconds" in recorded_metrics
        assert "genai.cost_dollars" in recorded_metrics  # OpenAI model

        client.shutdown()

    def test_track_llm_completion_uses_model_prefix_for_internal_models(self):
        """Verify track_llm_completion uses 'model' prefix for internal models."""
        # Create client with model_type='regression' (internal category)
        client = MCAClient(
            service_name="test-internal",
            model_id="mdl-internal-001",
            team_name="test",
            model_type="regression",
            model_version="1.0",
        )

        # Track what metrics were recorded
        recorded_metrics = []
        original_record_metric = client.record_metric

        def capture_record_metric(name, value, **attrs):
            recorded_metrics.append(name)
            return original_record_metric(name, value, **attrs)

        client.record_metric = capture_record_metric

        # Call track_llm_completion (even though it's not typical for internal models)
        track_llm_completion(
            client=client,
            model="custom-model",
            prompt_tokens=100,
            completion_tokens=50,
            latency=1.5,
            status="success",
        )

        # Verify all metrics use 'model.' prefix
        for metric_name in recorded_metrics:
            assert metric_name.startswith(
                "model."
            ), f"Metric '{metric_name}' should use 'model.' prefix for internal models"

        client.shutdown()

    def test_track_llm_embedding_uses_correct_prefix(self):
        """Verify track_llm_embedding uses correct prefix based on model type."""
        client = MCAClient(
            service_name="test-genai",
            model_id="mdl-genai-002",
            team_name="test",
            model_type="generative",
            llm_provider="openai",
            llm_model="text-embedding-ada-002",
        )

        # Track what metrics were recorded
        recorded_metrics = []
        original_record_metric = client.record_metric

        def capture_record_metric(name, value, **attrs):
            recorded_metrics.append(name)
            return original_record_metric(name, value, **attrs)

        client.record_metric = capture_record_metric

        # Call track_llm_embedding
        track_llm_embedding(
            client=client,
            model="text-embedding-ada-002",
            tokens=100,
            latency=0.5,
            dimensions=1536,
            status="success",
        )

        # Verify all metrics use 'genai.' prefix
        for metric_name in recorded_metrics:
            assert metric_name.startswith(
                "genai."
            ), f"Metric '{metric_name}' should use 'genai.' prefix for generative models"

        # Verify expected embedding metrics
        assert "genai.embedding_requests_total" in recorded_metrics
        assert "genai.embedding_tokens_total" in recorded_metrics
        assert "genai.embedding_latency_seconds" in recorded_metrics
        assert "genai.embedding_dimensions_count" in recorded_metrics

        client.shutdown()

    def test_track_llm_completion_with_error_uses_correct_prefix(self):
        """Verify error metrics also use correct prefix."""
        client = MCAClient(
            service_name="test-genai",
            model_id="mdl-genai-003",
            team_name="test",
            model_type="generative",
            llm_provider="openai",
            llm_model="gpt-4",
        )

        # Track what metrics were recorded
        recorded_metrics = []
        original_record_metric = client.record_metric

        def capture_record_metric(name, value, **attrs):
            recorded_metrics.append(name)
            return original_record_metric(name, value, **attrs)

        client.record_metric = capture_record_metric

        # Call track_llm_completion with error
        track_llm_completion(
            client=client,
            model="gpt-4",
            prompt_tokens=100,
            completion_tokens=0,
            latency=1.5,
            status="error",
            error="API timeout",
        )

        # Verify error metric uses correct prefix
        assert "genai.errors_total" in recorded_metrics, "Error metric should use correct prefix"
        assert (
            "llm.errors_total" not in recorded_metrics
        ), "Error metric should not use hardcoded 'llm.' prefix"

        client.shutdown()

    def test_metric_prefix_matches_client_configuration(self):
        """Verify metric prefix matches _metric_prefix attribute."""
        # Test generative model
        genai_client = MCAClient(
            service_name="test",
            model_id="mdl-001",
            team_name="test",
            model_type="generative",
            llm_provider="openai",
            llm_model="gpt-4",
        )

        assert (
            genai_client._metric_prefix == "genai"
        ), "Generative models should have 'genai' prefix"

        # Test internal model
        internal_client = MCAClient(
            service_name="test",
            model_id="mdl-002",
            team_name="test",
            model_type="regression",
            model_version="1.0",
        )

        assert (
            internal_client._metric_prefix == "model"
        ), "Internal models should have 'model' prefix"

        genai_client.shutdown()
        internal_client.shutdown()


class TestLiteLLMCallbackMetricNaming:
    """Test MCALiteLLMCallback metric names conform to schema validation (AC #4 Story 2-8, AC #3/#7 Story 2-9).

    This class verifies:
    - (Story 2-8) Every metric name passes MetricNamingConvention strict validation.
    - (Story 2-9) Both genai.py and litellm_callback.py paths emit genai.tokens_total
      with token_type attribute, so a single dashboard query captures traffic from both.
    """

    def test_litellm_callback_metric_names_pass_schema_validation(self):
        """Verify all MCALiteLLMCallback instrument names pass strict MetricNamingConvention."""
        from mca_sdk.validation.schema import MetricNamingConvention

        client = Mock()
        metric_reader = InMemoryMetricReader()
        resource = Resource.create({"service.name": "test-genai-service"})
        meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
        client.meter = meter_provider.get_meter("test-meter")

        tracer_provider = TracerProvider(resource=resource)
        client.tracer = tracer_provider.get_tracer("test-tracer")
        client.logger = Mock()

        try:
            callback = MCALiteLLMCallback(client)
            validator = MetricNamingConvention(
                model_category="internal", model_type="generative", strict=True
            )

            instrument_names = [
                callback._requests_counter.name,
                callback._requests_failed_counter.name,
                callback._tokens_counter.name,
                callback._cost_histogram.name,
                callback._latency_histogram.name,
            ]

            for name in instrument_names:
                validator.validate(name)
        finally:
            meter_provider.shutdown()
            tracer_provider.shutdown()

    def test_litellm_callback_token_counter_unified_structure(self):
        """AC #7: Both genai.py and litellm_callback.py emit genai.tokens_total with token_type.

        Confirms no separate genai.tokens_prompt or genai.tokens_completion
        instrument exists on the callback after Story 2-9 unification.
        """
        client = Mock()
        metric_reader = InMemoryMetricReader()
        resource = Resource.create({"service.name": "test-genai-service"})
        meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
        client.meter = meter_provider.get_meter("test-meter")

        tracer_provider = TracerProvider(resource=resource)
        client.tracer = tracer_provider.get_tracer("test-tracer")
        client.logger = Mock()

        try:
            callback = MCALiteLLMCallback(client)

            # Unified counter exists with correct name
            assert callback._tokens_counter.name == "genai.tokens_total"

            # Separate prompt/completion counters must NOT exist
            assert not hasattr(callback, "_prompt_tokens_counter")
            assert not hasattr(callback, "_completion_tokens_counter")

            # Collect all instrument names from the callback
            all_names = [
                callback._requests_counter.name,
                callback._requests_failed_counter.name,
                callback._tokens_counter.name,
                callback._cost_histogram.name,
                callback._latency_histogram.name,
            ]
            assert "genai.tokens_prompt" not in all_names
            assert "genai.tokens_completion" not in all_names
        finally:
            meter_provider.shutdown()
            tracer_provider.shutdown()

    def test_cross_path_both_emit_token_type_attribute(self):
        """AC #3 Story 2-9: Both genai.py and litellm_callback.py paths write to a shared OTel pipeline.

        Uses a real InMemoryMetricReader to verify genai.tokens_total is exported
        with token_type data points from both instrumentation paths, so a single
        PromQL/MQL query filtering genai.tokens_total{token_type="prompt"} captures both.
        """
        # Shared OTel infrastructure — both paths write to the same pipeline
        metric_reader = InMemoryMetricReader()
        resource = Resource.create({"service.name": "test-cross-path"})
        meter_provider = MeterProvider(resource=resource, metric_readers=[metric_reader])
        tracer_provider = TracerProvider(resource=resource)
        shared_meter = meter_provider.get_meter("test-cross")

        client_cb = Mock()
        client_cb.meter = shared_meter
        client_cb.tracer = tracer_provider.get_tracer("test-cross")
        client_cb.logger = Mock()

        try:
            # --- litellm_callback.py path (real counter, no Mock substitution) ---
            callback = MCALiteLLMCallback(client_cb)

            mock_response = Mock()
            mock_response.usage = Mock(prompt_tokens=50, completion_tokens=25, total_tokens=75)
            mock_response._hidden_params = {}

            start_time = datetime.now()
            end_time = start_time + timedelta(seconds=1)
            with patch("litellm.completion_cost", return_value=None):
                callback.log_success_event(
                    {"model": "gpt-4", "litellm_call_id": "xp-test"},
                    mock_response,
                    start_time,
                    end_time,
                )

            # --- genai.py path (real counter on shared meter) ---
            genai_client = Mock()
            genai_client._metric_prefix = "genai"

            # Proxy record_metric to a real counter on the shared meter so writes enter the pipeline
            _genai_token_counter = shared_meter.create_counter("genai.tokens_total")

            def real_record_metric(name, value, **kw):
                if name == "genai.tokens_total":
                    _genai_token_counter.add(value, attributes=kw)

            genai_client.record_metric = real_record_metric

            track_llm_completion(
                genai_client,
                model="gpt-4",
                prompt_tokens=50,
                completion_tokens=25,
                latency=0.5,
            )

            # Verify via real OTel pipeline — both paths must appear in exported metrics
            metrics_data = metric_reader.get_metrics_data()

            token_type_values = set()
            all_token_points = []
            found_metric = False

            for rm in metrics_data.resource_metrics:
                for sm in rm.scope_metrics:
                    for metric in sm.metrics:
                        if metric.name == "genai.tokens_total":
                            found_metric = True
                            for dp in metric.data.data_points:
                                if "token_type" in dp.attributes:
                                    token_type_values.add(dp.attributes["token_type"])
                                    all_token_points.append(dp.attributes)

            assert found_metric, "genai.tokens_total must be exported by at least one path"
            assert {
                "prompt",
                "completion",
                "total",
            } <= token_type_values, (
                f"token_type values must include prompt, completion, total; got {token_type_values}"
            )
            # litellm path emits {model, provider, status, token_type}; genai path emits
            # {model, token_type} — distinct attribute sets prevent OTel from collapsing them.
            litellm_points = [p for p in all_token_points if "provider" in p]
            genai_points = [p for p in all_token_points if "provider" not in p]
            assert (
                len(litellm_points) == 3
            ), "litellm_callback.py path must emit 3 data points (one per token_type)"
            assert (
                len(genai_points) == 3
            ), "genai.py path must emit 3 data points (one per token_type)"

        finally:
            meter_provider.shutdown()
            tracer_provider.shutdown()
