"""Comprehensive tests for core.ops.mlflow_trace.mlflow_trace module."""

from __future__ import annotations

import json
import os
from datetime import UTC, datetime
from types import SimpleNamespace
from unittest.mock import MagicMock, patch

import pytest

from core.ops.entities.config_entity import DatabricksConfig, MLflowConfig
from core.ops.entities.trace_entity import (
    DatasetRetrievalTraceInfo,
    GenerateNameTraceInfo,
    MessageTraceInfo,
    ModerationTraceInfo,
    SuggestedQuestionTraceInfo,
    ToolTraceInfo,
    WorkflowTraceInfo,
)
from core.ops.mlflow_trace.mlflow_trace import MLflowDataTrace, datetime_to_nanoseconds
from dify_graph.enums import BuiltinNodeTypes

# ── Helpers ──────────────────────────────────────────────────────────────────


def _dt() -> datetime:
    return datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC)


def _make_workflow_trace_info(**overrides) -> WorkflowTraceInfo:
    defaults = {
        "workflow_id": "wf-id",
        "tenant_id": "tenant",
        "workflow_run_id": "run-1",
        "workflow_run_elapsed_time": 1.0,
        "workflow_run_status": "succeeded",
        "workflow_run_inputs": {"key": "val"},
        "workflow_run_outputs": {"answer": "42"},
        "workflow_run_version": "v1",
        "total_tokens": 10,
        "file_list": [],
        "query": "hello",
        "metadata": {"user_id": "u1", "conversation_id": "c1"},
        "start_time": _dt(),
        "end_time": _dt(),
    }
    defaults.update(overrides)
    return WorkflowTraceInfo(**defaults)


def _make_message_trace_info(**overrides) -> MessageTraceInfo:
    defaults = {
        "conversation_model": "chat",
        "message_tokens": 5,
        "answer_tokens": 10,
        "total_tokens": 15,
        "conversation_mode": "chat",
        "metadata": {"conversation_id": "c1", "from_account_id": "a1"},
        "message_id": "msg-1",
        "message_data": SimpleNamespace(
            model_provider="openai",
            model_id="gpt-4",
            total_price=0.01,
            answer="response text",
        ),
        "inputs": {"prompt": "hi"},
        "outputs": "ok",
        "start_time": _dt(),
        "end_time": _dt(),
        "error": None,
    }
    defaults.update(overrides)
    return MessageTraceInfo(**defaults)


def _make_tool_trace_info(**overrides) -> ToolTraceInfo:
    defaults = {
        "tool_name": "my_tool",
        "tool_inputs": {"x": 1},
        "tool_outputs": "output",
        "tool_config": {"desc": "d"},
        "tool_parameters": {"p": "v"},
        "time_cost": 0.5,
        "metadata": {"user_id": "u1"},
        "message_id": "msg-1",
        "inputs": {"i": "v"},
        "outputs": {"o": "v"},
        "start_time": _dt(),
        "end_time": _dt(),
        "error": None,
    }
    defaults.update(overrides)
    return ToolTraceInfo(**defaults)


def _make_moderation_trace_info(**overrides) -> ModerationTraceInfo:
    defaults = {
        "flagged": False,
        "action": "allow",
        "preset_response": "",
        "query": "test",
        "metadata": {"user_id": "u1"},
        "message_id": "msg-1",
    }
    defaults.update(overrides)
    return ModerationTraceInfo(**defaults)


def _make_dataset_retrieval_trace_info(**overrides) -> DatasetRetrievalTraceInfo:
    defaults = {
        "metadata": {"user_id": "u1"},
        "message_id": "msg-1",
        "message_data": SimpleNamespace(),
        "inputs": "query",
        "documents": [{"content": "doc"}],
        "start_time": _dt(),
        "end_time": _dt(),
    }
    defaults.update(overrides)
    return DatasetRetrievalTraceInfo(**defaults)


def _make_suggested_question_trace_info(**overrides) -> SuggestedQuestionTraceInfo:
    defaults = {
        "suggested_question": ["q1", "q2"],
        "level": "info",
        "total_tokens": 5,
        "metadata": {"user_id": "u1"},
        "message_id": "msg-1",
        "message_data": SimpleNamespace(created_at=_dt(), updated_at=_dt()),
        "inputs": {"i": 1},
        "start_time": _dt(),
        "end_time": _dt(),
        "error": None,
    }
    defaults.update(overrides)
    return SuggestedQuestionTraceInfo(**defaults)


def _make_generate_name_trace_info(**overrides) -> GenerateNameTraceInfo:
    defaults = {
        "tenant_id": "t1",
        "metadata": {"user_id": "u1"},
        "message_id": "msg-1",
        "inputs": {"i": 1},
        "outputs": {"name": "test"},
        "start_time": _dt(),
        "end_time": _dt(),
    }
    defaults.update(overrides)
    return GenerateNameTraceInfo(**defaults)


def _make_node(**overrides):
    """Create a mock workflow node execution row."""
    defaults = {
        "id": "node-1",
        "tenant_id": "t1",
        "app_id": "app-1",
        "title": "Node Title",
        "node_type": BuiltinNodeTypes.CODE,
        "status": "succeeded",
        "inputs": '{"key": "value"}',
        "outputs": '{"result": "ok"}',
        "created_at": _dt(),
        "elapsed_time": 1.0,
        "process_data": None,
        "execution_metadata": None,
    }
    defaults.update(overrides)
    return SimpleNamespace(**defaults)


# ── Fixtures ─────────────────────────────────────────────────────────────────


@pytest.fixture
def mock_mlflow():
    with patch("core.ops.mlflow_trace.mlflow_trace.mlflow") as mock:
        yield mock


@pytest.fixture
def mock_tracing():
    """Patch all MLflow tracing functions used by the module."""
    with (
        patch("core.ops.mlflow_trace.mlflow_trace.start_span_no_context") as mock_start,
        patch("core.ops.mlflow_trace.mlflow_trace.update_current_trace") as mock_update,
        patch("core.ops.mlflow_trace.mlflow_trace.set_span_in_context") as mock_set,
        patch("core.ops.mlflow_trace.mlflow_trace.detach_span_from_context") as mock_detach,
    ):
        yield {
            "start": mock_start,
            "update": mock_update,
            "set": mock_set,
            "detach": mock_detach,
        }


@pytest.fixture
def mock_db():
    with patch("core.ops.mlflow_trace.mlflow_trace.db") as mock:
        yield mock


@pytest.fixture
def trace_instance(mock_mlflow):
    """Create an MLflowDataTrace using a basic MLflowConfig (no auth)."""
    config = MLflowConfig(tracking_uri="http://localhost:5000", experiment_id="0")
    return MLflowDataTrace(config)


# ── datetime_to_nanoseconds ─────────────────────────────────────────────────


class TestDatetimeToNanoseconds:
    def test_none_returns_none(self):
        assert datetime_to_nanoseconds(None) is None

    def test_converts_datetime(self):
        dt = datetime(2024, 1, 1, 0, 0, 0, tzinfo=UTC)
        expected = int(dt.timestamp() * 1_000_000_000)
        assert datetime_to_nanoseconds(dt) == expected


# ── __init__ / setup ─────────────────────────────────────────────────────────


class TestInit:
    def test_mlflow_config_no_auth(self, mock_mlflow):
        config = MLflowConfig(tracking_uri="http://localhost:5000", experiment_id="0")
        trace = MLflowDataTrace(config)
        mock_mlflow.set_tracking_uri.assert_called_with("http://localhost:5000")
        mock_mlflow.set_experiment.assert_called_with(experiment_id="0")
        assert trace.get_project_url() == "http://localhost:5000/#/experiments/0/traces"
        assert os.environ["MLFLOW_ENABLE_ASYNC_TRACE_LOGGING"] == "true"

    def test_mlflow_config_with_auth(self, mock_mlflow):
        config = MLflowConfig(
            tracking_uri="http://localhost:5000",
            experiment_id="1",
            username="user",
            password="pass",
        )
        MLflowDataTrace(config)
        assert os.environ["MLFLOW_TRACKING_USERNAME"] == "user"
        assert os.environ["MLFLOW_TRACKING_PASSWORD"] == "pass"

    def test_databricks_oauth(self, mock_mlflow):
        config = DatabricksConfig(
            host="https://db.com/",
            experiment_id="42",
            client_id="cid",
            client_secret="csec",
        )
        trace = MLflowDataTrace(config)
        assert os.environ["DATABRICKS_HOST"] == "https://db.com/"
        assert os.environ["DATABRICKS_CLIENT_ID"] == "cid"
        assert os.environ["DATABRICKS_CLIENT_SECRET"] == "csec"
        mock_mlflow.set_tracking_uri.assert_called_with("databricks")
        # Trailing slash stripped
        assert trace.get_project_url() == "https://db.com/ml/experiments/42/traces"

    def test_databricks_pat(self, mock_mlflow):
        config = DatabricksConfig(
            host="https://db.com",
            experiment_id="1",
            personal_access_token="pat",
        )
        trace = MLflowDataTrace(config)
        assert os.environ["DATABRICKS_TOKEN"] == "pat"
        assert "db.com/ml/experiments/1/traces" in trace.get_project_url()

    def test_databricks_no_creds_raises(self, mock_mlflow):
        config = DatabricksConfig(host="https://db.com", experiment_id="1")
        with pytest.raises(ValueError, match="Either Databricks token"):
            MLflowDataTrace(config)


# ── trace dispatcher ────────────────────────────────────────────────────────


class TestTraceDispatcher:
    def test_dispatches_workflow(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "workflow_trace") as mock_wt:
            trace_instance.trace(_make_workflow_trace_info())
            mock_wt.assert_called_once()

    def test_dispatches_message(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "message_trace") as mock_mt:
            trace_instance.trace(_make_message_trace_info())
            mock_mt.assert_called_once()

    def test_dispatches_tool(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "tool_trace") as mock_tt:
            trace_instance.trace(_make_tool_trace_info())
            mock_tt.assert_called_once()

    def test_dispatches_moderation(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "moderation_trace") as mock_mod:
            trace_instance.trace(_make_moderation_trace_info(message_data=SimpleNamespace(created_at=_dt())))
            mock_mod.assert_called_once()

    def test_dispatches_dataset_retrieval(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "dataset_retrieval_trace") as mock_dr:
            trace_instance.trace(_make_dataset_retrieval_trace_info())
            mock_dr.assert_called_once()

    def test_dispatches_suggested_question(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "suggested_question_trace") as mock_sq:
            trace_instance.trace(_make_suggested_question_trace_info())
            mock_sq.assert_called_once()

    def test_dispatches_generate_name(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "generate_name_trace") as mock_gn:
            trace_instance.trace(_make_generate_name_trace_info())
            mock_gn.assert_called_once()

    def test_reraises_exception(self, trace_instance, mock_tracing, mock_db):
        with patch.object(trace_instance, "workflow_trace", side_effect=RuntimeError("boom")):
            with pytest.raises(RuntimeError, match="boom"):
                trace_instance.trace(_make_workflow_trace_info())


# ── workflow_trace ───────────────────────────────────────────────────────────


class TestWorkflowTrace:
    def test_basic_workflow_no_nodes(self, trace_instance, mock_tracing, mock_db):
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = []
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"

        trace_info = _make_workflow_trace_info(conversation_id="sess-1")
        trace_instance.workflow_trace(trace_info)

        # Workflow span started and ended
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()

    def test_workflow_filters_sys_inputs_and_adds_query(self, trace_instance, mock_tracing, mock_db):
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = []
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"

        trace_info = _make_workflow_trace_info(
            workflow_run_inputs={"sys.app_id": "x", "user_input": "hi"},
            query="hello",
        )
        trace_instance.workflow_trace(trace_info)

        call_kwargs = mock_tracing["start"].call_args
        inputs = call_kwargs.kwargs["inputs"]
        assert "sys.app_id" not in inputs
        assert inputs["user_input"] == "hi"
        assert inputs["query"] == "hello"

    def test_workflow_with_llm_node(self, trace_instance, mock_tracing, mock_db):
        llm_node = _make_node(
            node_type=BuiltinNodeTypes.LLM,
            process_data=json.dumps(
                {
                    "prompts": [{"role": "user", "text": "hi"}],
                    "model_name": "gpt-4",
                    "model_provider": "openai",
                    "finish_reason": "stop",
                    "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
                }
            ),
            outputs='{"text": "hello world"}',
        )
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = [llm_node]

        workflow_span = MagicMock()
        node_span = MagicMock()
        mock_tracing["start"].side_effect = [workflow_span, node_span]
        mock_tracing["set"].return_value = "token"

        trace_instance.workflow_trace(_make_workflow_trace_info())
        assert mock_tracing["start"].call_count == 2
        node_span.end.assert_called_once()
        workflow_span.end.assert_called_once()

    def test_workflow_with_question_classifier_node(self, trace_instance, mock_tracing, mock_db):
        qc_node = _make_node(
            node_type=BuiltinNodeTypes.QUESTION_CLASSIFIER,
            process_data=json.dumps(
                {
                    "prompts": "classify this",
                    "model_name": "gpt-4",
                    "model_provider": "openai",
                }
            ),
        )
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = [qc_node]
        workflow_span = MagicMock()
        node_span = MagicMock()
        mock_tracing["start"].side_effect = [workflow_span, node_span]
        mock_tracing["set"].return_value = "token"

        trace_instance.workflow_trace(_make_workflow_trace_info())
        assert mock_tracing["start"].call_count == 2

    def test_workflow_with_http_request_node(self, trace_instance, mock_tracing, mock_db):
        http_node = _make_node(
            node_type=BuiltinNodeTypes.HTTP_REQUEST,
            process_data='{"url": "https://api.com"}',
        )
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = [http_node]
        workflow_span = MagicMock()
        node_span = MagicMock()
        mock_tracing["start"].side_effect = [workflow_span, node_span]
        mock_tracing["set"].return_value = "token"

        trace_instance.workflow_trace(_make_workflow_trace_info())
        # HTTP_REQUEST uses process_data as inputs
        node_start_call = mock_tracing["start"].call_args_list[1]
        assert node_start_call.kwargs["inputs"] == '{"url": "https://api.com"}'

    def test_workflow_with_knowledge_retrieval_node(self, trace_instance, mock_tracing, mock_db):
        kr_node = _make_node(
            node_type=BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL,
            outputs=json.dumps(
                {
                    "result": [
                        {"content": "doc1", "metadata": {"source": "s1"}},
                        {"content": "doc2", "metadata": {}},
                    ]
                }
            ),
        )
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = [kr_node]
        workflow_span = MagicMock()
        node_span = MagicMock()
        mock_tracing["start"].side_effect = [workflow_span, node_span]
        mock_tracing["set"].return_value = "token"

        trace_instance.workflow_trace(_make_workflow_trace_info())
        # outputs should be parsed to Document objects
        end_call = node_span.end.call_args
        outputs = end_call.kwargs["outputs"]
        assert len(outputs) == 2

    def test_workflow_with_failed_node(self, trace_instance, mock_tracing, mock_db):
        failed_node = _make_node(status="failed")
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = [failed_node]
        workflow_span = MagicMock()
        node_span = MagicMock()
        mock_tracing["start"].side_effect = [workflow_span, node_span]
        mock_tracing["set"].return_value = "token"

        trace_instance.workflow_trace(_make_workflow_trace_info())
        node_span.set_status.assert_called_once()
        node_span.add_event.assert_called_once()

    def test_workflow_with_workflow_error(self, trace_instance, mock_tracing, mock_db):
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = []
        workflow_span = MagicMock()
        mock_tracing["start"].return_value = workflow_span
        mock_tracing["set"].return_value = "token"

        trace_info = _make_workflow_trace_info(error="workflow failed")
        trace_instance.workflow_trace(trace_info)
        workflow_span.set_status.assert_called_once()
        workflow_span.add_event.assert_called_once()
        # Still ends the span via finally
        workflow_span.end.assert_called_once()

    def test_workflow_node_no_inputs_no_outputs(self, trace_instance, mock_tracing, mock_db):
        node = _make_node(inputs=None, outputs=None)
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = [node]
        workflow_span = MagicMock()
        node_span = MagicMock()
        mock_tracing["start"].side_effect = [workflow_span, node_span]
        mock_tracing["set"].return_value = "token"

        trace_instance.workflow_trace(_make_workflow_trace_info())
        node_call = mock_tracing["start"].call_args_list[1]
        assert node_call.kwargs["inputs"] == {}
        end_call = node_span.end.call_args
        assert end_call.kwargs["outputs"] == {}

    def test_workflow_no_user_id_no_conversation_id(self, trace_instance, mock_tracing, mock_db):
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = []
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"

        trace_info = _make_workflow_trace_info(
            metadata={},
            conversation_id=None,
        )
        trace_instance.workflow_trace(trace_info)
        # _set_trace_metadata still called with empty metadata
        mock_tracing["update"].assert_called_once()

    def test_workflow_empty_query(self, trace_instance, mock_tracing, mock_db):
        """When query is empty string, it's falsy so no query key added."""
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = []
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"

        trace_info = _make_workflow_trace_info(query="")
        trace_instance.workflow_trace(trace_info)
        call_kwargs = mock_tracing["start"].call_args
        inputs = call_kwargs.kwargs["inputs"]
        assert "query" not in inputs


# ── _parse_llm_inputs_and_attributes ─────────────────────────────────────────


class TestParseLlmInputsAndAttributes:
    def test_none_process_data(self, trace_instance):
        node = _make_node(process_data=None)
        inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node)
        assert inputs == {}
        assert attrs == {}

    def test_invalid_json(self, trace_instance):
        node = _make_node(process_data="not json")
        inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node)
        assert inputs == {}
        assert attrs == {}

    def test_valid_process_data_with_usage(self, trace_instance):
        node = _make_node(
            process_data=json.dumps(
                {
                    "prompts": [{"role": "user", "text": "hi"}],
                    "model_name": "gpt-4",
                    "model_provider": "openai",
                    "finish_reason": "stop",
                    "usage": {"prompt_tokens": 5, "completion_tokens": 10, "total_tokens": 15},
                }
            )
        )
        inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node)
        assert isinstance(inputs, list)
        assert attrs["model_name"] == "gpt-4"
        assert "usage" in attrs

    def test_valid_process_data_without_usage(self, trace_instance):
        node = _make_node(
            process_data=json.dumps(
                {
                    "prompts": "simple prompt",
                    "model_name": "gpt-3.5",
                }
            )
        )
        inputs, attrs = trace_instance._parse_llm_inputs_and_attributes(node)
        assert inputs == "simple prompt"
        assert attrs["model_name"] == "gpt-3.5"


# ── _parse_knowledge_retrieval_outputs ───────────────────────────────────────


class TestParseKnowledgeRetrievalOutputs:
    def test_with_results(self, trace_instance):
        outputs = {"result": [{"content": "c1", "metadata": {"s": "1"}}]}
        docs = trace_instance._parse_knowledge_retrieval_outputs(outputs)
        assert len(docs) == 1
        assert docs[0].page_content == "c1"

    def test_empty_result(self, trace_instance):
        outputs = {"result": []}
        result = trace_instance._parse_knowledge_retrieval_outputs(outputs)
        assert result == outputs

    def test_no_result_key(self, trace_instance):
        outputs = {"other": "data"}
        result = trace_instance._parse_knowledge_retrieval_outputs(outputs)
        assert result == outputs

    def test_result_not_list(self, trace_instance):
        outputs = {"result": "not a list"}
        result = trace_instance._parse_knowledge_retrieval_outputs(outputs)
        assert result == outputs


# ── message_trace ────────────────────────────────────────────────────────────


class TestMessageTrace:
    def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing, mock_db):
        trace_info = _make_message_trace_info(message_data=None)
        trace_instance.message_trace(trace_info)
        mock_tracing["start"].assert_not_called()

    def test_basic_message_trace(self, trace_instance, mock_tracing, mock_db):
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"
        mock_db.session.query.return_value.where.return_value.first.return_value = None

        trace_instance.message_trace(_make_message_trace_info())
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()

    def test_message_trace_with_error(self, trace_instance, mock_tracing, mock_db):
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"
        mock_db.session.query.return_value.where.return_value.first.return_value = None

        trace_info = _make_message_trace_info(error="something broke")
        trace_instance.message_trace(trace_info)
        span.set_status.assert_called_once()
        span.add_event.assert_called_once()

    def test_message_trace_with_file_data(self, trace_instance, mock_tracing, mock_db, monkeypatch):
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"
        mock_db.session.query.return_value.where.return_value.first.return_value = None
        monkeypatch.setenv("FILES_URL", "http://files.test")

        file_data = SimpleNamespace(url="path/to/file.png")
        trace_info = _make_message_trace_info(
            message_file_data=file_data,
            file_list=["existing_file.txt"],
        )
        trace_instance.message_trace(trace_info)
        call_kwargs = mock_tracing["start"].call_args
        attrs = call_kwargs.kwargs["attributes"]
        assert "http://files.test/path/to/file.png" in attrs["file_list"]
        assert "existing_file.txt" in attrs["file_list"]

    def test_message_trace_file_list_none(self, trace_instance, mock_tracing, mock_db):
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"
        mock_db.session.query.return_value.where.return_value.first.return_value = None

        trace_info = _make_message_trace_info(file_list=None, message_file_data=None)
        trace_instance.message_trace(trace_info)
        mock_tracing["start"].assert_called_once()

    def test_message_trace_with_end_user(self, trace_instance, mock_tracing, mock_db):
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"

        end_user = MagicMock()
        end_user.session_id = "session-xyz"
        mock_db.session.query.return_value.where.return_value.first.return_value = end_user

        trace_info = _make_message_trace_info(
            metadata={"from_end_user_id": "eu-1", "conversation_id": "c1"},
        )
        trace_instance.message_trace(trace_info)
        # update_current_trace called with user id from EndUser
        mock_tracing["update"].assert_called_once()

    def test_message_trace_with_no_conversation_id(self, trace_instance, mock_tracing, mock_db):
        span = MagicMock()
        mock_tracing["start"].return_value = span
        mock_tracing["set"].return_value = "token"
        mock_db.session.query.return_value.where.return_value.first.return_value = None

        trace_info = _make_message_trace_info(
            metadata={"from_account_id": "acc-1"},
        )
        trace_instance.message_trace(trace_info)
        mock_tracing["update"].assert_called_once()


# ── _get_message_user_id ─────────────────────────────────────────────────────


class TestGetMessageUserId:
    def test_returns_end_user_session_id(self, trace_instance, mock_db):
        end_user = MagicMock()
        end_user.session_id = "session-1"
        mock_db.session.query.return_value.where.return_value.first.return_value = end_user
        result = trace_instance._get_message_user_id({"from_end_user_id": "eu-1"})
        assert result == "session-1"

    def test_returns_account_id_when_no_end_user(self, trace_instance, mock_db):
        mock_db.session.query.return_value.where.return_value.first.return_value = None
        result = trace_instance._get_message_user_id({"from_end_user_id": "eu-1", "from_account_id": "acc-1"})
        assert result == "acc-1"

    def test_returns_account_id_when_no_end_user_id(self, trace_instance, mock_db):
        result = trace_instance._get_message_user_id({"from_account_id": "acc-1"})
        assert result == "acc-1"

    def test_returns_none_when_nothing(self, trace_instance, mock_db):
        result = trace_instance._get_message_user_id({})
        assert result is None


# ── tool_trace ───────────────────────────────────────────────────────────────


class TestToolTrace:
    def test_basic_tool_trace(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_instance.tool_trace(_make_tool_trace_info())
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()
        span.set_status.assert_not_called()

    def test_tool_trace_with_error(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_instance.tool_trace(_make_tool_trace_info(error="tool failed"))
        span.set_status.assert_called_once()
        span.add_event.assert_called_once()
        span.end.assert_called_once()


# ── moderation_trace ─────────────────────────────────────────────────────────


class TestModerationTrace:
    def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing):
        trace_info = _make_moderation_trace_info(message_data=None)
        trace_instance.moderation_trace(trace_info)
        mock_tracing["start"].assert_not_called()

    def test_basic_moderation_trace(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_info = _make_moderation_trace_info(
            message_data=SimpleNamespace(created_at=_dt()),
            start_time=_dt(),
            end_time=_dt(),
        )
        trace_instance.moderation_trace(trace_info)
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()
        end_kwargs = span.end.call_args.kwargs["outputs"]
        assert end_kwargs["action"] == "allow"
        assert end_kwargs["flagged"] is False

    def test_moderation_uses_message_data_created_at_if_no_start_time(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_info = _make_moderation_trace_info(
            message_data=SimpleNamespace(created_at=_dt()),
            start_time=None,
            end_time=_dt(),
        )
        trace_instance.moderation_trace(trace_info)
        mock_tracing["start"].assert_called_once()


# ── dataset_retrieval_trace ──────────────────────────────────────────────────


class TestDatasetRetrievalTrace:
    def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing):
        trace_info = _make_dataset_retrieval_trace_info(message_data=None)
        trace_instance.dataset_retrieval_trace(trace_info)
        mock_tracing["start"].assert_not_called()

    def test_basic_dataset_retrieval_trace(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_instance.dataset_retrieval_trace(_make_dataset_retrieval_trace_info())
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()


# ── suggested_question_trace ─────────────────────────────────────────────────


class TestSuggestedQuestionTrace:
    def test_returns_early_if_no_message_data(self, trace_instance, mock_tracing):
        trace_info = _make_suggested_question_trace_info(message_data=None)
        trace_instance.suggested_question_trace(trace_info)
        mock_tracing["start"].assert_not_called()

    def test_basic_suggested_question_trace(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_instance.suggested_question_trace(_make_suggested_question_trace_info())
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()

    def test_suggested_question_with_error(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_info = _make_suggested_question_trace_info(error="failed")
        trace_instance.suggested_question_trace(trace_info)
        span.set_status.assert_called_once()
        span.add_event.assert_called_once()

    def test_uses_message_data_times_when_no_start_end(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_info = _make_suggested_question_trace_info(
            start_time=None,
            end_time=None,
        )
        trace_instance.suggested_question_trace(trace_info)
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()


# ── generate_name_trace ──────────────────────────────────────────────────────


class TestGenerateNameTrace:
    def test_basic_generate_name_trace(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["start"].return_value = span

        trace_instance.generate_name_trace(_make_generate_name_trace_info())
        mock_tracing["start"].assert_called_once()
        span.end.assert_called_once()


# ── _get_workflow_nodes ──────────────────────────────────────────────────────


class TestGetWorkflowNodes:
    def test_queries_db(self, trace_instance, mock_db):
        mock_db.session.query.return_value.filter.return_value.order_by.return_value.all.return_value = ["n1", "n2"]
        result = trace_instance._get_workflow_nodes("run-1")
        assert result == ["n1", "n2"]


# ── _get_node_span_type ─────────────────────────────────────────────────────


class TestGetNodeSpanType:
    @pytest.mark.parametrize(
        ("node_type", "expected_contains"),
        [
            (BuiltinNodeTypes.LLM, "LLM"),
            (BuiltinNodeTypes.QUESTION_CLASSIFIER, "LLM"),
            (BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL, "RETRIEVER"),
            (BuiltinNodeTypes.TOOL, "TOOL"),
            (BuiltinNodeTypes.CODE, "TOOL"),
            (BuiltinNodeTypes.HTTP_REQUEST, "TOOL"),
            (BuiltinNodeTypes.AGENT, "AGENT"),
        ],
    )
    def test_mapped_types(self, trace_instance, node_type, expected_contains):
        result = trace_instance._get_node_span_type(node_type)
        assert expected_contains in str(result)

    def test_unknown_type_returns_chain(self, trace_instance):
        result = trace_instance._get_node_span_type("unknown_node")
        assert result == "CHAIN"


# ── _set_trace_metadata ─────────────────────────────────────────────────────


class TestSetTraceMetadata:
    def test_sets_and_detaches(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["set"].return_value = "token"

        trace_instance._set_trace_metadata(span, {"key": "val"})
        mock_tracing["set"].assert_called_once_with(span)
        mock_tracing["update"].assert_called_once_with(metadata={"key": "val"})
        mock_tracing["detach"].assert_called_once_with("token")

    def test_detaches_even_on_error(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["set"].return_value = "token"
        mock_tracing["update"].side_effect = RuntimeError("fail")

        with pytest.raises(RuntimeError):
            trace_instance._set_trace_metadata(span, {})
        mock_tracing["detach"].assert_called_once_with("token")

    def test_no_detach_when_token_is_none(self, trace_instance, mock_tracing):
        span = MagicMock()
        mock_tracing["set"].return_value = None

        trace_instance._set_trace_metadata(span, {})
        mock_tracing["detach"].assert_not_called()


# ── _parse_prompts ───────────────────────────────────────────────────────────


class TestParsePrompts:
    def test_string_input(self, trace_instance):
        assert trace_instance._parse_prompts("hello") == "hello"

    def test_dict_input(self, trace_instance):
        result = trace_instance._parse_prompts({"role": "user", "text": "hi"})
        assert result == {"role": "user", "content": "hi"}

    def test_list_input(self, trace_instance):
        prompts = [
            {"role": "user", "text": "hi"},
            {"role": "assistant", "text": "hello"},
        ]
        result = trace_instance._parse_prompts(prompts)
        assert len(result) == 2
        assert result[0]["role"] == "user"

    def test_none_input(self, trace_instance):
        assert trace_instance._parse_prompts(None) is None

    def test_int_passthrough(self, trace_instance):
        assert trace_instance._parse_prompts(42) == 42


# ── _parse_single_message ───────────────────────────────────────────────────


class TestParseSingleMessage:
    def test_basic_message(self, trace_instance):
        result = trace_instance._parse_single_message({"role": "user", "text": "hello"})
        assert result == {"role": "user", "content": "hello"}

    def test_default_role(self, trace_instance):
        result = trace_instance._parse_single_message({"text": "hello"})
        assert result["role"] == "user"

    def test_with_tool_calls(self, trace_instance):
        item = {
            "role": "assistant",
            "text": "",
            "tool_calls": [{"id": "tc1", "function": {"name": "fn"}}],
        }
        result = trace_instance._parse_single_message(item)
        assert "tool_calls" in result

    def test_tool_role_ignores_tool_calls(self, trace_instance):
        item = {
            "role": "tool",
            "text": "result",
            "tool_calls": [{"id": "tc1"}],
        }
        result = trace_instance._parse_single_message(item)
        assert "tool_calls" not in result

    def test_with_files(self, trace_instance):
        item = {"role": "user", "text": "look", "files": ["f1.png"]}
        result = trace_instance._parse_single_message(item)
        assert result["files"] == ["f1.png"]

    def test_no_files(self, trace_instance):
        result = trace_instance._parse_single_message({"role": "user", "text": "hi"})
        assert "files" not in result


# ── _resolve_tool_call_ids ───────────────────────────────────────────────────


class TestResolveToolCallIds:
    def test_resolves_tool_call_ids(self, trace_instance):
        messages = [
            {
                "role": "assistant",
                "content": "",
                "tool_calls": [{"id": "tc1"}, {"id": "tc2"}],
            },
            {"role": "tool", "content": "result1"},
            {"role": "tool", "content": "result2"},
        ]
        result = trace_instance._resolve_tool_call_ids(messages)
        assert result[1]["tool_call_id"] == "tc1"
        assert result[2]["tool_call_id"] == "tc2"

    def test_no_tool_calls(self, trace_instance):
        messages = [
            {"role": "user", "content": "hi"},
            {"role": "assistant", "content": "hello"},
        ]
        result = trace_instance._resolve_tool_call_ids(messages)
        assert "tool_call_id" not in result[0]
        assert "tool_call_id" not in result[1]

    def test_tool_message_no_ids_available(self, trace_instance):
        """Tool message with no preceding tool_calls should not crash."""
        messages = [
            {"role": "tool", "content": "result"},
        ]
        result = trace_instance._resolve_tool_call_ids(messages)
        assert "tool_call_id" not in result[0]


# ── api_check ────────────────────────────────────────────────────────────────


class TestApiCheck:
    def test_success(self, trace_instance, mock_mlflow):
        mock_mlflow.search_experiments.return_value = []
        assert trace_instance.api_check() is True

    def test_failure(self, trace_instance, mock_mlflow):
        mock_mlflow.search_experiments.side_effect = ConnectionError("refused")
        with pytest.raises(ValueError, match="MLflow connection failed"):
            trace_instance.api_check()


# ── get_project_url ──────────────────────────────────────────────────────────


class TestGetProjectUrl:
    def test_returns_url(self, trace_instance):
        assert "experiments" in trace_instance.get_project_url()
