"""Unit tests for integrations/genai.py module.

Tests GenAI/LLM integration helpers including client creation, cost estimation,
completion tracking, embedding tracking, and token rate calculations.
"""

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

from mca_sdk.integrations.genai import (
    create_genai_client,
    estimate_openai_cost,
    track_llm_completion,
    track_llm_embedding,
    calculate_token_rate,
    OPENAI_PRICING,
)
from mca_sdk.core.client import MCAClient


def _make_completion_instruments():
    """Return a dict of mock OTel instruments for track_llm_completion."""
    return {
        "requests_total": MagicMock(),
        "requests_failed_total": MagicMock(),
        "tokens_total": MagicMock(),
        "latency_seconds": MagicMock(),
        "cost_dollars": MagicMock(),
    }


def _make_embedding_instruments():
    """Return a dict of mock OTel instruments for track_llm_embedding."""
    return {
        "embedding_requests_total": MagicMock(),
        "embedding_requests_failed_total": MagicMock(),
        "embedding_tokens_total": MagicMock(),
        "embedding_latency_seconds": MagicMock(),
        "embedding_dimensions_count": MagicMock(),
    }


class TestCreateGenAIClient:
    """Test create_genai_client() function."""

    @patch("mca_sdk.integrations.genai.MCAClient")
    def test_create_with_required_params(self, mock_client_class):
        """Test creating GenAI client with required parameters."""
        mock_client = Mock()
        mock_client_class.return_value = mock_client

        result = create_genai_client(
            service_name="test-service",
            model_id="test-model-id",
            llm_provider="openai",
            llm_model="gpt-4",
        )

        # Should call MCAClient with correct params including model_id and taxonomy defaults
        mock_client_class.assert_called_once_with(
            service_name="test-service",
            model_id="test-model-id",
            team_name=None,
            model_category="vendor",
            model_type="generative",
            llm_provider="openai",
            llm_model="gpt-4",
        )
        assert result == mock_client

    @patch("mca_sdk.integrations.genai.MCAClient")
    def test_create_with_team_name(self, mock_client_class):
        """Test creating GenAI client with team name."""
        mock_client = Mock()
        mock_client_class.return_value = mock_client

        result = create_genai_client(
            service_name="test-service",
            model_id="test-model-id",
            llm_provider="openai",
            llm_model="gpt-4",
            team_name="ai-team",
        )

        # Should include team_name
        mock_client_class.assert_called_once()
        call_kwargs = mock_client_class.call_args[1]
        assert call_kwargs["team_name"] == "ai-team"

    @patch("mca_sdk.integrations.genai.MCAClient")
    def test_create_with_additional_kwargs(self, mock_client_class):
        """Test that additional kwargs are passed through."""
        mock_client = Mock()
        mock_client_class.return_value = mock_client

        result = create_genai_client(
            service_name="test-service",
            model_id="test-model-id",
            llm_provider="openai",
            llm_model="gpt-4",
            team_name="ai-team",
            otel_endpoint="http://custom:4318",
            debug_mode=True,
        )

        # Should pass through additional kwargs
        call_kwargs = mock_client_class.call_args[1]
        assert call_kwargs["otel_endpoint"] == "http://custom:4318"
        assert call_kwargs["debug_mode"] is True

    @patch("mca_sdk.integrations.genai.MCAClient")
    def test_create_with_explicit_phi_capture(self, mock_client_class):
        """Test explicit opt-in for PHI capture (HIPAA compliance)."""
        mock_client = Mock()
        mock_client_class.return_value = mock_client

        result = create_genai_client(
            service_name="test-service",
            model_id="test-model-id",
            llm_provider="openai",
            llm_model="gpt-4",
            capture_input_data=True,  # Explicit opt-in via kwargs
            capture_output_data=True,  # Explicit opt-in via kwargs
        )

        # Should pass explicit PHI capture flags through to MCAClient
        call_kwargs = mock_client_class.call_args[1]
        assert call_kwargs["capture_input_data"] is True
        assert call_kwargs["capture_output_data"] is True

    @patch("mca_sdk.integrations.genai.MCAClient")
    def test_create_defaults_to_capture(self, mock_client_class):
        """Test default behavior captures data (capture_phi=True)."""
        mock_client = MagicMock()
        mock_client_class.return_value = mock_client

        create_genai_client(
            service_name="test-service",
            model_id="test-model-id",
            llm_provider="openai",
            llm_model="gpt-4",
        )

        # By default, capture_phi=True is attached to the returned client
        assert mock_client._capture_phi is True

    @patch("mca_sdk.integrations.genai.MCAClient")
    def test_create_always_sets_model_type_to_generative(self, mock_client_class):
        """Test that model_type is always set to generative."""
        mock_client = Mock()
        mock_client_class.return_value = mock_client

        result = create_genai_client(
            service_name="test-service",
            model_id="test-model-id",
            llm_provider="anthropic",
            llm_model="claude-2",
        )

        call_kwargs = mock_client_class.call_args[1]
        assert call_kwargs["model_type"] == "generative"


class TestEstimateOpenAICost:
    """Test estimate_openai_cost() function."""

    def test_cost_calculation_gpt4(self):
        """Test cost calculation for GPT-4."""
        cost = estimate_openai_cost(model="gpt-4", prompt_tokens=1000, completion_tokens=500)

        # 1000/1000 * 0.03 + 500/1000 * 0.06 = 0.03 + 0.03 = 0.06
        assert cost == pytest.approx(0.06, rel=1e-6)

    def test_cost_calculation_gpt35_turbo(self):
        """Test cost calculation for GPT-3.5-turbo."""
        cost = estimate_openai_cost(
            model="gpt-3.5-turbo", prompt_tokens=1000, completion_tokens=500
        )

        # 1000/1000 * 0.0015 + 500/1000 * 0.002 = 0.0015 + 0.001 = 0.0025
        assert cost == pytest.approx(0.0025, rel=1e-6)

    def test_cost_calculation_gpt4_32k(self):
        """Test cost calculation for GPT-4-32k."""
        cost = estimate_openai_cost(model="gpt-4-32k", prompt_tokens=2000, completion_tokens=1000)

        # 2000/1000 * 0.06 + 1000/1000 * 0.12 = 0.12 + 0.12 = 0.24
        assert cost == pytest.approx(0.24, rel=1e-6)

    def test_cost_calculation_with_zero_tokens(self):
        """Test cost calculation with zero tokens."""
        cost = estimate_openai_cost(model="gpt-4", prompt_tokens=0, completion_tokens=0)

        assert cost == 0.0

    def test_unknown_model_defaults_to_gpt35_pricing(self):
        """Test that unknown models default to GPT-3.5-turbo pricing."""
        cost = estimate_openai_cost(
            model="unknown-model", prompt_tokens=1000, completion_tokens=500
        )

        # Should use gpt-3.5-turbo pricing
        expected = (1000 / 1000) * 0.0015 + (500 / 1000) * 0.002
        assert cost == pytest.approx(expected, rel=1e-6)

    def test_fractional_tokens(self):
        """Test cost calculation with fractional token counts."""
        cost = estimate_openai_cost(model="gpt-4", prompt_tokens=150, completion_tokens=75)

        # 150/1000 * 0.03 + 75/1000 * 0.06 = 0.0045 + 0.0045 = 0.009
        assert cost == pytest.approx(0.009, rel=1e-6)


class TestTrackLLMCompletion:
    """Test track_llm_completion() function."""

    def test_track_successful_completion(self):
        """Test tracking a successful LLM completion."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"

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

        # requests_total recorded with value 1
        mock_client.record_metric.assert_any_call(
            "genai.requests_total", 1, model="gpt-4", provider="openai", status="success"
        )

    def test_track_completion_records_token_metrics(self):
        """Test that all token metrics are recorded."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"

        track_llm_completion(
            client=mock_client,
            model="gpt-4",
            prompt_tokens=100,
            completion_tokens=50,
            latency=1.5,
        )

        # tokens_total recorded 3 times: total=150, prompt=100, completion=50
        token_calls = [
            c for c in mock_client.record_metric.call_args_list if c[0][0] == "genai.tokens_total"
        ]
        assert len(token_calls) == 3
        recorded_values = [c[0][1] for c in token_calls]
        assert 150 in recorded_values
        assert 100 in recorded_values
        assert 50 in recorded_values

    def test_track_completion_records_latency(self):
        """Test that latency is recorded."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"

        track_llm_completion(
            client=mock_client,
            model="gpt-4",
            prompt_tokens=100,
            completion_tokens=50,
            latency=2.5,
        )

        mock_client.record_metric.assert_any_call(
            "genai.latency_seconds", 2.5, model="gpt-4", provider="openai"
        )

    def test_track_completion_records_cost_for_openai_models(self):
        """Test that cost is recorded for OpenAI models."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"

        track_llm_completion(
            client=mock_client,
            model="gpt-4",
            prompt_tokens=1000,
            completion_tokens=500,
            latency=1.5,
        )

        cost_calls = [
            c for c in mock_client.record_metric.call_args_list if c[0][0] == "genai.cost_dollars"
        ]
        assert len(cost_calls) == 1
        assert cost_calls[0][0][1] == pytest.approx(0.06, rel=1e-6)

    @patch("mca_sdk.integrations.genai._get_instruments")
    def test_track_completion_no_cost_for_unknown_models(self, mock_get_instr):
        """Test that no cost is recorded for unknown models."""
        instruments = _make_completion_instruments()
        mock_get_instr.return_value = instruments
        mock_client = MagicMock()

        track_llm_completion(
            client=mock_client,
            model="claude-2",
            prompt_tokens=100,
            completion_tokens=50,
            latency=1.5,
        )

        instruments["cost_dollars"].record.assert_not_called()

    def test_track_completion_with_error(self):
        """Test tracking completion with error."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"
        mock_client._sdk_logger = MagicMock()

        track_llm_completion(
            client=mock_client,
            model="gpt-4",
            prompt_tokens=100,
            completion_tokens=0,
            latency=0.5,
            status="error",
            error="Rate limit exceeded",
        )

        # Error metric recorded, success counter not recorded
        metric_names = [c[0][0] for c in mock_client.record_metric.call_args_list]
        assert "genai.errors_total" in metric_names
        assert "genai.requests_total" not in metric_names
        # Error logged on client's sdk logger
        mock_client._sdk_logger.error.assert_called_once()
        assert "Rate limit exceeded" in str(mock_client._sdk_logger.error.call_args)


class TestTrackLLMEmbedding:
    """Test track_llm_embedding() function."""

    def test_track_embedding_success(self):
        """Test tracking successful embedding generation."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"

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

        mock_client.record_metric.assert_any_call(
            "genai.embedding_requests_total",
            1,
            model="text-embedding-ada-002",
            provider="openai",
            status="success",
        )
        mock_client.record_metric.assert_any_call(
            "genai.embedding_tokens_total",
            100,
            model="text-embedding-ada-002",
            provider="openai",
        )
        mock_client.record_metric.assert_any_call(
            "genai.embedding_latency_seconds",
            0.5,
            model="text-embedding-ada-002",
            provider="openai",
        )
        mock_client.record_metric.assert_any_call(
            "genai.embedding_dimensions_count",
            1536,
            model="text-embedding-ada-002",
            provider="openai",
        )

    def test_track_embedding_includes_model_in_attributes(self):
        """Test that model is included in metric attributes."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"

        track_llm_embedding(
            client=mock_client,
            model="custom-embedding-model",
            tokens=200,
            latency=1.0,
            dimensions=768,
        )

        # Verify model kwarg appears in tokens and latency record_metric calls
        tokens_call = next(
            c
            for c in mock_client.record_metric.call_args_list
            if c[0][0] == "genai.embedding_tokens_total"
        )
        assert tokens_call[1]["model"] == "custom-embedding-model"

        latency_call = next(
            c
            for c in mock_client.record_metric.call_args_list
            if c[0][0] == "genai.embedding_latency_seconds"
        )
        assert latency_call[1]["model"] == "custom-embedding-model"

    def test_track_embedding_with_error_status(self):
        """Test tracking embedding with error status."""
        mock_client = MagicMock()
        mock_client._metric_prefix = "genai"
        mock_client._sdk_logger = MagicMock()

        track_llm_embedding(
            client=mock_client,
            model="text-embedding-ada-002",
            tokens=50,
            latency=0.2,
            dimensions=1536,
            status="error",
        )

        # Error metric recorded, success counter not recorded
        metric_names = [c[0][0] for c in mock_client.record_metric.call_args_list]
        assert "genai.embedding_errors_total" in metric_names
        assert "genai.embedding_requests_total" not in metric_names

        error_call = next(
            c
            for c in mock_client.record_metric.call_args_list
            if c[0][0] == "genai.embedding_errors_total"
        )
        assert error_call[1]["status"] == "error"


class TestCalculateTokenRate:
    """Test calculate_token_rate() function."""

    def test_calculate_rates_normal_case(self):
        """Test token rate calculation with normal values."""
        rates = calculate_token_rate(prompt_tokens=100, completion_tokens=50, latency=1.5)

        assert rates["total_rate"] == pytest.approx(100.0, rel=1e-6)  # 150 / 1.5
        assert rates["prompt_rate"] == pytest.approx(66.67, rel=1e-2)  # 100 / 1.5
        assert rates["completion_rate"] == pytest.approx(33.33, rel=1e-2)  # 50 / 1.5

    def test_calculate_rates_with_zero_latency(self):
        """Test that zero latency returns zero rates."""
        rates = calculate_token_rate(prompt_tokens=100, completion_tokens=50, latency=0.0)

        assert rates["total_rate"] == 0.0
        assert rates["prompt_rate"] == 0.0
        assert rates["completion_rate"] == 0.0

    def test_calculate_rates_with_zero_tokens(self):
        """Test rate calculation with zero tokens."""
        rates = calculate_token_rate(prompt_tokens=0, completion_tokens=0, latency=1.0)

        assert rates["total_rate"] == 0.0
        assert rates["prompt_rate"] == 0.0
        assert rates["completion_rate"] == 0.0

    def test_calculate_rates_with_high_token_count(self):
        """Test rate calculation with high token counts."""
        rates = calculate_token_rate(prompt_tokens=10000, completion_tokens=5000, latency=10.0)

        assert rates["total_rate"] == pytest.approx(1500.0, rel=1e-6)
        assert rates["prompt_rate"] == pytest.approx(1000.0, rel=1e-6)
        assert rates["completion_rate"] == pytest.approx(500.0, rel=1e-6)

    def test_calculate_rates_with_very_fast_response(self):
        """Test rate calculation with very low latency."""
        rates = calculate_token_rate(prompt_tokens=100, completion_tokens=50, latency=0.01)

        assert rates["total_rate"] == pytest.approx(15000.0, rel=1e-6)  # 150 / 0.01
        assert rates["prompt_rate"] == pytest.approx(10000.0, rel=1e-6)  # 100 / 0.01
        assert rates["completion_rate"] == pytest.approx(5000.0, rel=1e-6)  # 50 / 0.01


class TestOpenAIPricingConstant:
    """Test OPENAI_PRICING constant."""

    def test_pricing_has_expected_models(self):
        """Test that pricing dict has expected models."""
        assert "gpt-4" in OPENAI_PRICING
        assert "gpt-4-32k" in OPENAI_PRICING
        assert "gpt-3.5-turbo" in OPENAI_PRICING
        assert "gpt-3.5-turbo-16k" in OPENAI_PRICING

    def test_each_model_has_prompt_and_completion_pricing(self):
        """Test that each model has both prompt and completion pricing."""
        for model, pricing in OPENAI_PRICING.items():
            assert "prompt" in pricing
            assert "completion" in pricing
            assert isinstance(pricing["prompt"], (int, float))
            assert isinstance(pricing["completion"], (float, int))

    def test_gpt4_pricing_higher_than_gpt35(self):
        """Test that GPT-4 pricing is higher than GPT-3.5."""
        gpt4_pricing = OPENAI_PRICING["gpt-4"]
        gpt35_pricing = OPENAI_PRICING["gpt-3.5-turbo"]

        assert gpt4_pricing["prompt"] > gpt35_pricing["prompt"]
        assert gpt4_pricing["completion"] > gpt35_pricing["completion"]
