"""Tests for OpikDataTrace workflow_trace changes.

Covers:
- _seed_to_uuid4 helper: produces valid UUID4 strings deterministically
- prepare_opik_uuid helper: basic contract
- workflow_trace without message_id now creates a root span parented to None
- workflow_trace without message_id: node spans parent to root_span_id (not workflow_app_log_id)
- workflow_trace with message_id still creates root span keyed on workflow_run_id (unchanged path)
"""

from __future__ import annotations

import uuid
from datetime import datetime
from typing import cast
from unittest.mock import MagicMock, patch

from dify_trace_opik.opik_trace import OpikDataTrace, _seed_to_uuid4, prepare_opik_uuid

from core.ops.entities.trace_entity import TraceTaskName, WorkflowTraceInfo

# A stable UUID4 used as the workflow_run_id throughout all tests.
_WORKFLOW_RUN_ID = "a3f1b2c4-d5e6-4f78-9a0b-c1d2e3f4a5b6"


# ---------------------------------------------------------------------------
# Helpers
# ---------------------------------------------------------------------------


def _make_workflow_trace_info(
    *,
    message_id: str | None = None,
    workflow_app_log_id: str | None = None,
    workflow_run_id: str = _WORKFLOW_RUN_ID,
) -> WorkflowTraceInfo:
    """Return a minimal WorkflowTraceInfo suitable for unit testing."""
    return WorkflowTraceInfo(
        message_id=message_id,
        workflow_id="wf-id",
        tenant_id="tenant-id",
        workflow_run_id=workflow_run_id,
        workflow_app_log_id=workflow_app_log_id,
        workflow_run_elapsed_time=1.5,
        workflow_run_status="succeeded",
        workflow_run_inputs={"query": "hello"},
        workflow_run_outputs={"result": "world"},
        workflow_run_version="1",
        total_tokens=42,
        file_list=[],
        query="hello",
        start_time=datetime(2025, 1, 1, 12, 0, 0),
        end_time=datetime(2025, 1, 1, 12, 0, 1),
        metadata={"app_id": "app-abc"},
        conversation_id=None,
    )


def _make_opik_trace_instance() -> OpikDataTrace:
    """Construct an OpikDataTrace with the Opik SDK client mocked out."""
    with patch("dify_trace_opik.opik_trace.Opik"):
        from dify_trace_opik.config import OpikConfig

        config = OpikConfig(api_key="key", project="test-project", url="https://www.comet.com/opik/api/")
        instance = OpikDataTrace(config)

    instance.add_trace = MagicMock(return_value=MagicMock(id="mock-trace-id"))
    instance.add_span = MagicMock()
    instance.get_service_account_with_tenant = MagicMock(return_value=MagicMock())
    return instance


def _add_trace_mock(instance: OpikDataTrace) -> MagicMock:
    return cast(MagicMock, instance.add_trace)


def _add_span_mock(instance: OpikDataTrace) -> MagicMock:
    return cast(MagicMock, instance.add_span)


# ---------------------------------------------------------------------------
# _seed_to_uuid4
# ---------------------------------------------------------------------------


class TestSeedToUuid4:
    def test_returns_valid_uuid4_string(self):
        result = _seed_to_uuid4("some-arbitrary-seed")
        parsed = uuid.UUID(result)
        assert parsed.version == 4

    def test_is_deterministic(self):
        assert _seed_to_uuid4("seed-abc") == _seed_to_uuid4("seed-abc")

    def test_different_seeds_give_different_results(self):
        assert _seed_to_uuid4("seed-1") != _seed_to_uuid4("seed-2")

    def test_workflow_run_id_with_root_suffix_is_valid_uuid4(self):
        """The primary use-case: deriving a root-span UUID from workflow_run_id + '-root'."""
        seed = _WORKFLOW_RUN_ID + "-root"
        result = _seed_to_uuid4(seed)
        parsed = uuid.UUID(result)
        assert parsed.version == 4

    def test_seed_and_seed_root_produce_different_uuids(self):
        """Root span UUID must differ from the base workflow UUID to avoid ID collisions."""
        base = _seed_to_uuid4(_WORKFLOW_RUN_ID)
        with_root = _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root")
        assert base != with_root


# ---------------------------------------------------------------------------
# prepare_opik_uuid
# ---------------------------------------------------------------------------


class TestPrepareOpikUuid:
    def test_is_deterministic(self):
        dt = datetime(2025, 6, 15, 10, 30, 0)
        uid = str(uuid.uuid4())
        assert prepare_opik_uuid(dt, uid) == prepare_opik_uuid(dt, uid)

    def test_different_uuids_give_different_results(self):
        dt = datetime(2025, 6, 15, 10, 30, 0)
        assert prepare_opik_uuid(dt, str(uuid.uuid4())) != prepare_opik_uuid(dt, str(uuid.uuid4()))

    def test_none_datetime_does_not_raise(self):
        assert prepare_opik_uuid(None, str(uuid.uuid4())) is not None

    def test_none_uuid_does_not_raise(self):
        assert prepare_opik_uuid(datetime(2025, 1, 1), None) is not None


# ---------------------------------------------------------------------------
# workflow_trace — no message_id (new code path)
# ---------------------------------------------------------------------------


class TestWorkflowTraceWithoutMessageId:
    def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None):
        instance = _make_opik_trace_instance()
        fake_repo = MagicMock()
        fake_repo.get_by_workflow_execution.return_value = node_executions or []

        with (
            patch("dify_trace_opik.opik_trace.db") as mock_db,
            patch("dify_trace_opik.opik_trace.sessionmaker"),
            patch(
                "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository",
                return_value=fake_repo,
            ),
        ):
            mock_db.engine = MagicMock()
            instance.workflow_trace(trace_info)

        return instance

    def _expected_root_span_id(self, trace_info: WorkflowTraceInfo):
        return prepare_opik_uuid(
            trace_info.start_time,
            _seed_to_uuid4(trace_info.workflow_run_id + "-root"),
        )

    def test_root_span_is_created(self):
        trace_info = _make_workflow_trace_info(message_id=None)
        instance = self._run(trace_info)
        assert _add_span_mock(instance).called

    def test_root_span_id_matches_expected(self):
        trace_info = _make_workflow_trace_info(message_id=None)
        instance = self._run(trace_info)

        expected = self._expected_root_span_id(trace_info)
        root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0]
        assert root_span_kwargs["id"] == expected

    def test_root_span_has_no_parent(self):
        trace_info = _make_workflow_trace_info(message_id=None)
        instance = self._run(trace_info)

        root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0]
        assert root_span_kwargs["parent_span_id"] is None

    def test_trace_name_is_workflow_trace(self):
        """Without message_id, the Opik trace itself should be named WORKFLOW_TRACE."""
        trace_info = _make_workflow_trace_info(message_id=None)
        instance = self._run(trace_info)

        trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0]
        assert trace_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE

    def test_root_span_name_is_workflow_trace(self):
        trace_info = _make_workflow_trace_info(message_id=None)
        instance = self._run(trace_info)

        root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0]
        assert root_span_kwargs["name"] == TraceTaskName.WORKFLOW_TRACE

    def test_root_span_has_workflow_tag(self):
        trace_info = _make_workflow_trace_info(message_id=None)
        instance = self._run(trace_info)

        root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0]
        assert "workflow" in root_span_kwargs["tags"]

    def test_node_execution_spans_are_parented_to_root(self):
        """Node spans must use root_span_id as parent, not any other ID."""
        trace_info = _make_workflow_trace_info(message_id=None)
        expected_root_span_id = self._expected_root_span_id(trace_info)

        node_exec = MagicMock()
        node_exec.id = str(uuid.uuid4())
        node_exec.title = "LLM Node"
        node_exec.node_type = "llm"
        node_exec.status = "succeeded"
        node_exec.process_data = {}
        node_exec.inputs = {"prompt": "hi"}
        node_exec.outputs = {"text": "hello"}
        node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0)
        node_exec.elapsed_time = 0.5
        node_exec.metadata = {}

        instance = self._run(trace_info, node_executions=[node_exec])

        # call_args_list[0] = root span, [1] = node execution span
        add_span = _add_span_mock(instance)
        assert add_span.call_count == 2
        node_span_kwargs = add_span.call_args_list[1][0][0]
        assert node_span_kwargs["parent_span_id"] == expected_root_span_id

    def test_node_span_not_parented_to_workflow_app_log_id(self):
        """Old behaviour derived parent from workflow_app_log_id; that must no longer apply."""
        trace_info = _make_workflow_trace_info(
            message_id=None,
            workflow_app_log_id=str(uuid.uuid4()),
        )

        node_exec = MagicMock()
        node_exec.id = str(uuid.uuid4())
        node_exec.title = "Tool Node"
        node_exec.node_type = "tool"
        node_exec.status = "succeeded"
        node_exec.process_data = {}
        node_exec.inputs = {}
        node_exec.outputs = {}
        node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0)
        node_exec.elapsed_time = 0.2
        node_exec.metadata = {}

        instance = self._run(trace_info, node_executions=[node_exec])

        old_parent_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_app_log_id)
        node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0]
        assert node_span_kwargs["parent_span_id"] != old_parent_id

    def test_root_span_id_differs_from_trace_id(self):
        """The root span must have a different ID from the Opik trace to maintain correct hierarchy."""
        trace_info = _make_workflow_trace_info(message_id=None)
        dify_trace_id = trace_info.trace_id or trace_info.workflow_run_id
        opik_trace_id = prepare_opik_uuid(trace_info.start_time, dify_trace_id)
        root_span_id = self._expected_root_span_id(trace_info)
        assert root_span_id != opik_trace_id


# ---------------------------------------------------------------------------
# workflow_trace — with message_id (unchanged path, guard against regression)
# ---------------------------------------------------------------------------


class TestWorkflowTraceWithMessageId:
    _MESSAGE_ID = str(uuid.uuid4())

    def _run(self, trace_info: WorkflowTraceInfo, node_executions: list | None = None):
        instance = _make_opik_trace_instance()
        fake_repo = MagicMock()
        fake_repo.get_by_workflow_execution.return_value = node_executions or []

        with (
            patch("dify_trace_opik.opik_trace.db") as mock_db,
            patch("dify_trace_opik.opik_trace.sessionmaker"),
            patch(
                "dify_trace_opik.opik_trace.DifyCoreRepositoryFactory.create_workflow_node_execution_repository",
                return_value=fake_repo,
            ),
        ):
            mock_db.engine = MagicMock()
            instance.workflow_trace(trace_info)

        return instance

    def test_trace_name_is_message_trace(self):
        """With message_id, the Opik trace should be named MESSAGE_TRACE."""
        trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID)
        instance = self._run(trace_info)

        trace_kwargs = _add_trace_mock(instance).call_args_list[0][0][0]
        assert trace_kwargs["name"] == TraceTaskName.MESSAGE_TRACE

    def test_root_span_uses_workflow_run_id_directly(self):
        """When message_id is set, root_span_id = prepare_opik_uuid(start_time, workflow_run_id)."""
        trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID)
        instance = self._run(trace_info)

        expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id)
        root_span_kwargs = _add_span_mock(instance).call_args_list[0][0][0]
        assert root_span_kwargs["id"] == expected_root_span_id

    def test_root_span_id_differs_from_no_message_id_case(self):
        """The two branches must produce different root span IDs for the same workflow_run_id."""
        id_with_message = prepare_opik_uuid(
            datetime(2025, 1, 1, 12, 0, 0),
            _WORKFLOW_RUN_ID,
        )
        id_without_message = prepare_opik_uuid(
            datetime(2025, 1, 1, 12, 0, 0),
            _seed_to_uuid4(_WORKFLOW_RUN_ID + "-root"),
        )
        assert id_with_message != id_without_message

    def test_node_spans_parented_to_workflow_run_root_span(self):
        """Node spans must still parent to root_span_id derived from workflow_run_id."""
        trace_info = _make_workflow_trace_info(message_id=self._MESSAGE_ID)
        expected_root_span_id = prepare_opik_uuid(trace_info.start_time, trace_info.workflow_run_id)

        node_exec = MagicMock()
        node_exec.id = str(uuid.uuid4())
        node_exec.title = "LLM"
        node_exec.node_type = "llm"
        node_exec.status = "succeeded"
        node_exec.process_data = {}
        node_exec.inputs = {}
        node_exec.outputs = {}
        node_exec.created_at = datetime(2025, 1, 1, 12, 0, 0)
        node_exec.elapsed_time = 0.3
        node_exec.metadata = {}

        instance = self._run(trace_info, node_executions=[node_exec])

        node_span_kwargs = _add_span_mock(instance).call_args_list[1][0][0]
        assert node_span_kwargs["parent_span_id"] == expected_root_span_id
