import json
from unittest.mock import MagicMock

from opentelemetry.trace import Link, StatusCode

from core.ops.aliyun_trace.entities.semconv import (
    GEN_AI_FRAMEWORK,
    GEN_AI_SESSION_ID,
    GEN_AI_SPAN_KIND,
    GEN_AI_USER_ID,
    INPUT_VALUE,
    OUTPUT_VALUE,
)
from core.ops.aliyun_trace.utils import (
    create_common_span_attributes,
    create_links_from_trace_id,
    create_status_from_error,
    extract_retrieval_documents,
    format_input_messages,
    format_output_messages,
    format_retrieval_documents,
    get_user_id_from_message_data,
    get_workflow_node_status,
    serialize_json_data,
)
from core.rag.models.document import Document
from dify_graph.entities import WorkflowNodeExecution
from dify_graph.enums import WorkflowNodeExecutionStatus
from models import EndUser


def test_get_user_id_from_message_data_no_end_user(monkeypatch):
    message_data = MagicMock()
    message_data.from_account_id = "account_id"
    message_data.from_end_user_id = None

    assert get_user_id_from_message_data(message_data) == "account_id"


def test_get_user_id_from_message_data_with_end_user(monkeypatch):
    message_data = MagicMock()
    message_data.from_account_id = "account_id"
    message_data.from_end_user_id = "end_user_id"

    end_user_data = MagicMock(spec=EndUser)
    end_user_data.session_id = "session_id"

    mock_query = MagicMock()
    mock_query.where.return_value.first.return_value = end_user_data

    mock_session = MagicMock()
    mock_session.query.return_value = mock_query

    from core.ops.aliyun_trace.utils import db

    monkeypatch.setattr(db, "session", mock_session)

    assert get_user_id_from_message_data(message_data) == "session_id"


def test_get_user_id_from_message_data_end_user_not_found(monkeypatch):
    message_data = MagicMock()
    message_data.from_account_id = "account_id"
    message_data.from_end_user_id = "end_user_id"

    mock_query = MagicMock()
    mock_query.where.return_value.first.return_value = None

    mock_session = MagicMock()
    mock_session.query.return_value = mock_query

    from core.ops.aliyun_trace.utils import db

    monkeypatch.setattr(db, "session", mock_session)

    assert get_user_id_from_message_data(message_data) == "account_id"


def test_create_status_from_error():
    # Case OK
    status_ok = create_status_from_error(None)
    assert status_ok.status_code == StatusCode.OK

    # Case Error
    status_err = create_status_from_error("some error")
    assert status_err.status_code == StatusCode.ERROR
    assert status_err.description == "some error"


def test_get_workflow_node_status():
    node_execution = MagicMock(spec=WorkflowNodeExecution)

    # SUCCEEDED
    node_execution.status = WorkflowNodeExecutionStatus.SUCCEEDED
    status = get_workflow_node_status(node_execution)
    assert status.status_code == StatusCode.OK

    # FAILED
    node_execution.status = WorkflowNodeExecutionStatus.FAILED
    node_execution.error = "node fail"
    status = get_workflow_node_status(node_execution)
    assert status.status_code == StatusCode.ERROR
    assert status.description == "node fail"

    # EXCEPTION
    node_execution.status = WorkflowNodeExecutionStatus.EXCEPTION
    node_execution.error = "node exception"
    status = get_workflow_node_status(node_execution)
    assert status.status_code == StatusCode.ERROR
    assert status.description == "node exception"

    # UNSET/OTHER
    node_execution.status = WorkflowNodeExecutionStatus.RUNNING
    status = get_workflow_node_status(node_execution)
    assert status.status_code == StatusCode.UNSET


def test_create_links_from_trace_id(monkeypatch):
    # Mock create_link
    mock_link = MagicMock(spec=Link)
    import core.ops.aliyun_trace.data_exporter.traceclient

    monkeypatch.setattr(core.ops.aliyun_trace.data_exporter.traceclient, "create_link", lambda trace_id_str: mock_link)

    # Trace ID None
    assert create_links_from_trace_id(None) == []

    # Trace ID Present
    links = create_links_from_trace_id("trace_id")
    assert len(links) == 1
    assert links[0] == mock_link


def test_extract_retrieval_documents():
    doc1 = MagicMock(spec=Document)
    doc1.page_content = "content1"
    doc1.metadata = {"dataset_id": "ds1", "doc_id": "di1", "document_id": "dd1", "score": 0.9}

    doc2 = MagicMock(spec=Document)
    doc2.page_content = "content2"
    doc2.metadata = {"dataset_id": "ds2"}  # Missing some keys

    documents = [doc1, doc2]
    extracted = extract_retrieval_documents(documents)

    assert len(extracted) == 2
    assert extracted[0]["content"] == "content1"
    assert extracted[0]["metadata"]["dataset_id"] == "ds1"
    assert extracted[0]["score"] == 0.9

    assert extracted[1]["content"] == "content2"
    assert extracted[1]["metadata"]["dataset_id"] == "ds2"
    assert extracted[1]["metadata"]["doc_id"] is None
    assert extracted[1]["score"] is None


def test_serialize_json_data():
    data = {"a": 1}
    # Test ensure_ascii default (False)
    assert serialize_json_data(data) == json.dumps(data, ensure_ascii=False)
    # Test ensure_ascii True
    assert serialize_json_data(data, ensure_ascii=True) == json.dumps(data, ensure_ascii=True)


def test_create_common_span_attributes():
    attrs = create_common_span_attributes(
        session_id="s1", user_id="u1", span_kind="kind1", framework="fw1", inputs="in1", outputs="out1"
    )
    assert attrs[GEN_AI_SESSION_ID] == "s1"
    assert attrs[GEN_AI_USER_ID] == "u1"
    assert attrs[GEN_AI_SPAN_KIND] == "kind1"
    assert attrs[GEN_AI_FRAMEWORK] == "fw1"
    assert attrs[INPUT_VALUE] == "in1"
    assert attrs[OUTPUT_VALUE] == "out1"


def test_format_retrieval_documents():
    # Not a list
    assert format_retrieval_documents("not a list") == []

    # Valid list
    docs = [
        {"metadata": {"score": 0.8, "document_id": "doc1", "source": "src1"}, "content": "c1", "title": "t1"},
        {
            "metadata": {"_source": "src2", "doc_metadata": {"extra": "val"}},
            "content": "c2",
            # Missing title
        },
        "not a dict",  # Should be skipped
    ]
    formatted = format_retrieval_documents(docs)

    assert len(formatted) == 2
    assert formatted[0]["document"]["content"] == "c1"
    assert formatted[0]["document"]["metadata"]["title"] == "t1"
    assert formatted[0]["document"]["metadata"]["source"] == "src1"
    assert formatted[0]["document"]["score"] == 0.8
    assert formatted[0]["document"]["id"] == "doc1"

    assert formatted[1]["document"]["content"] == "c2"
    assert formatted[1]["document"]["metadata"]["source"] == "src2"
    assert formatted[1]["document"]["metadata"]["extra"] == "val"
    assert "title" not in formatted[1]["document"]["metadata"]
    assert formatted[1]["document"]["score"] == 0.0  # Default

    # Exception handling
    # We can trigger an exception by passing something that causes an error in the loop logic,
    # but the try/except covers the whole function.
    # Passing a list that contains something that throws when calling .get() - though dicts won't.
    # Let's mock a dict that raises on get.
    class BadDict:
        def get(self, *args, **kwargs):
            raise Exception("boom")

    assert format_retrieval_documents([BadDict()]) == []


def test_format_input_messages():
    # Not a dict
    assert format_input_messages(None) == serialize_json_data([])

    # No prompts
    assert format_input_messages({}) == serialize_json_data([])

    # Valid prompts
    process_data = {
        "prompts": [
            {"role": "user", "text": "hello"},
            {"role": "assistant", "text": "hi"},
            {"role": "system", "text": "be helpful"},
            {"role": "tool", "text": "result"},
            {"role": "invalid", "text": "skip me"},
            "not a dict",
            {"role": "user", "text": ""},  # Empty text, should be skipped? Code says `if text: message = ...`
        ]
    }
    result = format_input_messages(process_data)
    result_list = json.loads(result)

    assert len(result_list) == 4
    assert result_list[0]["role"] == "user"
    assert result_list[0]["parts"][0]["content"] == "hello"
    assert result_list[1]["role"] == "assistant"
    assert result_list[2]["role"] == "system"
    assert result_list[3]["role"] == "tool"

    # Exception path
    assert format_input_messages({"prompts": [None]}) == serialize_json_data([])


def test_format_output_messages():
    # Not a dict
    assert format_output_messages(None) == serialize_json_data([])

    # No text
    assert format_output_messages({"finish_reason": "stop"}) == serialize_json_data([])

    # Valid
    outputs = {"text": "done", "finish_reason": "length"}
    result = format_output_messages(outputs)
    result_list = json.loads(result)
    assert len(result_list) == 1
    assert result_list[0]["role"] == "assistant"
    assert result_list[0]["parts"][0]["content"] == "done"
    assert result_list[0]["finish_reason"] == "length"

    # Invalid finish reason
    outputs2 = {"text": "done", "finish_reason": "unknown"}
    result2 = format_output_messages(outputs2)
    result_list2 = json.loads(result2)
    assert result_list2[0]["finish_reason"] == "stop"

    # Exception path
    # Trigger exception in serialize_json_data by passing non-serializable
    assert format_output_messages({"text": MagicMock()}) == serialize_json_data([])
