from __future__ import annotations

from types import SimpleNamespace

import pytest

from deerflow.agents.middlewares.title_middleware import (
    TitleMiddleware,
    finalize_thread_title,
)


class _FakeCheckpointer:
    def __init__(self, checkpoint: dict, metadata: dict | None = None):
        self._tuple = SimpleNamespace(
            checkpoint=checkpoint,
            metadata=metadata or {},
            config={"configurable": {"thread_id": "thread-1", "checkpoint_ns": ""}},
        )
        self.aput_calls: list[tuple[dict, dict, dict, dict]] = []
        self._version = 0

    async def aget_tuple(self, _config):
        return self._tuple

    def get_next_version(self, _current, _channel):
        self._version += 1
        return self._version

    async def aput(self, config, checkpoint, metadata, versions):
        self.aput_calls.append((config, checkpoint, metadata, versions))
        self._tuple = SimpleNamespace(
            checkpoint=checkpoint,
            metadata=metadata,
            config={
                "configurable": {
                    "thread_id": config["configurable"]["thread_id"],
                    "checkpoint_ns": config["configurable"]["checkpoint_ns"],
                    "checkpoint_id": "new-ckpt",
                }
            },
        )
        return self._tuple.config


class _FakeThreadStore:
    def __init__(self):
        self.updated: list[tuple[str, str, str | None]] = []

    async def update_display_name(self, thread_id: str, title: str, *, user_id=None):
        self.updated.append((thread_id, title, user_id))


def _source_state() -> dict:
    return {
        "messages": [
            SimpleNamespace(type="human", content="please analyze the first topic"),
            SimpleNamespace(type="ai", content="this is the first answer"),
        ],
        "title": "please analyze the first topic",
        "title_provisional": True,
    }


def _dict_source_state() -> dict:
    return {
        "messages": [
            {"type": "human", "content": "please analyze the first topic"},
            {"type": "ai", "content": "this is the first answer"},
        ],
        "title": "please analyze the first topic",
        "title_provisional": True,
    }


@pytest.mark.asyncio
async def test_finalize_thread_title_updates_checkpoint_and_display_name(monkeypatch) -> None:
    async def _fake_generate(self, state):
        assert state["title_provisional"] is True
        return {"title": "first topic analysis", "title_provisional": False}

    monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)

    checkpointer = _FakeCheckpointer({"channel_values": _source_state()})
    thread_store = _FakeThreadStore()

    updated = await finalize_thread_title(
        checkpointer=checkpointer,
        thread_store=thread_store,
        thread_id="thread-1",
        source_state=_source_state(),
    )

    assert updated is True
    assert len(checkpointer.aput_calls) == 1
    _, written_checkpoint, _, new_versions = checkpointer.aput_calls[0]
    assert written_checkpoint["channel_values"]["title"] == "first topic analysis"
    assert written_checkpoint["channel_values"]["title_provisional"] is False
    assert set(new_versions) == {"title", "title_provisional"}
    assert thread_store.updated == [("thread-1", "first topic analysis", None)]


@pytest.mark.asyncio
async def test_finalize_thread_title_skips_if_current_thread_already_finalized(monkeypatch) -> None:
    async def _fake_generate(self, state):
        return {"title": "first topic analysis", "title_provisional": False}

    monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)

    checkpointer = _FakeCheckpointer(
        {
            "channel_values": {
                **_source_state(),
                "title": "manually renamed title",
                "title_provisional": False,
            }
        }
    )
    thread_store = _FakeThreadStore()

    updated = await finalize_thread_title(
        checkpointer=checkpointer,
        thread_store=thread_store,
        thread_id="thread-1",
        source_state=_source_state(),
    )

    assert updated is False
    assert checkpointer.aput_calls == []
    assert thread_store.updated == []


def test_sync_after_model_keeps_provisional_title_for_background_finalize() -> None:
    middleware = TitleMiddleware()

    result = middleware.after_model(_source_state(), runtime=SimpleNamespace())

    assert result is None


@pytest.mark.asyncio
async def test_async_after_model_keeps_provisional_title_for_background_finalize() -> None:
    middleware = TitleMiddleware()

    result = await middleware.aafter_model(_source_state(), runtime=SimpleNamespace())

    assert result is None


@pytest.mark.asyncio
async def test_finalize_thread_title_accepts_serialized_dict_messages(monkeypatch) -> None:
    async def _fake_generate(self, state):
        return {"title": "first topic analysis", "title_provisional": False}

    monkeypatch.setattr(TitleMiddleware, "_agenerate_title_result", _fake_generate)

    checkpointer = _FakeCheckpointer({"channel_values": _dict_source_state()})

    updated = await finalize_thread_title(
        checkpointer=checkpointer,
        thread_store=None,
        thread_id="thread-1",
        source_state=_dict_source_state(),
    )

    assert updated is True