# requires: mcp, pydantic
"""一个内部知识库 / 工单的 MCP 服务：tools（搜索、创建工单）、resources（文档）、prompts（模板）。

基于 MCP Python SDK 2.x 编写（1.x 中的 FastMCP 在 2.x 改名为 MCPServer），以官方文档为准：
https://py.sdk.modelcontextprotocol.io/
运行方式：
  python kb_server.py           # stdio 方式启动，供 Cursor / Claude Desktop 等客户端接入
  python test_kb_server.py      # 进程内客户端测试（不需要启动服务）
"""
from __future__ import annotations

import os

from pydantic import BaseModel, Field

from mcp.server.mcpserver import MCPServer

DOCS = {
    "vpn": "VPN 连接失败：先确认账号未过期，再重新安装客户端；仍失败请提交 IT 工单。",
    "expense": "差旅报销：一线城市住宿标准每晚 600 元，需在出差结束后 30 天内提交。",
}
TICKETS: dict[str, dict] = {}

mcp = MCPServer("internal-kb", instructions="公司内部知识库和 IT 工单工具。只读工具可以直接调用，创建工单前先确认问题描述。")


class SearchHit(BaseModel):
    doc_id: str
    snippet: str


@mcp.tool()
def search_docs(query: str, top_k: int = 3) -> list[SearchHit]:
    """在内部知识库中按关键词搜索文档，返回文档 ID 和摘要。用于回答公司制度和 IT 问题。"""
    hits = [SearchHit(doc_id=k, snippet=v[:40]) for k, v in DOCS.items()
            if any(w and w in v for w in query.split())]
    return hits[:top_k]


class TicketIn(BaseModel):
    title: str = Field(min_length=4, max_length=80)
    category: str = Field(pattern="^(it|hr|finance)$")
    request_id: str = Field(description="客户端生成的唯一 ID，用于幂等")


@mcp.tool()
def create_ticket(ticket: TicketIn) -> dict:
    """创建一张工单。会产生副作用：同一个 request_id 只创建一次。"""
    if ticket.request_id not in TICKETS:
        TICKETS[ticket.request_id] = {"id": f"T{len(TICKETS) + 1:04d}", **ticket.model_dump()}
    return TICKETS[ticket.request_id]


@mcp.resource("kb://docs/{doc_id}")
def read_doc(doc_id: str) -> str:
    """读取一篇知识库文档的全文。"""
    return DOCS.get(doc_id, "文档不存在")


@mcp.prompt()
def answer_with_citation(question: str) -> str:
    """基于知识库回答问题的提示词模板。"""
    return f"请先调用 search_docs 检索，再根据检索结果回答，并标注 doc_id。问题：{question}"


if __name__ == "__main__" and os.environ.get("RUN_EXAMPLES") != "1":
    mcp.run()          # 默认 stdio；也可以 mcp.run("streamable-http")