"""Tests for OpenTelemetry integration.

验证 OpenTelemetry 追踪功能的集成测试。
"""

from __future__ import annotations

import pytest
from opentelemetry import trace
from opentelemetry.sdk.trace import TracerProvider
from opentelemetry.sdk.trace.export import SimpleSpanProcessor
from opentelemetry.sdk.trace.export.in_memory_span_exporter import InMemorySpanExporter

from sagellm_gateway.telemetry import (
    TelemetryConfig,
    add_span_attributes,
    add_span_event,
    get_tracer,
    load_telemetry_config_from_env,
    setup_telemetry,
)


@pytest.fixture
def in_memory_exporter() -> InMemorySpanExporter:
    """创建内存中的 span exporter 用于测试."""
    # 保存原有的 TracerProvider
    original_provider = trace.get_tracer_provider()

    exporter = InMemorySpanExporter()

    # 创建 TracerProvider 并添加 exporter
    provider = TracerProvider()
    provider.add_span_processor(SimpleSpanProcessor(exporter))

    # 使用 _internal 方法强制设置（测试用）
    from opentelemetry.trace import (
        _TRACER_PROVIDER,  # noqa: F811
        _TRACER_PROVIDER_SET_ONCE,
    )

    _TRACER_PROVIDER = None  # noqa: F811
    _TRACER_PROVIDER_SET_ONCE._done = False
    trace.set_tracer_provider(provider)

    yield exporter

    # 清理
    exporter.clear()

    # 恢复原有的 TracerProvider
    _TRACER_PROVIDER = None
    _TRACER_PROVIDER_SET_ONCE._done = False
    if original_provider is not None:
        trace.set_tracer_provider(original_provider)


class TestTelemetryConfig:
    """测试 TelemetryConfig."""

    def test_default_config(self) -> None:
        """测试默认配置."""
        config = TelemetryConfig()
        assert config.service_name == "sagellm-gateway"
        assert config.enabled is True
        assert config.otlp_endpoint is None
        assert config.console_exporter is False
        assert config.sample_rate == 1.0

    def test_custom_config(self) -> None:
        """测试自定义配置."""
        config = TelemetryConfig(
            service_name="test-service",
            enabled=False,
            otlp_endpoint="http://localhost:4317",
            console_exporter=True,
            sample_rate=0.5,
        )
        assert config.service_name == "test-service"
        assert config.enabled is False
        assert config.otlp_endpoint == "http://localhost:4317"
        assert config.console_exporter is True
        assert config.sample_rate == 0.5

    def test_load_from_env(self, monkeypatch: pytest.MonkeyPatch) -> None:
        """测试从环境变量加载配置."""
        monkeypatch.setenv("SAGELLM_GATEWAY_TELEMETRY_ENABLED", "false")
        monkeypatch.setenv("SAGELLM_GATEWAY_SERVICE_NAME", "test-gateway")
        monkeypatch.setenv("OTEL_EXPORTER_OTLP_ENDPOINT", "http://test:4317")
        monkeypatch.setenv("SAGELLM_GATEWAY_TELEMETRY_CONSOLE", "true")
        monkeypatch.setenv("SAGELLM_GATEWAY_TELEMETRY_SAMPLE_RATE", "0.1")

        config = load_telemetry_config_from_env()
        assert config.enabled is False
        assert config.service_name == "test-gateway"
        assert config.otlp_endpoint == "http://test:4317"
        assert config.console_exporter is True
        assert config.sample_rate == 0.1


class TestSetupTelemetry:
    """测试 setup_telemetry."""

    def test_setup_disabled(self) -> None:
        """测试禁用追踪."""
        config = TelemetryConfig(enabled=False)
        provider = setup_telemetry(config)
        assert provider is None

    def test_setup_enabled(self) -> None:
        """测试启用追踪."""
        config = TelemetryConfig(enabled=True)
        provider = setup_telemetry(config)
        assert provider is not None
        assert isinstance(provider, TracerProvider)

    def test_setup_with_console_exporter(self) -> None:
        """测试启用控制台导出."""
        config = TelemetryConfig(enabled=True, console_exporter=True)
        provider = setup_telemetry(config)
        assert provider is not None


class TestTracerAPI:
    """测试 Tracer API."""

    def test_get_tracer(self, in_memory_exporter: InMemorySpanExporter) -> None:
        """测试获取 tracer."""
        tracer = get_tracer("test")
        assert tracer is not None

        # 创建一个 span
        with tracer.start_as_current_span("test-span"):
            pass

        # 验证 span 已导出
        spans = in_memory_exporter.get_finished_spans()
        assert len(spans) == 1
        assert spans[0].name == "test-span"

    def test_span_attributes(self, in_memory_exporter: InMemorySpanExporter) -> None:
        """测试添加 span 属性."""
        tracer = get_tracer("test")

        with tracer.start_as_current_span("test-span"):
            add_span_attributes(
                {
                    "model": "gpt2",
                    "temperature": 0.8,
                    "max_tokens": 100,
                }
            )

        spans = in_memory_exporter.get_finished_spans()
        assert len(spans) == 1

        span = spans[0]
        assert span.attributes["model"] == "gpt2"  # type: ignore
        assert span.attributes["temperature"] == 0.8  # type: ignore
        assert span.attributes["max_tokens"] == 100  # type: ignore

    def test_span_events(self, in_memory_exporter: InMemorySpanExporter) -> None:
        """测试添加 span 事件."""
        tracer = get_tracer("test")

        with tracer.start_as_current_span("test-span"):
            add_span_event("cache_miss")
            add_span_event("token_generated", {"count": 42})

        spans = in_memory_exporter.get_finished_spans()
        assert len(spans) == 1

        span = spans[0]
        events = span.events
        assert len(events) == 2
        assert events[0].name == "cache_miss"
        assert events[1].name == "token_generated"
        assert events[1].attributes["count"] == 42  # type: ignore

    def test_nested_spans(self, in_memory_exporter: InMemorySpanExporter) -> None:
        """测试嵌套 span."""
        tracer = get_tracer("test")

        with tracer.start_as_current_span("parent"):
            add_span_attributes({"parent": "true"})

            with tracer.start_as_current_span("child"):
                add_span_attributes({"child": "true"})

        spans = in_memory_exporter.get_finished_spans()
        assert len(spans) == 2

        # 子 span 应该先完成
        child_span = spans[0]
        parent_span = spans[1]

        assert child_span.name == "child"
        assert parent_span.name == "parent"

        # 验证父子关系
        assert child_span.parent is not None
        assert child_span.parent.span_id == parent_span.context.span_id

    def test_span_without_recording(self, in_memory_exporter: InMemorySpanExporter) -> None:
        """测试在没有活动 span 时添加属性（不应抛出异常）."""
        # 不应该抛出异常
        add_span_attributes({"key": "value"})
        add_span_event("event")

        # 应该没有 span 被记录
        spans = in_memory_exporter.get_finished_spans()
        assert len(spans) == 0


class TestIntegration:
    """集成测试."""

    def test_full_workflow(self, in_memory_exporter: InMemorySpanExporter) -> None:
        """测试完整的追踪工作流."""
        tracer = get_tracer(__name__)

        # 模拟一个 API 请求
        with tracer.start_as_current_span("api_request") as span:
            add_span_attributes(
                {
                    "http.method": "POST",
                    "http.url": "/v1/chat/completions",
                    "user.id": "user123",
                }
            )

            # 模拟速率限制检查
            with tracer.start_as_current_span("rate_limit_check"):
                add_span_event("rate_limit_passed")

            # 模拟推理
            with tracer.start_as_current_span("inference"):
                add_span_attributes(
                    {
                        "model": "gpt2",
                        "max_tokens": 100,
                    }
                )
                add_span_event("model_loaded")
                add_span_event("generation_completed", {"tokens": 42})

            span.set_status(trace.Status(trace.StatusCode.OK))

        # 验证所有 span
        spans = in_memory_exporter.get_finished_spans()
        assert len(spans) == 3

        span_names = {s.name for s in spans}
        assert span_names == {"api_request", "rate_limit_check", "inference"}

        # 验证属性和事件
        api_span = next(s for s in spans if s.name == "api_request")
        assert api_span.attributes["http.method"] == "POST"  # type: ignore
        assert api_span.attributes["user.id"] == "user123"  # type: ignore

        inference_span = next(s for s in spans if s.name == "inference")
        assert inference_span.attributes["model"] == "gpt2"  # type: ignore
        assert len(inference_span.events) == 2
