import gc
import logging
import warnings
from unittest.mock import AsyncMock, MagicMock, patch

import pytest
from dify_trace_tencent.config import TencentConfig
from dify_trace_tencent.tencent_trace import TencentDataTrace

from core.ops.entities.trace_entity import (
    DatasetRetrievalTraceInfo,
    GenerateNameTraceInfo,
    MessageTraceInfo,
    ModerationTraceInfo,
    SuggestedQuestionTraceInfo,
    ToolTraceInfo,
    WorkflowTraceInfo,
)
from graphon.entities import WorkflowNodeExecution
from graphon.enums import BuiltinNodeTypes
from models import Account, App, TenantAccountJoin

logger = logging.getLogger(__name__)


@pytest.fixture
def tencent_config():
    return TencentConfig(service_name="test-service", endpoint="https://test-endpoint", token="test-token")


@pytest.fixture
def mock_trace_client():
    with patch("dify_trace_tencent.tencent_trace.TencentTraceClient") as mock:
        yield mock


@pytest.fixture
def mock_span_builder():
    with patch("dify_trace_tencent.tencent_trace.TencentSpanBuilder") as mock:
        yield mock


@pytest.fixture
def mock_trace_utils():
    with patch("dify_trace_tencent.tencent_trace.TencentTraceUtils") as mock:
        yield mock


@pytest.fixture
def tencent_data_trace(tencent_config, mock_trace_client):
    return TencentDataTrace(tencent_config)


class TestTencentDataTrace:
    def test_init(self, tencent_config, mock_trace_client):
        trace = TencentDataTrace(tencent_config)
        mock_trace_client.assert_called_once_with(
            service_name=tencent_config.service_name,
            endpoint=tencent_config.endpoint,
            token=tencent_config.token,
            metrics_export_interval_sec=5,
        )
        assert trace.trace_client == mock_trace_client.return_value

    def test_trace_dispatch(self, tencent_data_trace):
        methods = [
            (
                WorkflowTraceInfo(
                    workflow_id="wf",
                    tenant_id="t",
                    workflow_run_id="run",
                    workflow_run_elapsed_time=1.0,
                    workflow_run_status="s",
                    workflow_run_inputs={},
                    workflow_run_outputs={},
                    workflow_run_version="v",
                    total_tokens=0,
                    file_list=[],
                    query="",
                    metadata={},
                ),
                "workflow_trace",
            ),
            (
                MessageTraceInfo(
                    message_id="msg",
                    message_data={},
                    inputs={},
                    outputs={},
                    start_time=None,
                    end_time=None,
                    conversation_mode="chat",
                    conversation_model="gpt-3.5-turbo",
                    message_tokens=0,
                    answer_tokens=0,
                    total_tokens=0,
                    metadata={},
                ),
                "message_trace",
            ),
            (
                ModerationTraceInfo(
                    flagged=False, action="a", preset_response="p", query="q", metadata={}, message_id="m"
                ),
                None,
            ),  # Pass
            (
                SuggestedQuestionTraceInfo(
                    suggested_question=[],
                    level="l",
                    total_tokens=0,
                    metadata={},
                    message_id="m",
                    message_data={},
                    inputs={},
                    start_time=None,
                    end_time=None,
                ),
                "suggested_question_trace",
            ),
            (
                DatasetRetrievalTraceInfo(
                    metadata={},
                    message_id="m",
                    message_data={},
                    inputs={},
                    documents=[],
                    start_time=None,
                    end_time=None,
                ),
                "dataset_retrieval_trace",
            ),
            (
                ToolTraceInfo(
                    tool_name="t",
                    tool_inputs={},
                    tool_outputs="",
                    tool_config={},
                    tool_parameters={},
                    time_cost=0,
                    metadata={},
                    message_id="m",
                    inputs={},
                    outputs={},
                    start_time=None,
                    end_time=None,
                ),
                "tool_trace",
            ),
            (
                GenerateNameTraceInfo(
                    tenant_id="t", metadata={}, message_id="m", inputs={}, outputs={}, start_time=None, end_time=None
                ),
                None,
            ),  # Pass
        ]

        for trace_info, method_name in methods:
            if method_name:
                with patch.object(tencent_data_trace, method_name) as mock_method:
                    tencent_data_trace.trace(trace_info)
                    mock_method.assert_called_once_with(trace_info)
            else:
                tencent_data_trace.trace(trace_info)

    def test_api_check(self, tencent_data_trace):
        tencent_data_trace.trace_client.api_check.return_value = True
        assert tencent_data_trace.api_check() is True
        tencent_data_trace.trace_client.api_check.assert_called_once()

    def test_get_project_url(self, tencent_data_trace):
        tencent_data_trace.trace_client.get_project_url.return_value = "http://url"
        assert tencent_data_trace.get_project_url() == "http://url"
        tencent_data_trace.trace_client.get_project_url.assert_called_once()

    def test_workflow_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.workflow_run_id = "run-id"
        trace_info.trace_id = "parent-trace-id"

        mock_trace_utils.convert_to_trace_id.return_value = 123
        mock_trace_utils.create_link.return_value = "link"

        with patch.object(tencent_data_trace, "_get_user_id", return_value="user-1"):
            with patch.object(tencent_data_trace, "_process_workflow_nodes") as mock_proc:
                with patch.object(tencent_data_trace, "_record_workflow_trace_duration") as mock_dur:
                    mock_span_builder.build_workflow_spans.return_value = [MagicMock(), MagicMock()]

                    tencent_data_trace.workflow_trace(trace_info)

                    mock_trace_utils.convert_to_trace_id.assert_called_once_with("run-id")
                    mock_trace_utils.create_link.assert_called_once_with("parent-trace-id")
                    mock_span_builder.build_workflow_spans.assert_called_once()
                    assert tencent_data_trace.trace_client.add_span.call_count == 2
                    mock_proc.assert_called_once_with(trace_info, 123)
                    mock_dur.assert_called_once_with(trace_info)

    def test_workflow_trace_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.workflow_run_id = "run-id"

        with patch(
            "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_trace_id", side_effect=Exception("error")
        ):
            with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                tencent_data_trace.workflow_trace(trace_info)
                mock_log.assert_called_once_with("[Tencent APM] Failed to process workflow trace")

    def test_message_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder):
        trace_info = MagicMock(spec=MessageTraceInfo)
        trace_info.message_id = "msg-id"
        trace_info.trace_id = "parent-trace-id"

        mock_trace_utils.convert_to_trace_id.return_value = 123
        mock_trace_utils.create_link.return_value = "link"

        with patch.object(tencent_data_trace, "_get_user_id", return_value="user-1"):
            with patch.object(tencent_data_trace, "_record_message_llm_metrics") as mock_metrics:
                with patch.object(tencent_data_trace, "_record_message_trace_duration") as mock_dur:
                    mock_span_builder.build_message_span.return_value = MagicMock()

                    tencent_data_trace.message_trace(trace_info)

                    mock_trace_utils.convert_to_trace_id.assert_called_once_with("msg-id")
                    mock_trace_utils.create_link.assert_called_once_with("parent-trace-id")
                    mock_span_builder.build_message_span.assert_called_once()
                    tencent_data_trace.trace_client.add_span.assert_called_once()
                    mock_metrics.assert_called_once_with(trace_info)
                    mock_dur.assert_called_once_with(trace_info)

    def test_message_trace_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)

        with patch(
            "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_trace_id", side_effect=Exception("error")
        ):
            with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                tencent_data_trace.message_trace(trace_info)
                mock_log.assert_called_once_with("[Tencent APM] Failed to process message trace")

    def test_tool_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder):
        trace_info = MagicMock(spec=ToolTraceInfo)
        trace_info.message_id = "msg-id"

        mock_trace_utils.convert_to_span_id.return_value = 456
        mock_trace_utils.convert_to_trace_id.return_value = 123

        tencent_data_trace.tool_trace(trace_info)

        mock_trace_utils.convert_to_span_id.assert_called_once_with("msg-id", "message")
        mock_trace_utils.convert_to_trace_id.assert_called_once_with("msg-id")
        mock_span_builder.build_tool_span.assert_called_once_with(trace_info, 123, 456)
        tencent_data_trace.trace_client.add_span.assert_called_once()

    def test_tool_trace_no_msg_id(self, tencent_data_trace):
        trace_info = MagicMock(spec=ToolTraceInfo)
        trace_info.message_id = None

        tencent_data_trace.tool_trace(trace_info)
        tencent_data_trace.trace_client.add_span.assert_not_called()

    def test_tool_trace_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=ToolTraceInfo)
        trace_info.message_id = "msg-id"

        with patch(
            "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_span_id", side_effect=Exception("error")
        ):
            with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                tencent_data_trace.tool_trace(trace_info)
                mock_log.assert_called_once_with("[Tencent APM] Failed to process tool trace")

    def test_dataset_retrieval_trace(self, tencent_data_trace, mock_trace_utils, mock_span_builder):
        trace_info = MagicMock(spec=DatasetRetrievalTraceInfo)
        trace_info.message_id = "msg-id"

        mock_trace_utils.convert_to_span_id.return_value = 456
        mock_trace_utils.convert_to_trace_id.return_value = 123

        tencent_data_trace.dataset_retrieval_trace(trace_info)

        mock_trace_utils.convert_to_span_id.assert_called_once_with("msg-id", "message")
        mock_trace_utils.convert_to_trace_id.assert_called_once_with("msg-id")
        mock_span_builder.build_retrieval_span.assert_called_once_with(trace_info, 123, 456)
        tencent_data_trace.trace_client.add_span.assert_called_once()

    def test_dataset_retrieval_trace_no_msg_id(self, tencent_data_trace):
        trace_info = MagicMock(spec=DatasetRetrievalTraceInfo)
        trace_info.message_id = None

        tencent_data_trace.dataset_retrieval_trace(trace_info)
        tencent_data_trace.trace_client.add_span.assert_not_called()

    def test_dataset_retrieval_trace_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=DatasetRetrievalTraceInfo)
        trace_info.message_id = "msg-id"

        with patch(
            "dify_trace_tencent.tencent_trace.TencentTraceUtils.convert_to_span_id", side_effect=Exception("error")
        ):
            with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                tencent_data_trace.dataset_retrieval_trace(trace_info)
                mock_log.assert_called_once_with("[Tencent APM] Failed to process dataset retrieval trace")

    def test_suggested_question_trace(self, tencent_data_trace):
        trace_info = MagicMock(spec=SuggestedQuestionTraceInfo)
        with patch("dify_trace_tencent.tencent_trace.logger.info") as mock_log:
            tencent_data_trace.suggested_question_trace(trace_info)
            mock_log.assert_called_once_with("[Tencent APM] Processing suggested question trace")

    def test_suggested_question_trace_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=SuggestedQuestionTraceInfo)
        with patch("dify_trace_tencent.tencent_trace.logger.info", side_effect=Exception("error")):
            with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                tencent_data_trace.suggested_question_trace(trace_info)
                mock_log.assert_called_once_with("[Tencent APM] Failed to process suggested question trace")

    def test_process_workflow_nodes(self, tencent_data_trace, mock_trace_utils):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.workflow_run_id = "run-id"
        mock_trace_utils.convert_to_span_id.return_value = 111

        node1 = MagicMock(spec=WorkflowNodeExecution)
        node1.id = "n1"
        node1.node_type = BuiltinNodeTypes.LLM
        node2 = MagicMock(spec=WorkflowNodeExecution)
        node2.id = "n2"
        node2.node_type = BuiltinNodeTypes.TOOL

        with patch.object(tencent_data_trace, "_get_workflow_node_executions", return_value=[node1, node2]):
            with patch.object(tencent_data_trace, "_build_workflow_node_span", side_effect=["span1", "span2"]):
                with patch.object(tencent_data_trace, "_record_llm_metrics") as mock_metrics:
                    tencent_data_trace._process_workflow_nodes(trace_info, 123)

                    assert tencent_data_trace.trace_client.add_span.call_count == 2
                    mock_metrics.assert_called_once_with(node1)

    def test_process_workflow_nodes_node_exception(self, tencent_data_trace, mock_trace_utils):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        mock_trace_utils.convert_to_span_id.return_value = 111

        node = MagicMock(spec=WorkflowNodeExecution)
        node.id = "n1"

        with patch.object(tencent_data_trace, "_get_workflow_node_executions", return_value=[node]):
            with patch.object(tencent_data_trace, "_build_workflow_node_span", side_effect=Exception("node error")):
                with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                    tencent_data_trace._process_workflow_nodes(trace_info, 123)
                    # The exception should be caught by the outer handler since convert_to_span_id is called first
                    mock_log.assert_called_once_with("[Tencent APM] Failed to process workflow nodes")

    def test_process_workflow_nodes_exception(self, tencent_data_trace, mock_trace_utils):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        mock_trace_utils.convert_to_span_id.side_effect = Exception("outer error")

        with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
            tencent_data_trace._process_workflow_nodes(trace_info, 123)
            mock_log.assert_called_once_with("[Tencent APM] Failed to process workflow nodes")

    def test_build_workflow_node_span(self, tencent_data_trace, mock_span_builder):
        trace_info = MagicMock(spec=WorkflowTraceInfo)

        nodes = [
            (BuiltinNodeTypes.LLM, mock_span_builder.build_workflow_llm_span),
            (BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL, mock_span_builder.build_workflow_retrieval_span),
            (BuiltinNodeTypes.TOOL, mock_span_builder.build_workflow_tool_span),
            (BuiltinNodeTypes.CODE, mock_span_builder.build_workflow_task_span),
        ]

        for node_type, builder_method in nodes:
            node = MagicMock(spec=WorkflowNodeExecution)
            node.node_type = node_type
            builder_method.return_value = "span"

            result = tencent_data_trace._build_workflow_node_span(node, 123, trace_info, 456)

            assert result == "span"
            builder_method.assert_called_once_with(123, 456, trace_info, node)

    def test_build_workflow_node_span_exception(self, tencent_data_trace, mock_span_builder):
        node = MagicMock(spec=WorkflowNodeExecution)
        node.node_type = BuiltinNodeTypes.LLM
        node.id = "n1"
        mock_span_builder.build_workflow_llm_span.side_effect = Exception("error")

        with patch("dify_trace_tencent.tencent_trace.logger.debug") as mock_log:
            result = tencent_data_trace._build_workflow_node_span(node, 123, MagicMock(), 456)
            assert result is None
            mock_log.assert_called_once()

    def test_get_workflow_node_executions(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.metadata = {"app_id": "app-1"}
        trace_info.workflow_run_id = "run-1"

        app = MagicMock(spec=App)
        app.id = "app-1"
        app.created_by = "user-1"

        account = MagicMock(spec=Account)
        account.id = "user-1"

        tenant_join = MagicMock(spec=TenantAccountJoin)
        tenant_join.tenant_id = "tenant-1"

        mock_executions = [MagicMock()]

        with patch("dify_trace_tencent.tencent_trace.db") as mock_db:
            mock_db.engine = "engine"
            with patch("dify_trace_tencent.tencent_trace.Session") as mock_session_ctx:
                session = mock_session_ctx.return_value.__enter__.return_value
                session.scalar.side_effect = [app, account, tenant_join]

                with patch("dify_trace_tencent.tencent_trace.SQLAlchemyWorkflowNodeExecutionRepository") as mock_repo:
                    mock_repo.return_value.get_by_workflow_execution.return_value = mock_executions

                    results = tencent_data_trace._get_workflow_node_executions(trace_info)

                    assert results == mock_executions
                    account.set_tenant_id.assert_called_once_with("tenant-1")

    def test_get_workflow_node_executions_no_app_id(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.metadata = {}

        with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
            results = tencent_data_trace._get_workflow_node_executions(trace_info)
            assert results == []
            mock_log.assert_called_once()

    def test_get_workflow_node_executions_app_not_found(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.metadata = {"app_id": "app-1"}

        with patch("dify_trace_tencent.tencent_trace.db") as mock_db:
            mock_db.init_app = MagicMock()  # Ensure init_app is mocked
            mock_db.engine = "engine"
            with patch("dify_trace_tencent.tencent_trace.Session") as mock_session_ctx:
                session = mock_session_ctx.return_value.__enter__.return_value
                session.scalar.return_value = None

                with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                    results = tencent_data_trace._get_workflow_node_executions(trace_info)
                    assert results == []
                    mock_log.assert_called_once()

    def test_get_user_id_workflow(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.tenant_id = "tenant-1"
        trace_info.metadata = {"user_id": "user-1"}

        with patch("dify_trace_tencent.tencent_trace.sessionmaker", side_effect=Exception("Database error")):
            with patch("dify_trace_tencent.tencent_trace.db") as mock_db:
                mock_db.init_app = MagicMock()
                mock_db.engine = MagicMock()

                user_id = tencent_data_trace._get_user_id(trace_info)
                assert user_id == "unknown"

    def test_get_user_id_only_user_id(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)
        trace_info.metadata = {"user_id": "user-1"}

        user_id = tencent_data_trace._get_user_id(trace_info)
        assert user_id == "user-1"

    def test_get_user_id_anonymous(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)
        trace_info.metadata = {}

        user_id = tencent_data_trace._get_user_id(trace_info)
        assert user_id == "anonymous"

    def test_get_user_id_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.tenant_id = "t"
        trace_info.metadata = {"user_id": "u"}

        with patch("dify_trace_tencent.tencent_trace.sessionmaker", side_effect=Exception("error")):
            with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
                user_id = tencent_data_trace._get_user_id(trace_info)
                assert user_id == "unknown"
                mock_log.assert_called_once_with("[Tencent APM] Failed to get user ID")

    def test_record_llm_metrics_usage_in_process_data(self, tencent_data_trace):
        node = MagicMock(spec=WorkflowNodeExecution)
        node.process_data = {
            "usage": {
                "latency": 2.5,
                "time_to_first_token": 0.5,
                "time_to_generate": 2.0,
                "prompt_tokens": 10,
                "completion_tokens": 20,
            },
            "model_provider": "openai",
            "model_name": "gpt-4",
            "model_mode": "chat",
        }
        node.outputs = {}

        tencent_data_trace._record_llm_metrics(node)

        tencent_data_trace.trace_client.record_llm_duration.assert_called_once()
        tencent_data_trace.trace_client.record_time_to_first_token.assert_called_once()
        tencent_data_trace.trace_client.record_time_to_generate.assert_called_once()
        assert tencent_data_trace.trace_client.record_token_usage.call_count == 2

    def test_record_llm_metrics_usage_in_outputs(self, tencent_data_trace):
        node = MagicMock(spec=WorkflowNodeExecution)
        node.process_data = {}
        node.outputs = {"usage": {"latency": 1.0, "prompt_tokens": 5}}

        tencent_data_trace._record_llm_metrics(node)
        tencent_data_trace.trace_client.record_llm_duration.assert_called_once()
        tencent_data_trace.trace_client.record_token_usage.assert_called_once()

    def test_record_llm_metrics_exception(self, tencent_data_trace):
        node = MagicMock(spec=WorkflowNodeExecution)
        node.process_data = None
        node.outputs = None

        with patch("dify_trace_tencent.tencent_trace.logger.debug") as mock_log:
            tencent_data_trace._record_llm_metrics(node)
            # Should not crash

    def test_record_message_llm_metrics(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)
        trace_info.metadata = {"ls_provider": "openai", "ls_model_name": "gpt-4"}
        trace_info.message_data = {"provider_response_latency": 1.1}
        trace_info.is_streaming_request = True
        trace_info.gen_ai_server_time_to_first_token = 0.2
        trace_info.llm_streaming_time_to_generate = 0.9
        trace_info.message_tokens = 15
        trace_info.answer_tokens = 25

        tencent_data_trace._record_message_llm_metrics(trace_info)

        tencent_data_trace.trace_client.record_llm_duration.assert_called_once()
        tencent_data_trace.trace_client.record_time_to_first_token.assert_called_once()
        tencent_data_trace.trace_client.record_time_to_generate.assert_called_once()
        assert tencent_data_trace.trace_client.record_token_usage.call_count == 2

    def test_record_message_llm_metrics_object_data(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)
        trace_info.metadata = {}
        msg_data = MagicMock()
        msg_data.provider_response_latency = 1.1
        msg_data.model_provider = "anthropic"
        msg_data.model_id = "claude"
        trace_info.message_data = msg_data
        trace_info.is_streaming_request = False

        tencent_data_trace._record_message_llm_metrics(trace_info)
        tencent_data_trace.trace_client.record_llm_duration.assert_called_once()

    def test_record_message_llm_metrics_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)
        trace_info.metadata = None

        with patch("dify_trace_tencent.tencent_trace.logger.debug") as mock_log:
            tencent_data_trace._record_message_llm_metrics(trace_info)
            # Should not crash

    def test_record_workflow_trace_duration(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        from datetime import datetime, timedelta

        now = datetime.now()
        trace_info.start_time = now
        trace_info.end_time = now + timedelta(seconds=3)
        trace_info.workflow_run_status = "succeeded"
        trace_info.conversation_id = "conv-1"

        # Mock the record_trace_duration method to capture arguments
        with patch.object(tencent_data_trace.trace_client, "record_trace_duration") as mock_record:
            tencent_data_trace._record_workflow_trace_duration(trace_info)

            # Assert the method was called once
            mock_record.assert_called_once()

            # Extract arguments passed to the method
            args, kwargs = mock_record.call_args

            # Validate the duration argument
            assert args[0] == 3.0

            # Validate the attributes dict in kwargs
            attributes = kwargs["attributes"] if "attributes" in kwargs else args[1] if len(args) > 1 else {}
            assert attributes["conversation_mode"] == "workflow"
            assert attributes["has_conversation"] == "true"

    def test_record_workflow_trace_duration_fallback(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.start_time = None
        trace_info.workflow_run_elapsed_time = 4.5
        trace_info.workflow_run_status = "failed"
        trace_info.conversation_id = None

        with patch.object(tencent_data_trace.trace_client, "record_trace_duration") as mock_record:
            tencent_data_trace._record_workflow_trace_duration(trace_info)
            mock_record.assert_called_once()
            args, kwargs = mock_record.call_args
            assert args[0] == 4.5
            # Check attributes dict (either in kwargs or as second positional arg)
            attributes = kwargs["attributes"] if "attributes" in kwargs else args[1] if len(args) > 1 else {}
            assert attributes["has_conversation"] == "false"

    def test_record_workflow_trace_duration_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=WorkflowTraceInfo)
        trace_info.start_time = MagicMock()  # This might cause total_seconds() to fail if not mocked right

        with patch("dify_trace_tencent.tencent_trace.logger.debug") as mock_log:
            tencent_data_trace._record_workflow_trace_duration(trace_info)

    def test_record_message_trace_duration(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)
        from datetime import datetime, timedelta

        now = datetime.now()
        trace_info.start_time = now
        trace_info.end_time = now + timedelta(seconds=2)
        trace_info.conversation_mode = "chat"
        trace_info.is_streaming_request = True

        tencent_data_trace._record_message_trace_duration(trace_info)
        tencent_data_trace.trace_client.record_trace_duration.assert_called_once_with(
            2.0, {"conversation_mode": "chat", "stream": "true"}
        )

    def test_record_message_trace_duration_exception(self, tencent_data_trace):
        trace_info = MagicMock(spec=MessageTraceInfo)
        trace_info.start_time = None

        with patch("dify_trace_tencent.tencent_trace.logger.debug") as mock_log:
            tencent_data_trace._record_message_trace_duration(trace_info)

    def test_close(self, tencent_data_trace):
        client = tencent_data_trace.trace_client
        tencent_data_trace.close()
        client.shutdown.assert_called_once()

    def test_close_is_idempotent(self, tencent_data_trace):
        client = tencent_data_trace.trace_client

        tencent_data_trace.close()
        tencent_data_trace.close()

        client.shutdown.assert_called_once()

    def test_close_exception(self, tencent_data_trace):
        tencent_data_trace.trace_client.shutdown.side_effect = Exception("error")
        with patch("dify_trace_tencent.tencent_trace.logger.exception") as mock_log:
            tencent_data_trace.close()
            mock_log.assert_called_once_with("[Tencent APM] Failed to shutdown trace client during cleanup")

    def test_close_handles_async_shutdown_mock(self, tencent_data_trace):
        shutdown = AsyncMock()
        tencent_data_trace.trace_client.shutdown = shutdown

        with warnings.catch_warnings(record=True) as caught:
            warnings.simplefilter("always")
            tencent_data_trace.close()
            gc.collect()

        shutdown.assert_called_once()
        assert not [
            warning
            for warning in caught
            if issubclass(warning.category, RuntimeWarning)
            and "AsyncMockMixin._execute_mock_call" in str(warning.message)
        ]
