import json
from collections.abc import Sequence
from typing import Any

from langchain.messages import ToolMessage
from langchain.tools import ToolRuntime, tool
from langgraph.config import get_stream_writer

from kbqa.agents.context import ChatRunContext, ResearchContext
from kbqa.rag.types import SourceCandidate


def serialize_source(source: SourceCandidate) -> dict[str, Any]:
    return {
        "chunk_id": source.chunk_id,
        "document_id": source.document_id,
        "filename": source.filename,
        "order_index": source.order_index,
        "page_number": source.page_number,
        "content": source.content,
        "score": source.score,
    }


@tool(response_format="content_and_artifact")
async def search_knowledge_base(
    query: str,
    runtime: ToolRuntime[ResearchContext],
) -> tuple[str, list[dict[str, Any]]]:
    """搜索本地知识库并返回与问题相关的文档片段。

    Args:
        query: 简洁、独立、适合语义检索的查询语句。
    """
    sources = await runtime.context.retriever.retrieve(query, top_k=5)
    if not sources:
        return "未检索到相关知识库片段。", []
    artifact = [serialize_source(source) for source in sources]
    content = "\n\n".join(
        f"检索片段 {index + 1}（{source.filename}，第 {source.order_index + 1} 段）：\n"
        f"{source.content}"
        for index, source in enumerate(sources)
    )
    return content, artifact


def _message_text(content: Any) -> str:
    if isinstance(content, str):
        return content
    if isinstance(content, Sequence):
        texts: list[str] = []
        for block in content:
            if isinstance(block, str):
                texts.append(block)
            elif isinstance(block, dict) and block.get("type") == "text":
                texts.append(str(block.get("text", "")))
        return "".join(texts)
    return str(content)


def create_research_tool(research_agent):
    @tool(response_format="content_and_artifact")
    async def research(
        question: str,
        runtime: ToolRuntime[ChatRunContext],
    ) -> tuple[str, list[dict[str, Any]]]:
        """调用研究 Agent 检索知识库并整理可引用证据。

        Args:
            question: 结合会话语境改写后的独立知识问题。
        """
        result = await research_agent.ainvoke(
            {"messages": [{"role": "user", "content": question}]},
            context=ResearchContext(
                retriever=runtime.context.retriever,
                request_id=runtime.context.request_id,
            ),
            config={
                "recursion_limit": 24,
                "tags": ["research-agent"],
                "metadata": {"request_id": runtime.context.request_id},
            },
        )
        by_chunk: dict[str, dict[str, Any]] = {}
        for message in result["messages"]:
            if not isinstance(message, ToolMessage) or message.name != "search_knowledge_base":
                continue
            for source in message.artifact or []:
                chunk_id = str(source["chunk_id"])
                current = by_chunk.get(chunk_id)
                if current is None or float(source["score"]) > float(current["score"]):
                    by_chunk[chunk_id] = dict(source)
        sources = sorted(by_chunk.values(), key=lambda item: float(item["score"]), reverse=True)[:8]
        writer = get_stream_writer()
        if not sources:
            writer({"kind": "no_evidence"})
            return "知识库中没有足够资料。", []
        labeled = []
        evidence_blocks: list[str] = []
        for index, source in enumerate(sources, start=1):
            labeled_source = {**source, "label": f"S{index}"}
            labeled.append(labeled_source)
            evidence_blocks.append(
                f"[S{index}] {source['filename']} 第 {int(source['order_index']) + 1} 段\n"
                f"{source['content']}"
            )
        writer({"kind": "sources", "sources": labeled})
        summary = _message_text(result["messages"][-1].content)
        tool_content = f"研究员摘要：\n{summary}\n\n可引用证据：\n" + "\n\n".join(evidence_blocks)
        return tool_content, labeled

    return research


def artifact_json(artifact: list[dict[str, Any]]) -> str:
    return json.dumps(artifact, ensure_ascii=False, separators=(",", ":"))
