# Copyright (c) Meta Platforms, Inc. and affiliates.
# All rights reserved.
#
# This source code is licensed under the terms described in the LICENSE file in
# the root directory of this source tree.

"""
Unit tests for OTel Telemetry Provider.

These tests focus on the provider's functionality:
- Initialization and configuration
- FastAPI middleware setup
- SQLAlchemy instrumentation
- Environment variable handling
"""

from unittest.mock import MagicMock

import pytest

from llama_stack.providers.inline.telemetry.otel.config import OTelTelemetryConfig
from llama_stack.providers.inline.telemetry.otel.otel import OTelTelemetryProvider


@pytest.fixture
def otel_config():
    """Fixture providing a basic OTelTelemetryConfig."""
    return OTelTelemetryConfig(
        service_name="test-service",
        service_version="1.0.0",
        deployment_environment="test",
        span_processor="simple",
    )


@pytest.fixture
def otel_provider(otel_config, monkeypatch):
    """Fixture providing an OTelTelemetryProvider instance."""
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318")
    return OTelTelemetryProvider(config=otel_config)


class TestOTelProviderInitialization:
    """Tests for OTel provider initialization and configuration."""

    def test_provider_initializes_with_valid_config(self, otel_config, monkeypatch):
        """Test that provider initializes correctly with valid configuration."""
        monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318")

        provider = OTelTelemetryProvider(config=otel_config)

        assert provider.config == otel_config
        assert provider.config.service_name == "test-service"
        assert provider.config.service_version == "1.0.0"
        assert provider.config.deployment_environment == "test"

    def test_initialization_with_batch_processor(self, monkeypatch):
        """Test initialization with batch span processor."""
        monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318")
        config = OTelTelemetryConfig(
            service_name="test-service",
            service_version="1.0.0",
            deployment_environment="test",
            span_processor="batch",
        )

        provider = OTelTelemetryProvider(config=config)

        assert provider.config.span_processor == "batch"

    def test_warns_when_endpoints_missing(self, otel_config, monkeypatch, caplog):
        """Test that warnings are issued when OTLP endpoints are not set."""
        monkeypatch.delenv("OTEL_EXPORTER_OTLP_ENDPOINT", raising=False)
        monkeypatch.delenv("OTEL_EXPORTER_OTLP_TRACES_ENDPOINT", raising=False)
        monkeypatch.delenv("OTEL_EXPORTER_OTLP_METRICS_ENDPOINT", raising=False)

        OTelTelemetryProvider(config=otel_config)

        # Check that warnings were logged
        assert any("Traces will not be exported" in record.message for record in caplog.records)
        assert any("Metrics will not be exported" in record.message for record in caplog.records)


class TestOTelProviderMiddleware:
    """Tests for FastAPI and SQLAlchemy instrumentation."""

    def test_fastapi_middleware_can_be_applied(self, otel_provider):
        """Test that fastapi_middleware can be called without errors."""
        mock_app = MagicMock()

        # Should not raise an exception
        otel_provider.fastapi_middleware(mock_app)

        # Verify FastAPIInstrumentor was called (it patches the app)
        # The actual instrumentation is tested in E2E tests

    def test_sqlalchemy_instrumentation_without_engine(self, otel_provider):
        """
        Test that sqlalchemy_instrumentation can be called.

        Note: Testing with a real engine would require SQLAlchemy setup.
        The actual instrumentation is tested when used with real databases.
        """
        # Should not raise an exception
        otel_provider.sqlalchemy_instrumentation()


class TestOTelProviderConfiguration:
    """Tests for configuration and environment variable handling."""

    def test_service_metadata_configuration(self, otel_provider):
        """Test that service metadata is properly configured."""
        assert otel_provider.config.service_name == "test-service"
        assert otel_provider.config.service_version == "1.0.0"
        assert otel_provider.config.deployment_environment == "test"

    def test_span_processor_configuration(self, monkeypatch):
        """Test different span processor configurations."""
        monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://localhost:4318")

        # Test simple processor
        config_simple = OTelTelemetryConfig(
            service_name="test",
            span_processor="simple",
        )
        provider_simple = OTelTelemetryProvider(config=config_simple)
        assert provider_simple.config.span_processor == "simple"

        # Test batch processor
        config_batch = OTelTelemetryConfig(
            service_name="test",
            span_processor="batch",
        )
        provider_batch = OTelTelemetryProvider(config=config_batch)
        assert provider_batch.config.span_processor == "batch"

    def test_sample_run_config_generation(self):
        """Test that sample_run_config generates valid configuration."""
        sample_config = OTelTelemetryConfig.sample_run_config()

        assert "service_name" in sample_config
        assert "span_processor" in sample_config
        assert "${env.OTEL_SERVICE_NAME" in sample_config["service_name"]