import collections
from datetime import datetime, timedelta
from unittest.mock import MagicMock

import pytest
from dify_trace_langsmith.config import LangSmithConfig
from dify_trace_langsmith.entities.langsmith_trace_entity import (
    LangSmithRunModel,
    LangSmithRunType,
    LangSmithRunUpdateModel,
)
from dify_trace_langsmith.langsmith_trace import LangSmithDataTrace

from core.ops.entities.trace_entity import (
    DatasetRetrievalTraceInfo,
    GenerateNameTraceInfo,
    MessageTraceInfo,
    ModerationTraceInfo,
    SuggestedQuestionTraceInfo,
    ToolTraceInfo,
    TraceTaskName,
    WorkflowTraceInfo,
)
from graphon.enums import BuiltinNodeTypes, WorkflowNodeExecutionMetadataKey
from models import EndUser


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


@pytest.fixture
def langsmith_config():
    return LangSmithConfig(api_key="ls-123", project="default", endpoint="https://api.smith.langchain.com")


@pytest.fixture
def trace_instance(langsmith_config, monkeypatch: pytest.MonkeyPatch):
    # Mock LangSmith client
    mock_client = MagicMock()
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.Client", lambda **kwargs: mock_client)

    instance = LangSmithDataTrace(langsmith_config)
    return instance


def test_init(langsmith_config, monkeypatch: pytest.MonkeyPatch):
    mock_client_class = MagicMock()
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.Client", mock_client_class)
    monkeypatch.setenv("FILES_URL", "http://test.url")

    instance = LangSmithDataTrace(langsmith_config)

    mock_client_class.assert_called_once_with(api_key=langsmith_config.api_key, api_url=langsmith_config.endpoint)
    assert instance.langsmith_key == langsmith_config.api_key
    assert instance.project_name == langsmith_config.project
    assert instance.file_base_url == "http://test.url"


def test_trace_dispatch(trace_instance, monkeypatch: pytest.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(trace_instance, monkeypatch: pytest.MonkeyPatch):
    # Setup trace info
    workflow_data = MagicMock()
    workflow_data.created_at = _dt()
    workflow_data.finished_at = _dt() + timedelta(seconds=1)

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

    # Mock dependencies
    mock_session = MagicMock()
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.sessionmaker", lambda bind: lambda: mock_session)
    monkeypatch.setattr("dify_trace_langsmith.langsmith_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 = {WorkflowNodeExecutionMetadataKey.TOTAL_TOKENS: 30}

    node_other = MagicMock()
    node_other.id = "node-other"
    node_other.title = "Tool Node"
    node_other.node_type = BuiltinNodeTypes.TOOL
    node_other.status = "succeeded"
    node_other.process_data = None
    node_other.inputs = {"tool_input": "val"}
    node_other.outputs = {"tool_output": "val"}
    node_other.created_at = None  # Trigger datetime.now()
    node_other.elapsed_time = 0.2
    node_other.metadata = {}

    node_retrieval = MagicMock()
    node_retrieval.id = "node-retrieval"
    node_retrieval.title = "Retrieval Node"
    node_retrieval.node_type = BuiltinNodeTypes.KNOWLEDGE_RETRIEVAL
    node_retrieval.status = "succeeded"
    node_retrieval.process_data = None
    node_retrieval.inputs = {"query": "val"}
    node_retrieval.outputs = {"results": "val"}
    node_retrieval.created_at = _dt()
    node_retrieval.elapsed_time = 0.2
    node_retrieval.metadata = {}

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

    mock_factory = MagicMock()
    mock_factory.create_workflow_node_execution_repository.return_value = repo
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.DifyCoreRepositoryFactory", mock_factory)

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

    trace_instance.add_run = MagicMock()

    trace_instance.workflow_trace(trace_info)

    # Verify add_run calls
    # 1. message run (id="msg-1")
    # 2. workflow run (id="run-1")
    # 3. node llm run (id="node-llm")
    # 4. node other run (id="node-other")
    # 5. node retrieval run (id="node-retrieval")
    assert trace_instance.add_run.call_count == 5

    call_args = [call[0][0] for call in trace_instance.add_run.call_args_list]

    assert call_args[0].id == "msg-1"
    assert call_args[0].name == TraceTaskName.MESSAGE_TRACE

    assert call_args[1].id == "run-1"
    assert call_args[1].name == TraceTaskName.WORKFLOW_TRACE
    assert call_args[1].parent_run_id == "msg-1"

    assert call_args[2].id == "node-llm"
    assert call_args[2].run_type == LangSmithRunType.llm

    assert call_args[3].id == "node-other"
    assert call_args[3].run_type == LangSmithRunType.tool

    assert call_args[4].id == "node-retrieval"
    assert call_args[4].run_type == LangSmithRunType.retriever


def test_workflow_trace_no_start_time(trace_instance, monkeypatch: pytest.MonkeyPatch):
    workflow_data = MagicMock()
    workflow_data.created_at = _dt()
    workflow_data.finished_at = _dt() + timedelta(seconds=1)

    trace_info = WorkflowTraceInfo(
        tenant_id="tenant-1",
        workflow_id="wf-1",
        workflow_run_id="run-1",
        workflow_run_inputs={},
        workflow_run_outputs={},
        workflow_run_status="succeeded",
        workflow_run_version="1.0",
        workflow_run_elapsed_time=1.0,
        total_tokens=10,
        file_list=[],
        query="hi",
        message_id="msg-1",
        conversation_id="conv-1",
        start_time=None,
        end_time=None,
        trace_id="trace-1",
        metadata={"app_id": "app-1"},
        workflow_app_log_id="log-1",
        error="",
        workflow_data=workflow_data,
    )

    mock_session = MagicMock()
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.sessionmaker", lambda bind: lambda: mock_session)
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.db", MagicMock(engine="engine"))
    repo = MagicMock()
    repo.get_by_workflow_execution.return_value = []
    mock_factory = MagicMock()
    mock_factory.create_workflow_node_execution_repository.return_value = repo
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.DifyCoreRepositoryFactory", mock_factory)
    monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock())

    trace_instance.add_run = MagicMock()
    trace_instance.workflow_trace(trace_info)
    assert trace_instance.add_run.called


def test_workflow_trace_missing_app_id(trace_instance, monkeypatch: pytest.MonkeyPatch):
    trace_info = MagicMock(spec=WorkflowTraceInfo)
    trace_info.trace_id = "trace-1"
    trace_info.message_id = None
    trace_info.workflow_run_id = "run-1"
    trace_info.start_time = None
    trace_info.workflow_data = MagicMock()
    trace_info.workflow_data.created_at = _dt()
    trace_info.metadata = {}  # Empty metadata
    trace_info.workflow_app_log_id = "log-1"
    trace_info.file_list = []
    trace_info.total_tokens = 0
    trace_info.workflow_run_inputs = {}
    trace_info.workflow_run_outputs = {}
    trace_info.error = ""

    mock_session = MagicMock()
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.sessionmaker", lambda bind: lambda: mock_session)
    monkeypatch.setattr("dify_trace_langsmith.langsmith_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(trace_instance, monkeypatch: pytest.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.answer = "hello answer"

    trace_info = MessageTraceInfo(
        message_id="msg-1",
        message_data=message_data,
        inputs={"input": "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,
        message_file_data=MagicMock(url="file-url"),
    )

    # Mock EndUser lookup
    mock_end_user = MagicMock(spec=EndUser)
    mock_end_user.session_id = "session-id-123"
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.db.session.get", lambda model, pk: mock_end_user)

    trace_instance.add_run = MagicMock()

    trace_instance.message_trace(trace_info)

    # 1. message run
    # 2. llm run
    assert trace_instance.add_run.call_count == 2

    call_args = [call[0][0] for call in trace_instance.add_run.call_args_list]
    assert call_args[0].id == "msg-1"
    assert call_args[0].extra["metadata"]["end_user_id"] == "session-id-123"
    assert call_args[1].parent_run_id == "msg-1"
    assert call_args[1].name == "llm"


def test_message_trace_no_data(trace_instance):
    trace_info = MagicMock(spec=MessageTraceInfo)
    trace_info.message_data = None
    trace_info.file_list = []
    trace_info.message_file_data = None
    trace_info.metadata = {}
    trace_instance.add_run = MagicMock()
    trace_instance.message_trace(trace_info)
    trace_instance.add_run.assert_not_called()


def test_moderation_trace_no_data(trace_instance):
    trace_info = MagicMock(spec=ModerationTraceInfo)
    trace_info.message_data = None
    trace_instance.add_run = MagicMock()
    trace_instance.moderation_trace(trace_info)
    trace_instance.add_run.assert_not_called()


def test_suggested_question_trace_no_data(trace_instance):
    trace_info = MagicMock(spec=SuggestedQuestionTraceInfo)
    trace_info.message_data = None
    trace_instance.add_run = MagicMock()
    trace_instance.suggested_question_trace(trace_info)
    trace_instance.add_run.assert_not_called()


def test_dataset_retrieval_trace_no_data(trace_instance):
    trace_info = MagicMock(spec=DatasetRetrievalTraceInfo)
    trace_info.message_data = None
    trace_instance.add_run = MagicMock()
    trace_instance.dataset_retrieval_trace(trace_info)
    trace_instance.add_run.assert_not_called()


def test_moderation_trace(trace_instance):
    message_data = MagicMock()
    message_data.created_at = _dt()
    message_data.updated_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={},
        trace_id="trace-1",
        query="hi",
    )

    trace_instance.add_run = MagicMock()
    trace_instance.moderation_trace(trace_info)
    trace_instance.add_run.assert_called_once()
    assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.MODERATION_TRACE


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

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

    trace_instance.add_run = MagicMock()
    trace_instance.suggested_question_trace(trace_info)
    trace_instance.add_run.assert_called_once()
    assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.SUGGESTED_QUESTION_TRACE


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_run = MagicMock()
    trace_instance.dataset_retrieval_trace(trace_info)
    trace_instance.add_run.assert_called_once()
    assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.DATASET_RETRIEVAL_TRACE


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",
        time_cost=0.1,
        start_time=_dt(),
        end_time=_dt(),
        metadata={},
        trace_id="trace-1",
        tool_config={},
        tool_parameters={},
        file_url="http://file",
    )

    trace_instance.add_run = MagicMock()
    trace_instance.tool_trace(trace_info)
    trace_instance.add_run.assert_called_once()
    assert trace_instance.add_run.call_args[0][0].name == "my_tool"


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=None,
        end_time=None,
        metadata={},
        trace_id="trace-1",
    )

    trace_instance.add_run = MagicMock()
    trace_instance.generate_name_trace(trace_info)
    trace_instance.add_run.assert_called_once()
    assert trace_instance.add_run.call_args[0][0].name == TraceTaskName.GENERATE_NAME_TRACE


def test_add_run_success(trace_instance):
    run_data = LangSmithRunModel(
        id="run-1", name="test", inputs={}, outputs={}, run_type=LangSmithRunType.tool, start_time=_dt()
    )
    trace_instance.project_id = "proj-1"
    trace_instance.add_run(run_data)
    trace_instance.langsmith_client.create_run.assert_called_once()
    args, kwargs = trace_instance.langsmith_client.create_run.call_args
    assert kwargs["session_id"] == "proj-1"


def test_add_run_error(trace_instance):
    run_data = LangSmithRunModel(id="run-1", name="test", run_type=LangSmithRunType.tool, start_time=_dt())
    trace_instance.langsmith_client.create_run.side_effect = Exception("failed")
    with pytest.raises(ValueError, match="LangSmith Failed to create run: failed"):
        trace_instance.add_run(run_data)


def test_update_run_success(trace_instance):
    update_data = LangSmithRunUpdateModel(run_id="run-1", outputs={"out": "val"})
    trace_instance.update_run(update_data)
    trace_instance.langsmith_client.update_run.assert_called_once()


def test_update_run_error(trace_instance):
    update_data = LangSmithRunUpdateModel(run_id="run-1")
    trace_instance.langsmith_client.update_run.side_effect = Exception("failed")
    with pytest.raises(ValueError, match="LangSmith Failed to update run: failed"):
        trace_instance.update_run(update_data)


def test_workflow_trace_usage_extraction_error(trace_instance, monkeypatch: pytest.MonkeyPatch, caplog):
    workflow_data = MagicMock()
    workflow_data.created_at = _dt()
    workflow_data.finished_at = _dt() + timedelta(seconds=1)

    trace_info = WorkflowTraceInfo(
        tenant_id="tenant-1",
        workflow_id="wf-1",
        workflow_run_id="run-1",
        workflow_run_inputs={},
        workflow_run_outputs={},
        workflow_run_status="succeeded",
        workflow_run_version="1.0",
        workflow_run_elapsed_time=1.0,
        total_tokens=100,
        file_list=[],
        query="hi",
        message_id="msg-1",
        conversation_id="conv-1",
        start_time=_dt(),
        end_time=_dt(),
        trace_id="trace-1",
        metadata={"app_id": "app-1"},
        workflow_app_log_id="log-1",
        error="",
        workflow_data=workflow_data,
    )

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

    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 = BadDict({"model_mode": "chat", "model_name": "gpt-4", "usage": True, "prompts": ["p"]})
    node_llm.inputs = {}
    node_llm.outputs = {}
    node_llm.created_at = _dt()
    node_llm.elapsed_time = 0.5
    node_llm.metadata = {}

    repo = MagicMock()
    repo.get_by_workflow_execution.return_value = [node_llm]

    mock_factory = MagicMock()
    mock_factory.create_workflow_node_execution_repository.return_value = repo
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.DifyCoreRepositoryFactory", mock_factory)
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.sessionmaker", lambda bind: lambda: MagicMock())
    monkeypatch.setattr("dify_trace_langsmith.langsmith_trace.db", MagicMock(engine="engine"))
    monkeypatch.setattr(trace_instance, "get_service_account_with_tenant", lambda app_id: MagicMock())

    trace_instance.add_run = MagicMock()

    import logging

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

    assert "Failed to extract usage" in caplog.text


def test_api_check_success(trace_instance):
    assert trace_instance.api_check() is True
    assert trace_instance.langsmith_client.create_project.called
    assert trace_instance.langsmith_client.delete_project.called


def test_api_check_error(trace_instance):
    trace_instance.langsmith_client.create_project.side_effect = Exception("error")
    with pytest.raises(ValueError, match="LangSmith API check failed: error"):
        trace_instance.api_check()


def test_get_project_url_success(trace_instance):
    trace_instance.langsmith_client.get_run_url.return_value = "https://smith.langchain.com/o/org/p/proj/r/run"
    url = trace_instance.get_project_url()
    assert url == "https://smith.langchain.com/o/org/p/proj"


def test_get_project_url_error(trace_instance):
    trace_instance.langsmith_client.get_run_url.side_effect = Exception("error")
    with pytest.raises(ValueError, match="LangSmith get run url failed: error"):
        trace_instance.get_project_url()
