from __future__ import annotations

from langchain_core.messages import AIMessage, HumanMessage

from deerflow.agents.middlewares.message_timestamp_middleware import MessageTimestampMiddleware


def test_after_model_stamps_missing_message_timestamps():
    middleware = MessageTimestampMiddleware()
    state = {
        "messages": [
            HumanMessage(content="hello"),
            AIMessage(content="hi"),
        ]
    }

    middleware.after_model(state, runtime=None)  # type: ignore[arg-type]

    timestamps = [msg.additional_kwargs.get("deerflow_created_at") for msg in state["messages"]]
    assert all(isinstance(ts, str) and ts.endswith("+08:00") for ts in timestamps)


def test_after_model_keeps_existing_timestamp():
    middleware = MessageTimestampMiddleware()
    human = HumanMessage(content="hello")
    human.additional_kwargs["deerflow_created_at"] = "2026-04-22T01:00:00.000Z"
    state = {"messages": [human, AIMessage(content="hi")]}

    middleware.after_model(state, runtime=None)  # type: ignore[arg-type]

    assert state["messages"][0].additional_kwargs["deerflow_created_at"] == "2026-04-22T01:00:00.000Z"