from __future__ import annotations

from types import SimpleNamespace

import pytest

import deerflow.agents.middlewares.title_middleware as title_middleware_module
from deerflow.agents.middlewares.title_middleware import (
    TitleMiddleware,
    finalize_thread_title,
)
from deerflow.config.title_config import TitleConfig


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,
    }


def _knowledge_base_source_state(question: str = "这个库讲了什么") -> dict:
    return {
        "messages": [
            SimpleNamespace(
                type="human",
                content=question,
                additional_kwargs={
                    "llmwiki_knowledge_base_context": [
                        {"id": "kb-1", "name": "对话沉淀"}
                    ]
                },
            ),
            SimpleNamespace(type="ai", content="这个知识库汇总了对话资料。"),
        ],
        "title": question,
        "title_provisional": True,
    }


class _FakeTitleModel:
    def __init__(self, content: str):
        self.content = content
        self.extra_body: dict = {}
        self.prompt = ""

    async def ainvoke(self, prompt, config):
        self.prompt = prompt
        return SimpleNamespace(content=self.content)


def _title_config() -> TitleConfig:
    return TitleConfig(
        enabled=True,
        max_words=6,
        max_chars=60,
        prompt_template="User: {user_msg}\nAssistant: {assistant_msg}\nMax: {max_words}",
    )


def test_admin_question_title_switch_finalizes_question_before_first_model_call(monkeypatch) -> None:
    monkeypatch.setattr(
        title_middleware_module,
        "get_conversation_context_settings",
        lambda: SimpleNamespace(use_question_as_title=True),
    )
    middleware = TitleMiddleware(
        title_config=_title_config().model_copy(update={"enabled": False})
    )
    state = {
        "messages": [SimpleNamespace(type="human", content="请帮我分析这份合同")],
    }

    result = middleware.before_model(state, runtime=SimpleNamespace())

    assert result == {
        "title": "请帮我分析这份合同",
        "title_provisional": False,
    }


def test_admin_question_title_switch_off_keeps_async_title_generation(monkeypatch) -> None:
    monkeypatch.setattr(
        title_middleware_module,
        "get_conversation_context_settings",
        lambda: SimpleNamespace(use_question_as_title=False),
    )
    middleware = TitleMiddleware(title_config=_title_config())
    state = {
        "messages": [SimpleNamespace(type="human", content="请帮我分析这份合同")],
    }

    result = middleware.before_model(state, runtime=SimpleNamespace())

    assert result == {
        "title": "请帮我分析这份合同",
        "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


@pytest.mark.asyncio
async def test_deictic_knowledge_base_question_rejects_invented_model_title(
    monkeypatch,
) -> None:
    model = _FakeTitleModel("Rust异步编程库简介")
    monkeypatch.setattr(
        title_middleware_module,
        "create_chat_model",
        lambda **_kwargs: model,
    )
    middleware = TitleMiddleware(title_config=_title_config())

    result = await middleware._agenerate_title_result(
        _knowledge_base_source_state()
    )

    assert result == {
        "title": "对话沉淀：知识库概览",
        "title_provisional": False,
    }
    assert "服务端已验证的当前知识库名称" in model.prompt
    assert '"对话沉淀"' in model.prompt
    assert "禁止猜测或改写为其他库名" in model.prompt


@pytest.mark.asyncio
async def test_deictic_knowledge_base_question_keeps_title_with_selected_name(
    monkeypatch,
) -> None:
    model = _FakeTitleModel("对话沉淀内容概览")
    monkeypatch.setattr(
        title_middleware_module,
        "create_chat_model",
        lambda **_kwargs: model,
    )
    middleware = TitleMiddleware(title_config=_title_config())

    result = await middleware._agenerate_title_result(
        _knowledge_base_source_state("请介绍本知识库的主要内容")
    )

    assert result == {
        "title": "对话沉淀内容概览",
        "title_provisional": False,
    }


@pytest.mark.asyncio
async def test_ordinary_question_title_generation_is_unchanged(monkeypatch) -> None:
    model = _FakeTitleModel("量子计算发展概览")
    monkeypatch.setattr(
        title_middleware_module,
        "create_chat_model",
        lambda **_kwargs: model,
    )
    middleware = TitleMiddleware(title_config=_title_config())
    state = {
        "messages": [
            SimpleNamespace(type="human", content="请介绍量子计算的发展"),
            SimpleNamespace(type="ai", content="量子计算经历了多个发展阶段。"),
        ],
        "title": "请介绍量子计算的发展",
        "title_provisional": True,
    }

    result = await middleware._agenerate_title_result(state)

    assert result == {
        "title": "量子计算发展概览",
        "title_provisional": False,
    }
    assert "服务端已验证的当前知识库名称" not in model.prompt


def test_serialized_user_message_uses_knowledge_base_aware_local_title() -> None:
    middleware = TitleMiddleware(title_config=_title_config())
    state = {
        "messages": [
            {
                "type": "human",
                "content": "当前库主要说了什么？",
                "additional_kwargs": {
                    "llmwiki_knowledge_base_context": [
                        {"id": "kb-1", "name": "政策资料"}
                    ]
                },
            },
            {"type": "ai", "content": "该库收录了政策资料。"},
        ],
        "title": "当前库主要说了什么？",
        "title_provisional": True,
    }

    assert middleware._generate_title_result(state) == {
        "title": "政策资料：知识库概览",
        "title_provisional": False,
    }


def test_title_prompt_uses_last_non_empty_assistant_body_without_tool_calls() -> None:
    middleware = TitleMiddleware(title_config=_title_config())
    state = {
        "messages": [
            SimpleNamespace(type="human", content="请分析材料"),
            SimpleNamespace(
                type="ai",
                content="准备读取文件",
                tool_calls=[{"name": "read_file", "args": {}}],
            ),
            SimpleNamespace(type="tool", content="工具返回内容"),
            SimpleNamespace(type="ai", content="这是基于材料形成的最终结论。"),
            SimpleNamespace(type="ai", content="<think>仅内部思考</think>"),
        ],
        "title": "请分析材料",
        "title_provisional": True,
    }

    prompt, _ = middleware._build_title_prompt(state)

    assert "这是基于材料形成的最终结论。" in prompt
    assert "准备读取文件" not in prompt
    assert "仅内部思考" not in prompt