"""Tests for the TencentTraceClient helpers that drive tracing and metrics."""

from __future__ import annotations

import logging
import sys
import types
from types import SimpleNamespace
from typing import Any, TypedDict, cast
from unittest.mock import MagicMock

import pytest
from dify_trace_tencent import client as client_module
from dify_trace_tencent.client import TencentTraceClient, _get_opentelemetry_sdk_version
from dify_trace_tencent.entities.tencent_trace_entity import SpanData
from opentelemetry.sdk.trace import Event
from opentelemetry.trace import SpanContext, Status, StatusCode, TraceFlags

metric_reader_instances: list[DummyMetricReader] = []
meter_provider_instances: list[DummyMeterProvider] = []


class DummyHistogram:
    """Placeholder histogram type used by the stubbed metric stack."""


class AggregationTemporality:
    DELTA = "delta"


class DummyMeter:
    def __init__(self) -> None:
        self.created: list[tuple[dict[str, object], MagicMock]] = []

    def create_histogram(self, **kwargs: object) -> MagicMock:
        hist = MagicMock(name=f"hist-{kwargs.get('name')}")
        self.created.append((kwargs, hist))
        return hist


class DummyMeterProvider:
    def __init__(self, resource: object, metric_readers: list[object]) -> None:
        self.resource = resource
        self.metric_readers = metric_readers
        self.meter = DummyMeter()
        self.shutdown = MagicMock(name="meter_provider_shutdown")
        meter_provider_instances.append(self)

    def get_meter(self, name: str, version: str) -> DummyMeter:
        return self.meter


class DummyMetricReader:
    def __init__(self, exporter: object, export_interval_millis: int) -> None:
        self.exporter = exporter
        self.export_interval_millis = export_interval_millis
        self.shutdown = MagicMock(name="metric_reader_shutdown")
        metric_reader_instances.append(self)


class DummyGrpcMetricExporter:
    def __init__(self, **kwargs: object) -> None:
        self.kwargs = kwargs


class DummyHttpMetricExporter:
    def __init__(self, **kwargs: object) -> None:
        self.kwargs = kwargs


class DummyJsonMetricExporter:
    def __init__(self, **kwargs: object) -> None:
        self.kwargs = kwargs


class DummyJsonMetricExporterNoTemporality:
    """Exporter that rejects preferred_temporality to exercise fallback."""

    def __init__(self, **kwargs: object) -> None:
        if "preferred_temporality" in kwargs:
            raise RuntimeError("unsupported preferred_temporality")
        self.kwargs = kwargs


class PatchedCoreComponents(TypedDict):
    span_exporter: MagicMock
    span_processor: MagicMock
    tracer: MagicMock
    span: MagicMock
    tracer_provider: MagicMock
    trace_api: Any


def _add_stub_modules(monkeypatch: pytest.MonkeyPatch) -> None:
    """Drop fake metric modules into sys.modules so the client imports resolve."""

    metrics_module = cast(Any, types.ModuleType("opentelemetry.sdk.metrics"))
    metrics_module.Histogram = DummyHistogram
    metrics_module.MeterProvider = DummyMeterProvider
    monkeypatch.setitem(sys.modules, "opentelemetry.sdk.metrics", metrics_module)

    metrics_export_module = cast(Any, types.ModuleType("opentelemetry.sdk.metrics.export"))
    metrics_export_module.AggregationTemporality = AggregationTemporality
    metrics_export_module.PeriodicExportingMetricReader = DummyMetricReader
    monkeypatch.setitem(sys.modules, "opentelemetry.sdk.metrics.export", metrics_export_module)

    grpc_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.proto.grpc.metric_exporter"))
    grpc_module.OTLPMetricExporter = DummyGrpcMetricExporter
    monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.proto.grpc.metric_exporter", grpc_module)

    http_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.proto.http.metric_exporter"))
    http_module.OTLPMetricExporter = DummyHttpMetricExporter
    monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.proto.http.metric_exporter", http_module)

    http_json_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.http.json.metric_exporter"))
    http_json_module.OTLPMetricExporter = DummyJsonMetricExporter
    monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.http.json.metric_exporter", http_json_module)

    legacy_json_module = cast(Any, types.ModuleType("opentelemetry.exporter.otlp.json.metric_exporter"))
    legacy_json_module.OTLPMetricExporter = DummyJsonMetricExporter
    monkeypatch.setitem(sys.modules, "opentelemetry.exporter.otlp.json.metric_exporter", legacy_json_module)


@pytest.fixture(autouse=True)
def stub_metric_modules(monkeypatch: pytest.MonkeyPatch) -> None:
    metric_reader_instances.clear()
    meter_provider_instances.clear()
    _add_stub_modules(monkeypatch)


@pytest.fixture(autouse=True)
def patch_core_components(monkeypatch: pytest.MonkeyPatch) -> PatchedCoreComponents:
    span_exporter = MagicMock(name="span_exporter")
    monkeypatch.setattr(client_module, "OTLPSpanExporter", MagicMock(return_value=span_exporter))

    span_processor = MagicMock(name="span_processor")
    monkeypatch.setattr(client_module, "BatchSpanProcessor", MagicMock(return_value=span_processor))

    tracer = MagicMock(name="tracer")
    span = MagicMock(name="span")
    tracer.start_span.return_value = span

    tracer_provider = MagicMock(name="tracer_provider")
    tracer_provider.get_tracer.return_value = tracer
    tracer_provider.shutdown = MagicMock(name="tracer_provider_shutdown")
    monkeypatch.setattr(client_module, "TracerProvider", MagicMock(return_value=tracer_provider))

    resource = MagicMock(name="resource")
    monkeypatch.setattr(client_module, "Resource", MagicMock(return_value=resource))

    trace_api_stub = SimpleNamespace(
        set_span_in_context=MagicMock(name="set_span_in_context", return_value="trace-context"),
        NonRecordingSpan=MagicMock(name="non_recording_span", side_effect=lambda ctx: f"non-{ctx}"),
    )
    monkeypatch.setattr(client_module, "trace_api", trace_api_stub)

    fake_config = SimpleNamespace(
        project=SimpleNamespace(version="test"),
        COMMIT_SHA="sha",
        DEPLOY_ENV="dev",
        EDITION="cloud",
    )
    monkeypatch.setattr(client_module, "dify_config", fake_config)

    monkeypatch.setattr(client_module.socket, "gethostname", lambda: "fake-host")
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "")

    return {
        "span_exporter": span_exporter,
        "span_processor": span_processor,
        "tracer": tracer,
        "span": span,
        "tracer_provider": tracer_provider,
        "trace_api": trace_api_stub,
    }


def _make_span_context(trace_id: int = 1, span_id: int = 2) -> SpanContext:
    return SpanContext(
        trace_id=trace_id,
        span_id=span_id,
        is_remote=False,
        trace_flags=TraceFlags(TraceFlags.SAMPLED),
    )


def _build_client() -> TencentTraceClient:
    return TencentTraceClient(
        service_name="service",
        endpoint="https://trace.example.com:4317",
        token="token",
    )


def test_get_opentelemetry_sdk_version_reads_install(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(client_module, "version", lambda pkg: "2.0.0")
    assert _get_opentelemetry_sdk_version() == "2.0.0"


def test_get_opentelemetry_sdk_version_falls_back(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(client_module, "version", MagicMock(side_effect=RuntimeError("boom")))
    assert _get_opentelemetry_sdk_version() == "1.27.0"


@pytest.mark.parametrize(
    ("endpoint", "expected"),
    [
        (
            "https://example.com:9090",
            ("example.com:9090", False, "example.com", 9090),
        ),
        (
            "http://localhost",
            ("localhost:4317", True, "localhost", 4317),
        ),
        (
            "example.com:bad",
            ("example.com:4317", False, "example.com", 4317),
        ),
    ],
)
def test_resolve_grpc_target_parsable_variants(endpoint: str, expected: tuple[str, bool, str, int]) -> None:
    assert TencentTraceClient._resolve_grpc_target(endpoint) == expected


def test_resolve_grpc_target_handles_errors() -> None:
    assert TencentTraceClient._resolve_grpc_target(cast(str, 123)) == ("localhost:4317", True, "localhost", 4317)


@pytest.mark.parametrize(
    ("method", "attr_name", "args"),
    [
        ("record_llm_duration", "hist_llm_duration", (0.3, {"foo": object()})),
        ("record_token_usage", "hist_token_usage", (5, "input", "chat", "gpt", "gpt", "addr", "dify")),
        ("record_time_to_first_token", "hist_time_to_first_token", (0.4, "dify", "gpt")),
        ("record_time_to_generate", "hist_time_to_generate", (0.6, "dify", "gpt")),
        ("record_trace_duration", "hist_trace_duration", (1.0, {"meta": object()})),
    ],
)
def test_record_methods_call_histograms(method: str, attr_name: str, args: tuple[object, ...]) -> None:
    client = _build_client()
    hist_mock = MagicMock(name=attr_name)
    setattr(client, attr_name, hist_mock)

    getattr(client, method)(*args)
    hist_mock.record.assert_called_once()


def test_record_methods_skip_when_histogram_missing() -> None:
    client = _build_client()
    client.hist_llm_duration = None
    client.record_llm_duration(0.1)

    client.hist_token_usage = None
    client.record_token_usage(1, "go", "chat", "model", "model", "addr", "provider")

    client.hist_time_to_first_token = None
    client.record_time_to_first_token(0.2, "prov", "model")

    client.hist_time_to_generate = None
    client.record_time_to_generate(0.3, "prov", "model")

    client.hist_trace_duration = None
    client.record_trace_duration(0.5)


def test_record_llm_duration_handles_exceptions(caplog: pytest.LogCaptureFixture) -> None:
    client = _build_client()
    client.hist_llm_duration = MagicMock(name="hist_llm_duration")
    client.hist_llm_duration.record.side_effect = RuntimeError("boom")

    caplog.set_level(logging.DEBUG, logger=client_module.logger.name)
    client.record_llm_duration(0.2)

    assert "[Tencent APM] Failed to record LLM duration" in caplog.text


def test_create_and_export_span_sets_attributes(patch_core_components: PatchedCoreComponents) -> None:
    client = _build_client()
    span = patch_core_components["span"]
    ctx = _make_span_context(span_id=2)
    span.get_span_context.return_value = ctx

    data = SpanData(
        trace_id=1,
        parent_span_id=None,
        span_id=2,
        name="span",
        attributes={"key": "value"},
        events=[Event(name="evt", attributes={"k": "v"}, timestamp=123)],
        status=Status(StatusCode.OK),
        start_time=10,
        end_time=20,
    )

    client._create_and_export_span(data)
    span.set_attributes.assert_called_once()
    span.add_event.assert_called_once()
    span.set_status.assert_called_once()
    span.end.assert_called_once_with(end_time=20)
    assert client.span_contexts[2] == ctx


def test_create_and_export_span_uses_parent_context(patch_core_components: PatchedCoreComponents) -> None:
    client = _build_client()
    existing_context = _make_span_context(span_id=10)
    client.span_contexts[10] = existing_context
    span = patch_core_components["span"]
    span.get_span_context.return_value = _make_span_context(span_id=11)

    data = SpanData(
        trace_id=1,
        parent_span_id=10,
        span_id=11,
        name="span",
        attributes={},
        events=[],
        start_time=0,
        end_time=1,
    )

    client._create_and_export_span(data)
    trace_api = patch_core_components["trace_api"]
    trace_api.NonRecordingSpan.assert_called_once_with(existing_context)
    trace_api.set_span_in_context.assert_called_once()


def test_create_and_export_span_exception_logs_error(
    patch_core_components: PatchedCoreComponents, caplog: pytest.LogCaptureFixture
) -> None:
    client = _build_client()
    span = patch_core_components["span"]
    span.get_span_context.return_value = _make_span_context(span_id=2)
    client.tracer.start_span.side_effect = RuntimeError("boom")

    caplog.set_level(logging.DEBUG, logger=client_module.logger.name)
    client._create_and_export_span(
        SpanData(
            trace_id=1,
            parent_span_id=None,
            span_id=2,
            name="span",
            attributes={},
            events=[],
            start_time=0,
            end_time=1,
        )
    )

    error_records = [record for record in caplog.records if record.levelno == logging.ERROR]
    assert len(error_records) == 1
    assert error_records[0].getMessage() == "[Tencent APM] Error creating span: span"


def test_api_check_connects_successfully(monkeypatch: pytest.MonkeyPatch) -> None:
    client = _build_client()

    monkeypatch.setattr(
        TencentTraceClient,
        "_resolve_grpc_target",
        MagicMock(return_value=("host:123", False, "host", 123)),
    )

    socket_mock = MagicMock()
    socket_instance = MagicMock()
    socket_instance.connect_ex.return_value = 0
    socket_mock.return_value = socket_instance
    monkeypatch.setattr(client_module.socket, "socket", socket_mock)

    assert client.api_check()
    socket_instance.connect_ex.assert_called_once()


def test_api_check_returns_false_and_handles_local(monkeypatch: pytest.MonkeyPatch) -> None:
    client = _build_client()

    monkeypatch.setattr(
        TencentTraceClient,
        "_resolve_grpc_target",
        MagicMock(return_value=("host:123", False, "host", 123)),
    )

    socket_mock = MagicMock()
    socket_instance = MagicMock()
    socket_instance.connect_ex.return_value = 1
    socket_mock.return_value = socket_instance
    monkeypatch.setattr(client_module.socket, "socket", socket_mock)

    assert not client.api_check()

    monkeypatch.setattr(
        TencentTraceClient,
        "_resolve_grpc_target",
        MagicMock(return_value=("localhost:4317", True, "localhost", 4317)),
    )
    socket_instance.connect_ex.return_value = 1
    assert client.api_check()


def test_api_check_handles_exceptions(monkeypatch: pytest.MonkeyPatch) -> None:
    client = TencentTraceClient("svc", "https://localhost", "token")

    monkeypatch.setattr(client_module.socket, "socket", MagicMock(side_effect=RuntimeError("boom")))
    assert client.api_check()


def test_get_project_url() -> None:
    client = _build_client()
    assert client.get_project_url() == "https://console.cloud.tencent.com/apm"


def test_shutdown_flushes_all_components(patch_core_components: PatchedCoreComponents) -> None:
    client = _build_client()
    span_processor = patch_core_components["span_processor"]
    tracer_provider = patch_core_components["tracer_provider"]

    client.shutdown()
    span_processor.force_flush.assert_called_once()
    span_processor.shutdown.assert_called_once()
    tracer_provider.shutdown.assert_called_once()

    meter_provider = meter_provider_instances[-1]
    metric_reader = metric_reader_instances[-1]
    meter_provider.shutdown.assert_called_once()
    metric_reader.shutdown.assert_called_once()


def test_shutdown_logs_when_meter_provider_fails(caplog: pytest.LogCaptureFixture) -> None:
    client = _build_client()
    meter_provider = meter_provider_instances[-1]
    meter_provider.shutdown.side_effect = RuntimeError("boom")
    assert client.metric_reader is not None
    client.metric_reader.shutdown.side_effect = RuntimeError("boom")

    caplog.set_level(logging.DEBUG, logger=client_module.logger.name)
    client.shutdown()

    assert "[Tencent APM] Error shutting down meter provider" in caplog.text
    assert "[Tencent APM] Error shutting down metric reader" in caplog.text


def test_metrics_initialization_failure_sets_histogram_attributes(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setattr(DummyMeterProvider, "__init__", MagicMock(side_effect=RuntimeError("err")))
    client = _build_client()

    assert client.meter is None
    assert client.meter_provider is None
    assert client.hist_llm_duration is None
    assert client.hist_token_usage is None
    assert client.hist_time_to_first_token is None
    assert client.hist_time_to_generate is None
    assert client.hist_trace_duration is None
    assert client.metric_reader is None


def test_add_span_logs_exception(monkeypatch: pytest.MonkeyPatch, caplog: pytest.LogCaptureFixture) -> None:
    client = _build_client()
    monkeypatch.setattr(client, "_create_and_export_span", MagicMock(side_effect=RuntimeError("boom")))

    caplog.set_level(logging.DEBUG, logger=client_module.logger.name)
    client.add_span(
        SpanData(
            trace_id=1,
            parent_span_id=None,
            span_id=2,
            name="span",
            attributes={},
            events=[],
            start_time=0,
            end_time=1,
        )
    )

    error_records = [record for record in caplog.records if record.levelno == logging.ERROR]
    assert len(error_records) == 1
    assert error_records[0].getMessage() == "[Tencent APM] Failed to create span: span"


def test_create_and_export_span_converts_attribute_types(patch_core_components: PatchedCoreComponents) -> None:
    client = _build_client()
    span = patch_core_components["span"]
    span.get_span_context.return_value = _make_span_context(span_id=2)

    data = SpanData.model_construct(
        trace_id=1,
        parent_span_id=None,
        span_id=2,
        name="span",
        attributes={"num": 5, "flag": True, "pi": 3.14, "text": "value"},
        events=[],
        links=[],
        status=Status(StatusCode.OK),
        start_time=0,
        end_time=1,
    )

    client._create_and_export_span(data)
    (attrs,) = span.set_attributes.call_args.args
    assert attrs["num"] == 5
    assert attrs["flag"] is True
    assert attrs["pi"] == 3.14
    assert attrs["text"] == "value"


def test_record_llm_duration_converts_attributes() -> None:
    client = _build_client()
    hist_mock = MagicMock(name="hist_llm_duration")
    client.hist_llm_duration = hist_mock

    client.record_llm_duration(0.3, cast(dict[str, str], {"foo": object(), "bar": 2}))
    _, attrs = hist_mock.record.call_args.args
    assert isinstance(attrs["foo"], str)
    assert attrs["bar"] == 2


def test_record_trace_duration_converts_attributes() -> None:
    client = _build_client()
    hist_mock = MagicMock(name="hist_trace_duration")
    client.hist_trace_duration = hist_mock

    client.record_trace_duration(1.0, cast(dict[str, str], {"meta": object(), "ok": True}))
    _, attrs = hist_mock.record.call_args.args
    assert isinstance(attrs["meta"], str)
    assert attrs["ok"] is True


@pytest.mark.parametrize(
    ("method", "attr_name", "args"),
    [
        ("record_token_usage", "hist_token_usage", (5, "input", "chat", "gpt", "gpt", "addr", "dify")),
        ("record_time_to_first_token", "hist_time_to_first_token", (0.4, "dify", "gpt")),
        ("record_time_to_generate", "hist_time_to_generate", (0.6, "dify", "gpt")),
        ("record_trace_duration", "hist_trace_duration", (1.0, {"meta": object()})),
    ],
)
def test_record_methods_handle_exceptions(
    method: str, attr_name: str, args: tuple[object, ...], caplog: pytest.LogCaptureFixture
) -> None:
    client = _build_client()
    hist_mock = MagicMock(name=attr_name)
    hist_mock.record.side_effect = RuntimeError("boom")
    setattr(client, attr_name, hist_mock)

    caplog.set_level(logging.DEBUG, logger=client_module.logger.name)
    getattr(client, method)(*args)

    assert any(
        record.levelno == logging.DEBUG and record.getMessage().startswith("[Tencent APM] Failed to record")
        for record in caplog.records
    )


def test_metrics_initializes_grpc_metric_exporter() -> None:
    client = _build_client()
    metric_reader = metric_reader_instances[-1]
    exporter = cast(DummyGrpcMetricExporter, metric_reader.exporter)

    assert isinstance(exporter, DummyGrpcMetricExporter)
    assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000
    assert exporter.kwargs["endpoint"] == "trace.example.com:4317"
    assert exporter.kwargs["insecure"] is False
    assert cast(dict[str, dict[str, str]], exporter.kwargs)["headers"]["authorization"] == "Bearer token"


def test_metrics_initializes_http_protobuf_metric_exporter(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/protobuf")
    client = _build_client()
    metric_reader = metric_reader_instances[-1]
    exporter = cast(DummyHttpMetricExporter, metric_reader.exporter)

    assert isinstance(exporter, DummyHttpMetricExporter)
    assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000
    assert exporter.kwargs["endpoint"] == client.endpoint
    assert cast(dict[str, dict[str, str]], exporter.kwargs)["headers"]["authorization"] == "Bearer token"


def test_metrics_initializes_http_json_metric_exporter(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json")
    client = _build_client()
    metric_reader = metric_reader_instances[-1]
    exporter = cast(DummyJsonMetricExporter, metric_reader.exporter)

    assert isinstance(exporter, DummyJsonMetricExporter)
    assert metric_reader.export_interval_millis == client.metrics_export_interval_sec * 1000
    assert exporter.kwargs["endpoint"] == client.endpoint
    assert cast(dict[str, dict[str, str]], exporter.kwargs)["headers"]["authorization"] == "Bearer token"
    assert "preferred_temporality" in exporter.kwargs


def test_metrics_http_json_metric_exporter_falls_back_without_temporality(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json")
    exporter_module = sys.modules["opentelemetry.exporter.otlp.http.json.metric_exporter"]
    monkeypatch.setattr(exporter_module, "OTLPMetricExporter", DummyJsonMetricExporterNoTemporality)
    _ = _build_client()
    metric_reader = metric_reader_instances[-1]
    exporter = cast(DummyJsonMetricExporterNoTemporality, metric_reader.exporter)

    assert isinstance(exporter, DummyJsonMetricExporterNoTemporality)
    assert "preferred_temporality" not in exporter.kwargs


def test_metrics_http_json_uses_http_fallback_when_no_json_exporter(monkeypatch: pytest.MonkeyPatch) -> None:
    monkeypatch.setenv("OTEL_EXPORTER_OTLP_PROTOCOL", "http/json")

    def _fail_import(mod_path: str) -> types.ModuleType:
        raise ModuleNotFoundError(mod_path)

    monkeypatch.setattr(client_module.importlib, "import_module", _fail_import)

    _ = _build_client()
    metric_reader = metric_reader_instances[-1]
    assert isinstance(metric_reader.exporter, DummyHttpMetricExporter)
