import collections
import logging
from datetime import UTC, datetime, timedelta
from types import SimpleNamespace
from unittest.mock import MagicMock

import pytest

from core.ops.entities.config_entity import LangfuseConfig
from core.ops.entities.trace_entity import (
    DatasetRetrievalTraceInfo,
    GenerateNameTraceInfo,
    MessageTraceInfo,
    ModerationTraceInfo,
    SuggestedQuestionTraceInfo,
    ToolTraceInfo,
    TraceTaskName,
    WorkflowTraceInfo,
)
from core.ops.langfuse_trace.entities.langfuse_trace_entity import (
    LangfuseGeneration,
    LangfuseSpan,
    LangfuseTrace,
    LevelEnum,
    UnitEnum,
)
from core.ops.langfuse_trace.langfuse_trace import LangFuseDataTrace
from dify_graph.enums import BuiltinNodeTypes
from models import EndUser
from models.enums import MessageStatus


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


@pytest.fixture
def langfuse_config():
    return LangfuseConfig(public_key="pk-123", secret_key="sk-123", host="https://cloud.langfuse.com")


@pytest.fixture
def trace_instance(langfuse_config, monkeypatch):
    # Mock Langfuse client to avoid network calls
    mock_client = MagicMock()
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.Langfuse", lambda **kwargs: mock_client)

    instance = LangFuseDataTrace(langfuse_config)
    return instance


def test_init(langfuse_config, monkeypatch):
    mock_langfuse = MagicMock()
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.Langfuse", mock_langfuse)
    monkeypatch.setenv("FILES_URL", "http://test.url")

    instance = LangFuseDataTrace(langfuse_config)

    mock_langfuse.assert_called_once_with(
        public_key=langfuse_config.public_key,
        secret_key=langfuse_config.secret_key,
        host=langfuse_config.host,
    )
    assert instance.file_base_url == "http://test.url"


def test_trace_dispatch(trace_instance, monkeypatch):
    methods = [
        "workflow_trace",
        "message_trace",
        "moderation_trace",
        "suggested_question_trace",
        "dataset_retrieval_trace",
        "tool_trace",
        "generate_name_trace",
    ]
    mocks = {method: MagicMock() for method in methods}
    for method, m in mocks.items():
        monkeypatch.setattr(trace_instance, method, m)

    # WorkflowTraceInfo
    info = MagicMock(spec=WorkflowTraceInfo)
    trace_instance.trace(info)
    mocks["workflow_trace"].assert_called_once_with(info)

    # MessageTraceInfo
    info = MagicMock(spec=MessageTraceInfo)
    trace_instance.trace(info)
    mocks["message_trace"].assert_called_once_with(info)

    # ModerationTraceInfo
    info = MagicMock(spec=ModerationTraceInfo)
    trace_instance.trace(info)
    mocks["moderation_trace"].assert_called_once_with(info)

    # SuggestedQuestionTraceInfo
    info = MagicMock(spec=SuggestedQuestionTraceInfo)
    trace_instance.trace(info)
    mocks["suggested_question_trace"].assert_called_once_with(info)

    # DatasetRetrievalTraceInfo
    info = MagicMock(spec=DatasetRetrievalTraceInfo)
    trace_instance.trace(info)
    mocks["dataset_retrieval_trace"].assert_called_once_with(info)

    # ToolTraceInfo
    info = MagicMock(spec=ToolTraceInfo)
    trace_instance.trace(info)
    mocks["tool_trace"].assert_called_once_with(info)

    # GenerateNameTraceInfo
    info = MagicMock(spec=GenerateNameTraceInfo)
    trace_instance.trace(info)
    mocks["generate_name_trace"].assert_called_once_with(info)


def test_workflow_trace_with_message_id(trace_instance, monkeypatch):
    # Setup trace info
    trace_info = WorkflowTraceInfo(
        workflow_id="wf-1",
        tenant_id="tenant-1",
        workflow_run_id="run-1",
        workflow_run_elapsed_time=1.0,
        workflow_run_status="succeeded",
        workflow_run_inputs={"input": "hi"},
        workflow_run_outputs={"output": "hello"},
        workflow_run_version="1.0",
        message_id="msg-1",
        conversation_id="conv-1",
        total_tokens=100,
        file_list=[],
        query="hi",
        start_time=_dt(),
        end_time=_dt() + timedelta(seconds=1),
        trace_id="trace-1",
        metadata={"app_id": "app-1", "user_id": "user-1"},
        workflow_app_log_id="log-1",
        error="",
    )

    # Mock DB and Repositories
    mock_session = MagicMock()
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.sessionmaker", lambda bind: lambda: mock_session)
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.db", MagicMock(engine="engine"))

    # Mock node executions
    node_llm = MagicMock()
    node_llm.id = "node-llm"
    node_llm.title = "LLM Node"
    node_llm.node_type = BuiltinNodeTypes.LLM
    node_llm.status = "succeeded"
    node_llm.process_data = {
        "model_mode": "chat",
        "model_name": "gpt-4",
        "model_provider": "openai",
        "usage": {"prompt_tokens": 10, "completion_tokens": 20},
    }
    node_llm.inputs = {"prompts": "p"}
    node_llm.outputs = {"text": "t"}
    node_llm.created_at = _dt()
    node_llm.elapsed_time = 0.5
    node_llm.metadata = {"foo": "bar"}

    node_other = MagicMock()
    node_other.id = "node-other"
    node_other.title = "Other Node"
    node_other.node_type = BuiltinNodeTypes.CODE
    node_other.status = "failed"
    node_other.process_data = None
    node_other.inputs = {"code": "print"}
    node_other.outputs = {"result": "ok"}
    node_other.created_at = None  # Trigger datetime.now() branch
    node_other.elapsed_time = 0.2
    node_other.metadata = None

    repo = MagicMock()
    repo.get_by_workflow_run.return_value = [node_llm, node_other]

    mock_factory = MagicMock()
    mock_factory.create_workflow_node_execution_repository.return_value = repo
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.DifyCoreRepositoryFactory", mock_factory)

    monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock())

    # Track calls to add_trace, add_span, add_generation
    trace_instance.add_trace = MagicMock()
    trace_instance.add_span = MagicMock()
    trace_instance.add_generation = MagicMock()

    trace_instance.workflow_trace(trace_info)

    # Verify add_trace (Workflow Level)
    trace_instance.add_trace.assert_called_once()
    trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"]
    assert trace_data.id == "trace-1"
    assert trace_data.name == TraceTaskName.MESSAGE_TRACE
    assert "message" in trace_data.tags
    assert "workflow" in trace_data.tags

    # Verify add_span (Workflow Run Span)
    assert trace_instance.add_span.call_count >= 1
    # First span should be workflow run span because message_id is present
    workflow_span = trace_instance.add_span.call_args_list[0][1]["langfuse_span_data"]
    assert workflow_span.id == "run-1"
    assert workflow_span.name == TraceTaskName.WORKFLOW_TRACE

    # Verify Generation for LLM node
    trace_instance.add_generation.assert_called_once()
    gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"]
    assert gen_data.id == "node-llm"
    assert gen_data.usage.input == 10
    assert gen_data.usage.output == 20

    # Verify normal span for Other node
    # Second add_span call
    other_span = trace_instance.add_span.call_args_list[1][1]["langfuse_span_data"]
    assert other_span.id == "node-other"
    assert other_span.level == LevelEnum.ERROR


def test_workflow_trace_no_message_id(trace_instance, monkeypatch):
    trace_info = WorkflowTraceInfo(
        workflow_id="wf-1",
        tenant_id="tenant-1",
        workflow_run_id="run-1",
        workflow_run_elapsed_time=1.0,
        workflow_run_status="succeeded",
        workflow_run_inputs={},
        workflow_run_outputs={},
        workflow_run_version="1.0",
        total_tokens=0,
        file_list=[],
        query="",
        message_id=None,
        conversation_id="conv-1",
        start_time=_dt(),
        end_time=_dt(),
        trace_id=None,  # Should fallback to workflow_run_id
        metadata={"app_id": "app-1"},
        workflow_app_log_id="log-1",
        error="",
    )

    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.sessionmaker", lambda bind: lambda: MagicMock())
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.db", MagicMock(engine="engine"))
    repo = MagicMock()
    repo.get_by_workflow_run.return_value = []
    mock_factory = MagicMock()
    mock_factory.create_workflow_node_execution_repository.return_value = repo
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.DifyCoreRepositoryFactory", mock_factory)
    monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock())

    trace_instance.add_trace = MagicMock()
    trace_instance.workflow_trace(trace_info)

    trace_instance.add_trace.assert_called_once()
    trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"]
    assert trace_data.id == "run-1"
    assert trace_data.name == TraceTaskName.WORKFLOW_TRACE


def test_workflow_trace_missing_app_id(trace_instance, monkeypatch):
    trace_info = WorkflowTraceInfo(
        workflow_id="wf-1",
        tenant_id="tenant-1",
        workflow_run_id="run-1",
        workflow_run_elapsed_time=1.0,
        workflow_run_status="succeeded",
        workflow_run_inputs={},
        workflow_run_outputs={},
        workflow_run_version="1.0",
        total_tokens=0,
        file_list=[],
        query="",
        message_id=None,
        conversation_id="conv-1",
        start_time=_dt(),
        end_time=_dt(),
        metadata={},  # Missing app_id
        workflow_app_log_id="log-1",
        error="",
    )
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.sessionmaker", lambda bind: lambda: MagicMock())
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.db", MagicMock(engine="engine"))

    with pytest.raises(ValueError, match="No app_id found in trace_info metadata"):
        trace_instance.workflow_trace(trace_info)


def test_message_trace_basic(trace_instance, monkeypatch):
    message_data = MagicMock()
    message_data.id = "msg-1"
    message_data.from_account_id = "acc-1"
    message_data.from_end_user_id = None
    message_data.provider_response_latency = 0.5
    message_data.conversation_id = "conv-1"
    message_data.total_price = 0.01
    message_data.model_id = "gpt-4"
    message_data.answer = "hello"
    message_data.status = MessageStatus.NORMAL
    message_data.error = None

    trace_info = MessageTraceInfo(
        message_id="msg-1",
        message_data=message_data,
        inputs={"query": "hi"},
        outputs={"answer": "hello"},
        message_tokens=10,
        answer_tokens=20,
        total_tokens=30,
        start_time=_dt(),
        end_time=_dt() + timedelta(seconds=1),
        trace_id="trace-1",
        metadata={"foo": "bar"},
        conversation_mode="chat",
        conversation_model="gpt-4",
        file_list=[],
        error=None,
    )

    trace_instance.add_trace = MagicMock()
    trace_instance.add_generation = MagicMock()

    trace_instance.message_trace(trace_info)

    trace_instance.add_trace.assert_called_once()
    trace_instance.add_generation.assert_called_once()

    gen_data = trace_instance.add_generation.call_args[0][0]
    assert gen_data.name == "llm"
    assert gen_data.usage.total == 30


def test_message_trace_with_end_user(trace_instance, monkeypatch):
    message_data = MagicMock()
    message_data.id = "msg-1"
    message_data.from_account_id = "acc-1"
    message_data.from_end_user_id = "end-user-1"
    message_data.conversation_id = "conv-1"
    message_data.status = MessageStatus.NORMAL
    message_data.model_id = "gpt-4"
    message_data.error = ""
    message_data.answer = "hello"
    message_data.total_price = 0.0
    message_data.provider_response_latency = 0.1

    trace_info = MessageTraceInfo(
        message_id="msg-1",
        message_data=message_data,
        inputs={},
        outputs={},
        message_tokens=0,
        answer_tokens=0,
        total_tokens=0,
        start_time=_dt(),
        end_time=_dt(),
        metadata={},
        conversation_mode="chat",
        conversation_model="gpt-4",
        file_list=[],
        error=None,
    )

    # Mock DB session for EndUser lookup
    mock_end_user = MagicMock(spec=EndUser)
    mock_end_user.session_id = "session-id-123"

    mock_query = MagicMock()
    mock_query.where.return_value.first.return_value = mock_end_user
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.db.session.query", lambda model: mock_query)

    trace_instance.add_trace = MagicMock()
    trace_instance.add_generation = MagicMock()

    trace_instance.message_trace(trace_info)

    trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"]
    assert trace_data.user_id == "session-id-123"
    assert trace_data.metadata["user_id"] == "session-id-123"


def test_message_trace_none_data(trace_instance):
    trace_info = SimpleNamespace(message_data=None, file_list=[], metadata={})
    trace_instance.add_trace = MagicMock()
    trace_instance.message_trace(trace_info)
    trace_instance.add_trace.assert_not_called()


def test_moderation_trace(trace_instance):
    message_data = MagicMock()
    message_data.created_at = _dt()

    trace_info = ModerationTraceInfo(
        message_id="msg-1",
        message_data=message_data,
        inputs={"q": "hi"},
        action="stop",
        flagged=True,
        preset_response="blocked",
        start_time=None,
        end_time=None,
        metadata={"foo": "bar"},
        trace_id="trace-1",
        query="hi",
    )

    trace_instance.add_span = MagicMock()
    trace_instance.moderation_trace(trace_info)

    trace_instance.add_span.assert_called_once()
    span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"]
    assert span_data.name == TraceTaskName.MODERATION_TRACE
    assert span_data.output["flagged"] is True


def test_suggested_question_trace(trace_instance):
    message_data = MagicMock()
    message_data.status = MessageStatus.NORMAL
    message_data.error = None

    trace_info = SuggestedQuestionTraceInfo(
        message_id="msg-1",
        message_data=message_data,
        inputs="hi",
        suggested_question=["q1"],
        total_tokens=10,
        level="info",
        start_time=_dt(),
        end_time=_dt(),
        metadata={},
        trace_id="trace-1",
    )

    trace_instance.add_generation = MagicMock()
    trace_instance.suggested_question_trace(trace_info)

    trace_instance.add_generation.assert_called_once()
    gen_data = trace_instance.add_generation.call_args[1]["langfuse_generation_data"]
    assert gen_data.name == TraceTaskName.SUGGESTED_QUESTION_TRACE
    assert gen_data.usage.unit == UnitEnum.CHARACTERS


def test_dataset_retrieval_trace(trace_instance):
    message_data = MagicMock()
    message_data.created_at = _dt()
    message_data.updated_at = _dt()

    trace_info = DatasetRetrievalTraceInfo(
        message_id="msg-1",
        message_data=message_data,
        inputs="query",
        documents=[{"id": "doc1"}],
        start_time=None,
        end_time=None,
        metadata={},
        trace_id="trace-1",
    )

    trace_instance.add_span = MagicMock()
    trace_instance.dataset_retrieval_trace(trace_info)

    trace_instance.add_span.assert_called_once()
    span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"]
    assert span_data.name == TraceTaskName.DATASET_RETRIEVAL_TRACE
    assert span_data.output["documents"] == [{"id": "doc1"}]


def test_tool_trace(trace_instance):
    trace_info = ToolTraceInfo(
        message_id="msg-1",
        message_data=MagicMock(),
        inputs={},
        outputs={},
        tool_name="my_tool",
        tool_inputs={"a": 1},
        tool_outputs="result_string",
        time_cost=0.1,
        start_time=_dt(),
        end_time=_dt(),
        metadata={},
        trace_id="trace-1",
        tool_config={},
        tool_parameters={},
        error="some error",
    )

    trace_instance.add_span = MagicMock()
    trace_instance.tool_trace(trace_info)

    trace_instance.add_span.assert_called_once()
    span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"]
    assert span_data.name == "my_tool"
    assert span_data.level == LevelEnum.ERROR


def test_generate_name_trace(trace_instance):
    trace_info = GenerateNameTraceInfo(
        inputs={"q": "hi"},
        outputs={"name": "new"},
        tenant_id="tenant-1",
        conversation_id="conv-1",
        start_time=_dt(),
        end_time=_dt(),
        metadata={"m": 1},
    )

    trace_instance.add_trace = MagicMock()
    trace_instance.add_span = MagicMock()

    trace_instance.generate_name_trace(trace_info)

    trace_instance.add_trace.assert_called_once()
    trace_instance.add_span.assert_called_once()

    trace_data = trace_instance.add_trace.call_args[1]["langfuse_trace_data"]
    assert trace_data.name == TraceTaskName.GENERATE_NAME_TRACE
    assert trace_data.user_id == "tenant-1"

    span_data = trace_instance.add_span.call_args[1]["langfuse_span_data"]
    assert span_data.trace_id == "conv-1"


def test_add_trace_success(trace_instance):
    data = LangfuseTrace(id="t1", name="trace")
    trace_instance.add_trace(data)
    trace_instance.langfuse_client.trace.assert_called_once()


def test_add_trace_error(trace_instance):
    trace_instance.langfuse_client.trace.side_effect = Exception("error")
    data = LangfuseTrace(id="t1", name="trace")
    with pytest.raises(ValueError, match="LangFuse Failed to create trace: error"):
        trace_instance.add_trace(data)


def test_add_span_success(trace_instance):
    data = LangfuseSpan(id="s1", name="span", trace_id="t1")
    trace_instance.add_span(data)
    trace_instance.langfuse_client.span.assert_called_once()


def test_add_span_error(trace_instance):
    trace_instance.langfuse_client.span.side_effect = Exception("error")
    data = LangfuseSpan(id="s1", name="span", trace_id="t1")
    with pytest.raises(ValueError, match="LangFuse Failed to create span: error"):
        trace_instance.add_span(data)


def test_update_span(trace_instance):
    span = MagicMock()
    data = LangfuseSpan(id="s1", name="span", trace_id="t1")
    trace_instance.update_span(span, data)
    span.end.assert_called_once()


def test_add_generation_success(trace_instance):
    data = LangfuseGeneration(id="g1", name="gen", trace_id="t1")
    trace_instance.add_generation(data)
    trace_instance.langfuse_client.generation.assert_called_once()


def test_add_generation_error(trace_instance):
    trace_instance.langfuse_client.generation.side_effect = Exception("error")
    data = LangfuseGeneration(id="g1", name="gen", trace_id="t1")
    with pytest.raises(ValueError, match="LangFuse Failed to create generation: error"):
        trace_instance.add_generation(data)


def test_update_generation(trace_instance):
    gen = MagicMock()
    data = LangfuseGeneration(id="g1", name="gen", trace_id="t1")
    trace_instance.update_generation(gen, data)
    gen.end.assert_called_once()


def test_api_check_success(trace_instance):
    trace_instance.langfuse_client.auth_check.return_value = True
    assert trace_instance.api_check() is True


def test_api_check_error(trace_instance):
    trace_instance.langfuse_client.auth_check.side_effect = Exception("fail")
    with pytest.raises(ValueError, match="LangFuse API check failed: fail"):
        trace_instance.api_check()


def test_get_project_key_success(trace_instance):
    mock_data = MagicMock()
    mock_data.id = "proj-1"
    trace_instance.langfuse_client.client.projects.get.return_value = MagicMock(data=[mock_data])
    assert trace_instance.get_project_key() == "proj-1"


def test_get_project_key_error(trace_instance):
    trace_instance.langfuse_client.client.projects.get.side_effect = Exception("fail")
    with pytest.raises(ValueError, match="LangFuse get project key failed: fail"):
        trace_instance.get_project_key()


def test_moderation_trace_none(trace_instance):
    trace_info = ModerationTraceInfo(
        message_id="m",
        message_data=None,
        inputs={},
        action="s",
        flagged=False,
        preset_response="",
        query="",
        metadata={},
    )
    trace_instance.add_span = MagicMock()
    trace_instance.moderation_trace(trace_info)
    trace_instance.add_span.assert_not_called()


def test_suggested_question_trace_none(trace_instance):
    trace_info = SuggestedQuestionTraceInfo(
        message_id="m", message_data=None, inputs={}, suggested_question=[], total_tokens=0, level="i", metadata={}
    )
    trace_instance.add_generation = MagicMock()
    trace_instance.suggested_question_trace(trace_info)
    trace_instance.add_generation.assert_not_called()


def test_dataset_retrieval_trace_none(trace_instance):
    trace_info = DatasetRetrievalTraceInfo(message_id="m", message_data=None, inputs={}, documents=[], metadata={})
    trace_instance.add_span = MagicMock()
    trace_instance.dataset_retrieval_trace(trace_info)
    trace_instance.add_span.assert_not_called()


def test_langfuse_trace_entity_with_list_dict_input():
    # To cover lines 29-31 in langfuse_trace_entity.py
    # We need to mock replace_text_with_content or just check if it works
    # Actually replace_text_with_content is imported from core.ops.utils
    data = LangfuseTrace(id="t1", name="n", input=[{"text": "hello"}])
    assert isinstance(data.input, list)
    assert data.input[0]["content"] == "hello"


def test_workflow_trace_handles_usage_extraction_error(trace_instance, monkeypatch, caplog):
    # Setup trace info to trigger LLM node usage extraction
    trace_info = WorkflowTraceInfo(
        workflow_id="wf-1",
        tenant_id="t",
        workflow_run_id="r",
        workflow_run_elapsed_time=1.0,
        workflow_run_status="s",
        workflow_run_inputs={},
        workflow_run_outputs={},
        workflow_run_version="1",
        total_tokens=0,
        file_list=[],
        query="",
        message_id=None,
        conversation_id="c",
        start_time=_dt(),
        end_time=_dt(),
        metadata={"app_id": "app-1"},
        workflow_app_log_id="l",
        error="",
    )

    node = MagicMock()
    node.id = "n1"
    node.title = "LLM Node"
    node.node_type = BuiltinNodeTypes.LLM
    node.status = "succeeded"

    class BadDict(collections.UserDict):
        def get(self, key, default=None):
            if key == "usage":
                raise Exception("Usage extraction failed")
            return super().get(key, default)

    node.process_data = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]})
    node.created_at = _dt()
    node.elapsed_time = 0.1
    node.metadata = {}
    node.outputs = {}

    repo = MagicMock()
    repo.get_by_workflow_run.return_value = [node]
    mock_factory = MagicMock()
    mock_factory.create_workflow_node_execution_repository.return_value = repo
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.DifyCoreRepositoryFactory", mock_factory)
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.sessionmaker", lambda bind: lambda: MagicMock())
    monkeypatch.setattr("core.ops.langfuse_trace.langfuse_trace.db", MagicMock(engine="engine"))
    monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock())

    trace_instance.add_trace = MagicMock()
    trace_instance.add_generation = MagicMock()

    with caplog.at_level(logging.ERROR):
        trace_instance.workflow_trace(trace_info)

    assert "Failed to extract usage" in caplog.text
    trace_instance.add_generation.assert_called_once()
