#!/usr/bin/env python3
"""
RAGFlow MCP Server - stdio版本
为本地测试和打包部署提供stdio协议支持
"""
import asyncio
from mcp.server import Server
from mcp.server.stdio import stdio_server
from mcp.types import Tool, TextContent
from dotenv import load_dotenv
from configs.ragflow import ragflow
from services.chat_assistant import ask_ragflow, create_chat_session
from services.dataset import create_initial_dataset, get_dataset_by_name
import json

load_dotenv()

# 创建MCP Server实例（stdio模式）
server = Server("ragflow-mcp")


@server.list_tools()
async def list_tools() -> list[Tool]:
    """列出可用的工具"""
    return [
        Tool(
            name="get_ragflow_datasets",
            description="获取RAGFlow数据集列表",
            inputSchema={
                "type": "object",
                "properties": {},
                "required": []
            }
        ),
        Tool(
            name="create_rag",
            description="创建初始知识库和数据集",
            inputSchema={
                "type": "object",
                "properties": {
                    "name": {
                        "type": "string",
                        "description": "要创建的数据集名称"
                    }
                },
                "required": ["name"]
            }
        ),
        Tool(
            name="upload_rag",
            description="上传文档以增强数据集的知识库",
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "要上传文档的数据集名称"
                    },
                    "display_names": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "文档的显示名称列表"
                    },
                    "blobs": {
                        "type": "array",
                        "items": {"type": "string"},
                        "description": "文档内容字符串列表"
                    }
                },
                "required": ["dataset_name", "display_names", "blobs"]
            }
        ),
        Tool(
            name="query_rag",
            description="查询知识库中指定的数据集，基于提供的查询检索答案",
            inputSchema={
                "type": "object",
                "properties": {
                    "dataset_name": {
                        "type": "string",
                        "description": "要查询的数据集名称"
                    },
                    "query": {
                        "type": "string",
                        "description": "要在数据集中搜索的问题或查询字符串"
                    }
                },
                "required": ["dataset_name", "query"]
            }
        )
    ]


@server.call_tool()
async def call_tool(name: str, arguments: dict) -> list[TextContent]:
    """调用工具"""
    try:
        if name == "get_ragflow_datasets":
            result = ragflow.list_datasets()
            return [TextContent(type="text", text=str(result))]

        elif name == "create_rag":
            dataset_name = arguments.get("name")
            existed_datasets = get_dataset_by_name(dataset_name)
            if len(existed_datasets) > 0:
                result = f"Dataset '{dataset_name}' already exists"
            else:
                response = create_initial_dataset(dataset_name)
                result = f"Successfully created dataset '{dataset_name}': {response.id}"
            return [TextContent(type="text", text=result)]

        elif name == "upload_rag":
            dataset_name = arguments.get("dataset_name")
            display_names = arguments.get("display_names", [])
            blobs = arguments.get("blobs", [])

            dataset = ragflow.get_dataset(name=dataset_name)

            documents = []
            for display_name, blob in zip(display_names, blobs):
                documents.append({
                    "display_name": display_name,
                    "blob": blob
                })

            response = dataset.upload_documents(documents)

            doc_info = []
            for doc in response:
                dataset.async_parse_documents([doc.id])
                doc_info.append({
                    "name": doc.display_name if hasattr(doc, 'display_name') else display_names[0],
                    "id": doc.id
                })

            result = json.dumps({
                "status": "success",
                "message": f"Successfully uploaded {len(documents)} documents",
                "dataset": dataset_name,
                "documents": doc_info
            }, ensure_ascii=False)
            return [TextContent(type="text", text=result)]

        elif name == "query_rag":
            dataset_name = arguments.get("dataset_name")
            query = arguments.get("query")

            response = ask_ragflow(dataset_name, query)
            result = json.dumps({
                "reference": response['data']['reference'],
                "answer": response['data']['answer']
            }, ensure_ascii=False)
            return [TextContent(type="text", text=result)]

        else:
            return [TextContent(type="text", text=f"Unknown tool: {name}")]

    except Exception as e:
        return [TextContent(type="text", text=f"Error: {str(e)}")]


async def main():
    """主函数入口"""
    async with stdio_server() as (read_stream, write_stream):
        await server.run(
            read_stream,
            write_stream,
            server.create_initialization_options()
        )


if __name__ == "__main__":
    asyncio.run(main())