"""Unit tests for EnterpriseOtelTrace."""

from __future__ import annotations

import json
from datetime import UTC, datetime
from typing import Any
from unittest.mock import MagicMock, patch

import pytest

from core.ops.entities.trace_entity import (
    DatasetRetrievalTraceInfo,
    DraftNodeExecutionTrace,
    GenerateNameTraceInfo,
    MessageTraceInfo,
    ModerationTraceInfo,
    PromptGenerationTraceInfo,
    SuggestedQuestionTraceInfo,
    ToolTraceInfo,
    WorkflowNodeTraceInfo,
    WorkflowTraceInfo,
)
from enterprise.telemetry.entities import (
    EnterpriseTelemetryCounter,
    EnterpriseTelemetryEvent,
    EnterpriseTelemetryHistogram,
    EnterpriseTelemetrySpan,
)

# ---------------------------------------------------------------------------
# Fixtures
# ---------------------------------------------------------------------------


@pytest.fixture
def mock_exporter():
    exporter = MagicMock()
    exporter.include_content = True
    return exporter


@pytest.fixture
def trace_handler(mock_exporter):
    with patch("extensions.ext_enterprise_telemetry.get_enterprise_exporter", return_value=mock_exporter):
        from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace

        handler = EnterpriseOtelTrace()
    return handler


# ---------------------------------------------------------------------------
# Factory helpers
# ---------------------------------------------------------------------------

_T0 = datetime(2024, 1, 10, 12, 0, 0, tzinfo=UTC)
_T1 = datetime(2024, 1, 10, 12, 0, 5, tzinfo=UTC)


def make_workflow_info(**overrides) -> WorkflowTraceInfo:
    defaults: dict[str, Any] = {
        "workflow_id": "wf-001",
        "tenant_id": "tenant-abc",
        "workflow_run_id": "run-001",
        "workflow_run_elapsed_time": 5.0,
        "workflow_run_status": "succeeded",
        "workflow_run_inputs": {"query": "hello"},
        "workflow_run_outputs": {"answer": "world"},
        "workflow_run_version": "1",
        "total_tokens": 100,
        "prompt_tokens": 60,
        "completion_tokens": 40,
        "file_list": [],
        "query": "hello",
        "start_time": _T0,
        "end_time": _T1,
        "metadata": {
            "app_id": "app-001",
            "tenant_id": "tenant-abc",
            "app_name": "MyApp",
            "workspace_name": "WS",
            "triggered_from": "api",
        },
    }
    defaults.update(overrides)
    return WorkflowTraceInfo(**defaults)


def make_node_info(**overrides) -> WorkflowNodeTraceInfo:
    defaults: dict[str, Any] = {
        "workflow_id": "wf-001",
        "workflow_run_id": "run-001",
        "tenant_id": "tenant-abc",
        "node_execution_id": "ne-001",
        "node_id": "node-001",
        "node_type": "llm",
        "title": "LLM Node",
        "status": "succeeded",
        "elapsed_time": 2.5,
        "index": 1,
        "total_tokens": 80,
        "prompt_tokens": 50,
        "completion_tokens": 30,
        "model_provider": "openai",
        "model_name": "gpt-4",
        "start_time": _T0,
        "end_time": _T1,
        "metadata": {
            "app_id": "app-001",
            "tenant_id": "tenant-abc",
            "app_name": "MyApp",
        },
    }
    defaults.update(overrides)
    return WorkflowNodeTraceInfo(**defaults)


def make_draft_node_info(**overrides) -> DraftNodeExecutionTrace:
    defaults: dict[str, Any] = {
        "workflow_id": "wf-001",
        "workflow_run_id": "run-draft-001",
        "tenant_id": "tenant-abc",
        "node_execution_id": "ne-draft-001",
        "node_id": "node-001",
        "node_type": "llm",
        "title": "Draft LLM",
        "status": "succeeded",
        "elapsed_time": 1.2,
        "index": 0,
        "total_tokens": 50,
        "start_time": _T0,
        "end_time": _T1,
        "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"},
    }
    defaults.update(overrides)
    return DraftNodeExecutionTrace(**defaults)


def make_message_info(**overrides) -> MessageTraceInfo:
    defaults: dict[str, Any] = {
        "message_id": "msg-001",
        "conversation_model": "gpt-4",
        "message_tokens": 40,
        "answer_tokens": 60,
        "total_tokens": 100,
        "conversation_mode": "chat",
        "start_time": _T0,
        "end_time": _T1,
        "inputs": "user input",
        "outputs": "assistant output",
        "metadata": {
            "app_id": "app-001",
            "tenant_id": "tenant-abc",
            "from_source": "api",
            "ls_provider": "openai",
            "ls_model_name": "gpt-4",
            "status": "succeeded",
        },
    }
    defaults.update(overrides)
    return MessageTraceInfo(**defaults)


def make_tool_info(**overrides) -> ToolTraceInfo:
    defaults: dict[str, Any] = {
        "message_id": "msg-001",
        "tool_name": "web_search",
        "tool_inputs": {"query": "test"},
        "tool_outputs": "search results",
        "tool_config": {"max_results": 5},
        "tool_parameters": {"verbose": True},
        "time_cost": 1.5,
        "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"},
    }
    defaults.update(overrides)
    return ToolTraceInfo(**defaults)


def make_moderation_info(**overrides) -> ModerationTraceInfo:
    defaults: dict[str, Any] = {
        "message_id": "msg-001",
        "flagged": False,
        "action": "pass",
        "preset_response": "",
        "query": "is this ok?",
        "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"},
    }
    defaults.update(overrides)
    return ModerationTraceInfo(**defaults)


def make_suggested_question_info(**overrides) -> SuggestedQuestionTraceInfo:
    defaults: dict[str, Any] = {
        "message_id": "msg-001",
        "total_tokens": 30,
        "suggested_question": ["Question A?", "Question B?"],
        "level": "info",
        "status": "succeeded",
        "model_provider": "openai",
        "model_id": "gpt-3.5-turbo",
        "start_time": _T0,
        "end_time": _T1,
        "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"},
    }
    defaults.update(overrides)
    return SuggestedQuestionTraceInfo(**defaults)


def make_dataset_retrieval_info(**overrides) -> DatasetRetrievalTraceInfo:
    defaults: dict[str, Any] = {
        "message_id": "msg-001",
        "documents": [
            {
                "metadata": {
                    "dataset_id": "ds-001",
                    "dataset_name": "MyDataset",
                    "document_id": "doc-001",
                    "segment_id": "seg-001",
                    "score": 0.95,
                }
            }
        ],
        "inputs": "search query",
        "metadata": {
            "app_id": "app-001",
            "tenant_id": "tenant-abc",
            "embedding_models": {
                "ds-001": {
                    "embedding_model_provider": "openai",
                    "embedding_model": "text-embedding-3-small",
                }
            },
        },
    }
    defaults.update(overrides)
    return DatasetRetrievalTraceInfo(**defaults)


def make_generate_name_info(**overrides) -> GenerateNameTraceInfo:
    defaults: dict[str, Any] = {
        "message_id": "msg-001",
        "tenant_id": "tenant-abc",
        "conversation_id": "conv-001",
        "inputs": "some content",
        "outputs": "My Conversation",
        "start_time": _T0,
        "end_time": _T1,
        "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"},
    }
    defaults.update(overrides)
    return GenerateNameTraceInfo(**defaults)


def make_prompt_generation_info(**overrides) -> PromptGenerationTraceInfo:
    defaults: dict[str, Any] = {
        "tenant_id": "tenant-abc",
        "user_id": "user-001",
        "app_id": "app-001",
        "operation_type": "rule_generate",
        "instruction": "Generate a helpful prompt",
        "prompt_tokens": 50,
        "completion_tokens": 100,
        "total_tokens": 150,
        "model_provider": "openai",
        "model_name": "gpt-4",
        "latency": 3.2,
        "metadata": {"app_id": "app-001", "tenant_id": "tenant-abc"},
    }
    defaults.update(overrides)
    return PromptGenerationTraceInfo(**defaults)


# ---------------------------------------------------------------------------
# Constructor
# ---------------------------------------------------------------------------


def test_init_raises_when_exporter_is_none():
    with patch("extensions.ext_enterprise_telemetry.get_enterprise_exporter", return_value=None):
        from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace

        with pytest.raises(RuntimeError, match="exporter is not initialized"):
            EnterpriseOtelTrace()


def test_init_succeeds_with_valid_exporter(mock_exporter):
    with patch("extensions.ext_enterprise_telemetry.get_enterprise_exporter", return_value=mock_exporter):
        from enterprise.telemetry.enterprise_trace import EnterpriseOtelTrace

        handler = EnterpriseOtelTrace()
    assert handler._exporter is mock_exporter


# ---------------------------------------------------------------------------
# Helper methods
# ---------------------------------------------------------------------------


class TestSafePayloadValue:
    def test_string_passthrough(self, trace_handler):
        assert trace_handler._safe_payload_value("hello") == "hello"

    def test_dict_passthrough(self, trace_handler):
        d = {"key": "val"}
        assert trace_handler._safe_payload_value(d) == d

    def test_list_passthrough(self, trace_handler):
        lst = [1, 2, 3]
        assert trace_handler._safe_payload_value(lst) == lst

    def test_none_returns_none(self, trace_handler):
        assert trace_handler._safe_payload_value(None) is None

    def test_int_returns_none(self, trace_handler):
        assert trace_handler._safe_payload_value(42) is None

    def test_bool_returns_none(self, trace_handler):
        assert trace_handler._safe_payload_value(True) is None


class TestMaybeJson:
    def test_none_returns_none(self, trace_handler):
        assert trace_handler._maybe_json(None) is None

    def test_string_passthrough(self, trace_handler):
        assert trace_handler._maybe_json("hello") == "hello"

    def test_dict_serialised(self, trace_handler):
        result = trace_handler._maybe_json({"a": 1})
        assert result == json.dumps({"a": 1})

    def test_list_serialised(self, trace_handler):
        result = trace_handler._maybe_json([1, 2])
        assert result == "[1, 2]"

    def test_non_serialisable_falls_back_to_str(self, trace_handler):
        class Unserializable:
            def __repr__(self):
                return "Unserializable()"

        obj = Unserializable()
        result = trace_handler._maybe_json(obj)
        assert isinstance(result, str)


class TestContentOrRef:
    def test_returns_content_when_include_content_true(self, trace_handler, mock_exporter):
        mock_exporter.include_content = True
        result = trace_handler._content_or_ref("actual content", "ref:x=1")
        assert result == "actual content"

    def test_returns_ref_when_include_content_false(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        result = trace_handler._content_or_ref("actual content", "ref:x=1")
        assert result == "ref:x=1"

    def test_dict_serialised_when_include_content_true(self, trace_handler, mock_exporter):
        mock_exporter.include_content = True
        result = trace_handler._content_or_ref({"key": "val"}, "ref:x=1")
        assert result == json.dumps({"key": "val"})

    def test_none_returns_none_when_include_content_true(self, trace_handler, mock_exporter):
        mock_exporter.include_content = True
        result = trace_handler._content_or_ref(None, "ref:x=1")
        assert result is None


# ---------------------------------------------------------------------------
# trace() dispatcher
# ---------------------------------------------------------------------------


class TestTraceDispatcher:
    def test_dispatches_workflow_trace(self, trace_handler):
        with patch.object(trace_handler, "_workflow_trace") as mock_method:
            info = make_workflow_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_message_trace(self, trace_handler):
        with patch.object(trace_handler, "_message_trace") as mock_method:
            info = make_message_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_tool_trace(self, trace_handler):
        with patch.object(trace_handler, "_tool_trace") as mock_method:
            info = make_tool_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_draft_node_execution_trace(self, trace_handler):
        with patch.object(trace_handler, "_draft_node_execution_trace") as mock_method:
            info = make_draft_node_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_node_execution_trace(self, trace_handler):
        with patch.object(trace_handler, "_node_execution_trace") as mock_method:
            info = make_node_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_moderation_trace(self, trace_handler):
        with patch.object(trace_handler, "_moderation_trace") as mock_method:
            info = make_moderation_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_suggested_question_trace(self, trace_handler):
        with patch.object(trace_handler, "_suggested_question_trace") as mock_method:
            info = make_suggested_question_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_dataset_retrieval_trace(self, trace_handler):
        with patch.object(trace_handler, "_dataset_retrieval_trace") as mock_method:
            info = make_dataset_retrieval_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_generate_name_trace(self, trace_handler):
        with patch.object(trace_handler, "_generate_name_trace") as mock_method:
            info = make_generate_name_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_dispatches_prompt_generation_trace(self, trace_handler):
        with patch.object(trace_handler, "_prompt_generation_trace") as mock_method:
            info = make_prompt_generation_info()
            trace_handler.trace(info)
            mock_method.assert_called_once_with(info)

    def test_draft_node_dispatched_before_node(self, trace_handler):
        """DraftNodeExecutionTrace is a subclass of WorkflowNodeTraceInfo;
        it must be dispatched to _draft_node_execution_trace, not _node_execution_trace."""
        with (
            patch.object(trace_handler, "_draft_node_execution_trace") as mock_draft,
            patch.object(trace_handler, "_node_execution_trace") as mock_node,
        ):
            info = make_draft_node_info()
            trace_handler.trace(info)
            mock_draft.assert_called_once_with(info)
            mock_node.assert_not_called()


# ---------------------------------------------------------------------------
# _workflow_trace
# ---------------------------------------------------------------------------


class TestWorkflowTrace:
    def test_emits_correct_span_attributes(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log:
            info = make_workflow_info()
            trace_handler._workflow_trace(info)

        mock_exporter.export_span.assert_called_once()
        span_call = mock_exporter.export_span.call_args
        assert span_call[0][0] == EnterpriseTelemetrySpan.WORKFLOW_RUN
        attrs = span_call[0][1]
        assert attrs["dify.workflow.run_id"] == "run-001"
        assert attrs["dify.workflow.id"] == "wf-001"
        assert attrs["dify.tenant_id"] == "tenant-abc"
        assert attrs["dify.workflow.status"] == "succeeded"
        assert attrs["gen_ai.usage.total_tokens"] == 100

    def test_span_timing_passed_correctly(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_workflow_info()
            trace_handler._workflow_trace(info)

        span_call = mock_exporter.export_span.call_args
        assert span_call[1]["start_time"] == _T0
        assert span_call[1]["end_time"] == _T1

    def test_emits_companion_log_with_event_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log:
            trace_handler._workflow_trace(make_workflow_info())

        mock_log.assert_called_once()
        assert mock_log.call_args[1]["event_name"] == EnterpriseTelemetryEvent.WORKFLOW_RUN
        assert mock_log.call_args[1]["tenant_id"] == "tenant-abc"

    def test_companion_log_includes_content_when_enabled(self, trace_handler, mock_exporter):
        mock_exporter.include_content = True
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log:
            trace_handler._workflow_trace(make_workflow_info())

        log_attrs = mock_log.call_args[1]["attributes"]
        assert log_attrs["dify.workflow.inputs"] == json.dumps({"query": "hello"})
        assert log_attrs["dify.workflow.outputs"] == json.dumps({"answer": "world"})

    def test_companion_log_uses_ref_when_content_disabled(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log:
            trace_handler._workflow_trace(make_workflow_info())

        log_attrs = mock_log.call_args[1]["attributes"]
        assert log_attrs["dify.workflow.inputs"].startswith("ref:workflow_run_id=")
        assert log_attrs["dify.workflow.outputs"].startswith("ref:workflow_run_id=")

    def test_increments_token_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._workflow_trace(make_workflow_info())

        token_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.TOKENS
        ]
        assert len(token_calls) == 1
        assert token_calls[0][0][1] == 100

    def test_increments_input_and_output_token_counters(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._workflow_trace(make_workflow_info())

        all_calls = mock_exporter.increment_counter.call_args_list
        counter_names = [c[0][0] for c in all_calls]
        assert EnterpriseTelemetryCounter.INPUT_TOKENS in counter_names
        assert EnterpriseTelemetryCounter.OUTPUT_TOKENS in counter_names

    def test_no_input_token_counter_when_prompt_tokens_zero(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_workflow_info(prompt_tokens=0)
            trace_handler._workflow_trace(info)

        all_calls = mock_exporter.increment_counter.call_args_list
        counter_names = [c[0][0] for c in all_calls]
        assert EnterpriseTelemetryCounter.INPUT_TOKENS not in counter_names

    def test_records_workflow_duration_histogram(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._workflow_trace(make_workflow_info())

        mock_exporter.record_histogram.assert_called_once()
        hist_call = mock_exporter.record_histogram.call_args
        assert hist_call[0][0] == EnterpriseTelemetryHistogram.WORKFLOW_DURATION
        assert hist_call[0][1] == pytest.approx(5.0)

    def test_duration_falls_back_to_elapsed_time_when_timestamps_missing(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_workflow_info(start_time=None, end_time=None, workflow_run_elapsed_time=7.3)
            trace_handler._workflow_trace(info)

        hist_call = mock_exporter.record_histogram.call_args
        assert hist_call[0][1] == pytest.approx(7.3)

    def test_duration_defaults_to_zero_when_no_timing(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_workflow_info(start_time=None, end_time=None, workflow_run_elapsed_time=0)
            trace_handler._workflow_trace(info)

        hist_call = mock_exporter.record_histogram.call_args
        assert hist_call[0][1] == pytest.approx(0.0)

    def test_error_path_increments_error_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_workflow_info(error="Something went wrong", workflow_run_status="failed")
            trace_handler._workflow_trace(info)

        error_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.ERRORS
        ]
        assert len(error_calls) == 1

    def test_no_error_counter_on_success(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._workflow_trace(make_workflow_info())

        error_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.ERRORS
        ]
        assert len(error_calls) == 0

    def test_parent_trace_context_injected_into_span_attrs(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_workflow_info(
                metadata={
                    "app_id": "app-001",
                    "tenant_id": "tenant-abc",
                    "parent_trace_context": {
                        "trace_id": "outer-trace",
                        "parent_node_execution_id": "outer-ne-001",
                        "parent_workflow_run_id": "outer-run-001",
                        "parent_app_id": "outer-app-001",
                    },
                }
            )
            trace_handler._workflow_trace(info)

        attrs = mock_exporter.export_span.call_args[0][1]
        assert attrs["dify.parent.trace_id"] == "outer-trace"
        assert attrs["dify.parent.node.execution_id"] == "outer-ne-001"
        assert attrs["dify.parent.workflow.run_id"] == "outer-run-001"
        assert attrs["dify.parent.app.id"] == "outer-app-001"


# ---------------------------------------------------------------------------
# _node_execution_trace / _emit_node_execution_trace
# ---------------------------------------------------------------------------


class TestNodeExecutionTrace:
    def test_emits_span_with_node_execution_span_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._node_execution_trace(make_node_info())

        span_call = mock_exporter.export_span.call_args
        assert span_call[0][0] == EnterpriseTelemetrySpan.NODE_EXECUTION

    def test_span_contains_core_node_attributes(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._node_execution_trace(make_node_info())

        attrs = mock_exporter.export_span.call_args[0][1]
        assert attrs["dify.node.execution_id"] == "ne-001"
        assert attrs["dify.node.id"] == "node-001"
        assert attrs["dify.node.type"] == "llm"
        assert attrs["dify.node.status"] == "succeeded"
        assert attrs["gen_ai.request.model"] == "gpt-4"
        assert attrs["gen_ai.provider.name"] == "openai"

    def test_increments_token_counters_when_tokens_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._node_execution_trace(make_node_info())

        counter_names = [c[0][0] for c in mock_exporter.increment_counter.call_args_list]
        assert EnterpriseTelemetryCounter.TOKENS in counter_names
        assert EnterpriseTelemetryCounter.INPUT_TOKENS in counter_names
        assert EnterpriseTelemetryCounter.OUTPUT_TOKENS in counter_names

    def test_no_token_counters_when_total_tokens_zero(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._node_execution_trace(make_node_info(total_tokens=0))

        counter_names = [c[0][0] for c in mock_exporter.increment_counter.call_args_list]
        assert EnterpriseTelemetryCounter.TOKENS not in counter_names
        assert EnterpriseTelemetryCounter.INPUT_TOKENS not in counter_names

    def test_records_node_duration_histogram(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._node_execution_trace(make_node_info())

        hist_call = mock_exporter.record_histogram.call_args
        assert hist_call[0][0] == EnterpriseTelemetryHistogram.NODE_DURATION
        assert hist_call[0][1] == pytest.approx(2.5)

    def test_error_path_increments_error_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._node_execution_trace(make_node_info(error="Node failed", status="failed"))

        error_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.ERRORS
        ]
        assert len(error_calls) == 1

    def test_emits_companion_log_with_span_name_as_event(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log:
            trace_handler._node_execution_trace(make_node_info())

        mock_log.assert_called_once()
        assert mock_log.call_args[1]["event_name"] == EnterpriseTelemetrySpan.NODE_EXECUTION.value

    def test_plugin_name_added_to_duration_labels_for_tool_node(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_node_info(
                node_type="tool",
                metadata={
                    "app_id": "app-001",
                    "tenant_id": "tenant-abc",
                    "plugin_name": "my-plugin",
                },
            )
            trace_handler._node_execution_trace(info)

        hist_call = mock_exporter.record_histogram.call_args
        duration_labels = hist_call[0][2]
        assert duration_labels.get("plugin_name") == "my-plugin"

    def test_plugin_name_not_added_for_non_tool_node(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_node_info(
                node_type="llm",
                metadata={
                    "app_id": "app-001",
                    "tenant_id": "tenant-abc",
                    "plugin_name": "my-plugin",
                },
            )
            trace_handler._node_execution_trace(info)

        hist_call = mock_exporter.record_histogram.call_args
        duration_labels = hist_call[0][2]
        assert "plugin_name" not in duration_labels

    def test_companion_log_inputs_use_ref_when_content_disabled(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log:
            trace_handler._node_execution_trace(
                make_node_info(node_inputs={"prompt": "hello"}, node_outputs={"text": "world"})
            )

        log_attrs = mock_log.call_args[1]["attributes"]
        assert log_attrs["dify.node.inputs"].startswith("ref:node_execution_id=")
        assert log_attrs["dify.node.outputs"].startswith("ref:node_execution_id=")


# ---------------------------------------------------------------------------
# _draft_node_execution_trace
# ---------------------------------------------------------------------------


class TestDraftNodeExecutionTrace:
    def test_uses_draft_span_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            trace_handler._draft_node_execution_trace(make_draft_node_info())

        span_call = mock_exporter.export_span.call_args
        assert span_call[0][0] == EnterpriseTelemetrySpan.DRAFT_NODE_EXECUTION

    def test_correlation_id_is_node_execution_id(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_draft_node_info()
            trace_handler._draft_node_execution_trace(info)

        span_call = mock_exporter.export_span.call_args
        assert span_call[1]["correlation_id"] == "ne-draft-001"

    def test_trace_correlation_override_is_workflow_run_id(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log"):
            info = make_draft_node_info()
            trace_handler._draft_node_execution_trace(info)

        span_call = mock_exporter.export_span.call_args
        assert span_call[1]["trace_correlation_override"] == "run-draft-001"

    def test_companion_log_uses_draft_span_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_telemetry_log") as mock_log:
            trace_handler._draft_node_execution_trace(make_draft_node_info())

        assert mock_log.call_args[1]["event_name"] == EnterpriseTelemetrySpan.DRAFT_NODE_EXECUTION.value


# ---------------------------------------------------------------------------
# _message_trace
# ---------------------------------------------------------------------------


class TestMessageTrace:
    def test_emits_event_with_correct_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._message_trace(make_message_info())

        mock_emit.assert_called_once()
        assert mock_emit.call_args[1]["event_name"] == EnterpriseTelemetryEvent.MESSAGE_RUN

    def test_emits_correct_tenant_and_user(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._message_trace(make_message_info())

        assert mock_emit.call_args[1]["tenant_id"] == "tenant-abc"

    def test_duration_computed_from_timestamps(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._message_trace(make_message_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.message.duration"] == pytest.approx(5.0)

    def test_no_duration_when_timestamps_missing(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._message_trace(make_message_info(start_time=None, end_time=None))

        attrs = mock_emit.call_args[1]["attributes"]
        assert "dify.message.duration" not in attrs

    def test_records_duration_histogram_when_timestamps_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._message_trace(make_message_info())

        hist_calls = [
            c
            for c in mock_exporter.record_histogram.call_args_list
            if c[0][0] == EnterpriseTelemetryHistogram.MESSAGE_DURATION
        ]
        assert len(hist_calls) == 1
        assert hist_calls[0][0][1] == pytest.approx(5.0)

    def test_no_duration_histogram_when_timestamps_missing(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._message_trace(make_message_info(start_time=None, end_time=None))

        hist_names = [c[0][0] for c in mock_exporter.record_histogram.call_args_list]
        assert EnterpriseTelemetryHistogram.MESSAGE_DURATION not in hist_names

    def test_records_ttft_histogram_when_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._message_trace(make_message_info(gen_ai_server_time_to_first_token=0.42))

        ttft_calls = [
            c
            for c in mock_exporter.record_histogram.call_args_list
            if c[0][0] == EnterpriseTelemetryHistogram.MESSAGE_TTFT
        ]
        assert len(ttft_calls) == 1
        assert ttft_calls[0][0][1] == pytest.approx(0.42)

    def test_no_ttft_histogram_when_not_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._message_trace(make_message_info(gen_ai_server_time_to_first_token=None))

        hist_names = [c[0][0] for c in mock_exporter.record_histogram.call_args_list]
        assert EnterpriseTelemetryHistogram.MESSAGE_TTFT not in hist_names

    def test_increments_token_counters(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._message_trace(make_message_info())

        counter_names = [c[0][0] for c in mock_exporter.increment_counter.call_args_list]
        assert EnterpriseTelemetryCounter.TOKENS in counter_names
        assert EnterpriseTelemetryCounter.INPUT_TOKENS in counter_names
        assert EnterpriseTelemetryCounter.OUTPUT_TOKENS in counter_names

    def test_error_path_increments_error_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._message_trace(make_message_info(error="LLM failed"))

        error_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.ERRORS
        ]
        assert len(error_calls) == 1

    def test_inputs_and_outputs_gated_by_include_content(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._message_trace(make_message_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.message.inputs"].startswith("ref:message_id=")
        assert attrs["dify.message.outputs"].startswith("ref:message_id=")


# ---------------------------------------------------------------------------
# _tool_trace
# ---------------------------------------------------------------------------


class TestToolTrace:
    def test_emits_event_with_correct_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._tool_trace(make_tool_info())

        assert mock_emit.call_args[1]["event_name"] == EnterpriseTelemetryEvent.TOOL_EXECUTION

    def test_status_is_succeeded_on_success(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._tool_trace(make_tool_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.tool.status"] == "succeeded"

    def test_status_is_failed_on_error(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._tool_trace(make_tool_info(error="Tool error"))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.tool.status"] == "failed"

    def test_records_tool_duration_histogram(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._tool_trace(make_tool_info())

        hist_call = mock_exporter.record_histogram.call_args
        assert hist_call[0][0] == EnterpriseTelemetryHistogram.TOOL_DURATION
        assert hist_call[0][1] == pytest.approx(1.5)

    def test_error_increments_error_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._tool_trace(make_tool_info(error="Tool crashed"))

        error_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.ERRORS
        ]
        assert len(error_calls) == 1

    def test_inputs_and_outputs_gated_by_include_content(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._tool_trace(make_tool_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.tool.inputs"].startswith("ref:message_id=")
        assert attrs["dify.tool.outputs"].startswith("ref:message_id=")

    def test_inputs_present_when_include_content_true(self, trace_handler, mock_exporter):
        mock_exporter.include_content = True
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._tool_trace(make_tool_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.tool.inputs"] == json.dumps({"query": "test"})
        assert attrs["dify.tool.outputs"] == "search results"

    def test_increments_requests_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._tool_trace(make_tool_info())

        request_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.REQUESTS
        ]
        assert len(request_calls) == 1
        assert request_calls[0][0][2]["type"] == "tool"


# ---------------------------------------------------------------------------
# _moderation_trace
# ---------------------------------------------------------------------------


class TestModerationTrace:
    def test_emits_event_with_correct_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._moderation_trace(make_moderation_info())

        assert mock_emit.call_args[1]["event_name"] == EnterpriseTelemetryEvent.MODERATION_CHECK

    def test_flagged_true_sets_attribute(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._moderation_trace(make_moderation_info(flagged=True))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.moderation.flagged"] is True

    def test_flagged_false_sets_attribute(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._moderation_trace(make_moderation_info(flagged=False))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.moderation.flagged"] is False

    def test_query_gated_by_include_content(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._moderation_trace(make_moderation_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.moderation.query"].startswith("ref:message_id=")

    def test_query_present_when_include_content_true(self, trace_handler, mock_exporter):
        mock_exporter.include_content = True
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._moderation_trace(make_moderation_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.moderation.query"] == "is this ok?"

    def test_increments_requests_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._moderation_trace(make_moderation_info())

        request_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.REQUESTS
        ]
        assert len(request_calls) == 1
        assert request_calls[0][0][2]["type"] == "moderation"


# ---------------------------------------------------------------------------
# _suggested_question_trace
# ---------------------------------------------------------------------------


class TestSuggestedQuestionTrace:
    def test_emits_event_with_correct_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._suggested_question_trace(make_suggested_question_info())

        assert mock_emit.call_args[1]["event_name"] == EnterpriseTelemetryEvent.SUGGESTED_QUESTION_GENERATION

    def test_duration_computed_from_timestamps(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._suggested_question_trace(make_suggested_question_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.suggested_question.duration"] == pytest.approx(5.0)

    def test_duration_is_none_when_timestamps_missing(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._suggested_question_trace(make_suggested_question_info(start_time=None, end_time=None))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.suggested_question.duration"] is None

    def test_status_is_failed_when_error_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._suggested_question_trace(make_suggested_question_info(error="Generation failed"))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.suggested_question.status"] == "failed"

    def test_status_falls_back_to_succeeded_when_no_error(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._suggested_question_trace(make_suggested_question_info(status=None, error=None))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.suggested_question.status"] == "succeeded"

    def test_question_count_attribute(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._suggested_question_trace(make_suggested_question_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.suggested_question.count"] == 2

    def test_questions_gated_by_include_content(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._suggested_question_trace(make_suggested_question_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.suggested_question.questions"].startswith("ref:message_id=")

    def test_increments_requests_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._suggested_question_trace(make_suggested_question_info())

        request_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.REQUESTS
        ]
        assert len(request_calls) == 1
        assert request_calls[0][0][2]["type"] == "suggested_question"


# ---------------------------------------------------------------------------
# _dataset_retrieval_trace
# ---------------------------------------------------------------------------


class TestDatasetRetrievalTrace:
    def test_emits_event_with_correct_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info())

        assert mock_emit.call_args[1]["event_name"] == EnterpriseTelemetryEvent.DATASET_RETRIEVAL

    def test_document_count_attribute(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.retrieval.document_count"] == 1

    def test_dataset_ids_extracted(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert "ds-001" in attrs["dify.dataset.id"]

    def test_empty_documents_has_zero_count(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info(documents=[]))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.retrieval.document_count"] == 0

    def test_status_succeeded_when_no_error(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.retrieval.status"] == "succeeded"

    def test_status_failed_when_error_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info(error="DB error"))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.retrieval.status"] == "failed"

    def test_embedding_model_attributes_set_when_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert "dify.dataset.embedding_providers" in attrs
        assert "dify.dataset.embedding_models" in attrs

    def test_no_embedding_model_attributes_when_not_provided(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(
                make_dataset_retrieval_info(metadata={"app_id": "app-001", "tenant_id": "tenant-abc"})
            )

        attrs = mock_emit.call_args[1]["attributes"]
        assert "dify.dataset.embedding_providers" not in attrs
        assert "dify.dataset.embedding_models" not in attrs

    def test_rerank_attributes_set_when_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(
                make_dataset_retrieval_info(
                    metadata={
                        "app_id": "app-001",
                        "tenant_id": "tenant-abc",
                        "rerank_model_provider": "cohere",
                        "rerank_model_name": "rerank-english",
                    }
                )
            )

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.retrieval.rerank_provider"] == "cohere"
        assert attrs["dify.retrieval.rerank_model"] == "rerank-english"

    def test_no_rerank_attributes_when_not_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(
                make_dataset_retrieval_info(metadata={"app_id": "app-001", "tenant_id": "tenant-abc"})
            )

        attrs = mock_emit.call_args[1]["attributes"]
        assert "dify.retrieval.rerank_provider" not in attrs
        assert "dify.retrieval.rerank_model" not in attrs

    def test_dataset_retrieval_counter_incremented_per_dataset(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info())

        ds_calls = [
            c
            for c in mock_exporter.increment_counter.call_args_list
            if c[0][0] == EnterpriseTelemetryCounter.DATASET_RETRIEVALS
        ]
        assert len(ds_calls) == 1
        assert ds_calls[0][0][2]["dataset_id"] == "ds-001"

    def test_no_dataset_retrieval_counter_when_no_documents(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info(documents=[]))

        ds_calls = [
            c
            for c in mock_exporter.increment_counter.call_args_list
            if c[0][0] == EnterpriseTelemetryCounter.DATASET_RETRIEVALS
        ]
        assert len(ds_calls) == 0

    def test_query_gated_by_include_content(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._dataset_retrieval_trace(make_dataset_retrieval_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.retrieval.query"].startswith("ref:message_id=")


# ---------------------------------------------------------------------------
# _generate_name_trace
# ---------------------------------------------------------------------------


class TestGenerateNameTrace:
    def test_emits_event_with_correct_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._generate_name_trace(make_generate_name_info())

        assert mock_emit.call_args[1]["event_name"] == EnterpriseTelemetryEvent.GENERATE_NAME_EXECUTION

    def test_duration_computed_from_timestamps(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._generate_name_trace(make_generate_name_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.generate_name.duration"] == pytest.approx(5.0)

    def test_no_duration_when_timestamps_missing(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._generate_name_trace(make_generate_name_info(start_time=None, end_time=None))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.generate_name.duration"] is None

    def test_status_succeeded_on_success(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._generate_name_trace(make_generate_name_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.generate_name.status"] == "succeeded"

    def test_status_failed_when_metadata_has_error(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._generate_name_trace(
                make_generate_name_info(
                    metadata={
                        "app_id": "app-001",
                        "tenant_id": "tenant-abc",
                        "error": "Name generation failed",
                    }
                )
            )

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.generate_name.status"] == "failed"

    def test_inputs_and_outputs_gated_by_include_content(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._generate_name_trace(make_generate_name_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.generate_name.inputs"].startswith("ref:conversation_id=")
        assert attrs["dify.generate_name.outputs"].startswith("ref:conversation_id=")

    def test_increments_requests_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._generate_name_trace(make_generate_name_info())

        request_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.REQUESTS
        ]
        assert len(request_calls) == 1
        assert request_calls[0][0][2]["type"] == "generate_name"


# ---------------------------------------------------------------------------
# _prompt_generation_trace
# ---------------------------------------------------------------------------


class TestPromptGenerationTrace:
    def test_emits_event_with_correct_name(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._prompt_generation_trace(make_prompt_generation_info())

        assert mock_emit.call_args[1]["event_name"] == EnterpriseTelemetryEvent.PROMPT_GENERATION_EXECUTION

    def test_status_succeeded_on_success(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._prompt_generation_trace(make_prompt_generation_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.prompt_generation.status"] == "succeeded"

    def test_status_failed_when_error_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._prompt_generation_trace(make_prompt_generation_info(error="Generation error"))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.prompt_generation.status"] == "failed"

    def test_token_counters_incremented(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._prompt_generation_trace(make_prompt_generation_info())

        counter_names = [c[0][0] for c in mock_exporter.increment_counter.call_args_list]
        assert EnterpriseTelemetryCounter.TOKENS in counter_names
        assert EnterpriseTelemetryCounter.INPUT_TOKENS in counter_names
        assert EnterpriseTelemetryCounter.OUTPUT_TOKENS in counter_names

    def test_records_duration_histogram(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._prompt_generation_trace(make_prompt_generation_info())

        hist_calls = [
            c
            for c in mock_exporter.record_histogram.call_args_list
            if c[0][0] == EnterpriseTelemetryHistogram.PROMPT_GENERATION_DURATION
        ]
        assert len(hist_calls) == 1
        assert hist_calls[0][0][1] == pytest.approx(3.2)

    def test_total_price_attribute_set_when_present(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._prompt_generation_trace(make_prompt_generation_info(total_price=0.05, currency="USD"))

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.prompt_generation.total_price"] == pytest.approx(0.05)
        assert attrs["dify.prompt_generation.currency"] == "USD"

    def test_no_total_price_attribute_when_none(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._prompt_generation_trace(make_prompt_generation_info(total_price=None))

        attrs = mock_emit.call_args[1]["attributes"]
        assert "dify.prompt_generation.total_price" not in attrs

    def test_error_increments_error_counter(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._prompt_generation_trace(make_prompt_generation_info(error="Prompt failed"))

        error_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.ERRORS
        ]
        assert len(error_calls) == 1

    def test_no_error_counter_on_success(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._prompt_generation_trace(make_prompt_generation_info())

        error_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.ERRORS
        ]
        assert len(error_calls) == 0

    def test_instruction_gated_by_include_content(self, trace_handler, mock_exporter):
        mock_exporter.include_content = False
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._prompt_generation_trace(make_prompt_generation_info())

        attrs = mock_emit.call_args[1]["attributes"]
        assert attrs["dify.prompt_generation.instruction"].startswith("ref:trace_id=")

    def test_operation_type_label_used_in_token_counters(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event"):
            trace_handler._prompt_generation_trace(make_prompt_generation_info(operation_type="code_generate"))

        token_calls = [
            c for c in mock_exporter.increment_counter.call_args_list if c[0][0] == EnterpriseTelemetryCounter.TOKENS
        ]
        assert len(token_calls) == 1
        assert token_calls[0][0][2]["operation_type"] == "code_generate"

    def test_emits_correct_tenant_id(self, trace_handler, mock_exporter):
        with patch("enterprise.telemetry.enterprise_trace.emit_metric_only_event") as mock_emit:
            trace_handler._prompt_generation_trace(make_prompt_generation_info())

        assert mock_emit.call_args[1]["tenant_id"] == "tenant-abc"
