"""
RAG MCP Server - HTTP Stream传输模式
基于依赖注入容器的MCP服务器实现
"""
from fastapi import FastAPI, Request
from fastapi.responses import JSONResponse, StreamingResponse
from mcp.server import Server
import uvicorn
import json
import asyncio
from typing import AsyncIterator

from src.di import get_container
from src.utils import setup_logger
from src.config import get_config_manager


# 配置日志
config = get_config_manager()
setup_logger(config.paths.log_file, config.log.level)

# 创建FastAPI应用
app = FastAPI(
    title="RAG MCP Server",
    description="知识库检索增强生成系统 - MCP服务器 (HTTP Stream)",
    version="1.1.0"
)

# 创建MCP服务器
mcp_server = Server("rag-knowledge-base")


def get_handlers():
    """从依赖注入容器获取 MCP handlers"""
    container = get_container()
    
    # 检查容器是否已注册 mcp_handlers
    if not container.has("mcp_handlers"):
        # 注册 mcp_handlers
        def create_mcp_handlers():
            from src.mcp import MCPHandlers
            rag_service = container.get("rag_service")
            return MCPHandlers(rag_service=rag_service)
        
        container.register_singleton("mcp_handlers", create_mcp_handlers)
    
    return container.get("mcp_handlers")


# 懒加载 handlers
_handlers = None

def _get_handlers():
    """获取 handlers（懒加载）"""
    global _handlers
    if _handlers is None:
        _handlers = get_handlers()
    return _handlers


@mcp_server.list_tools()
async def list_tools():
    """列出所有可用的工具"""
    return _get_handlers().tools.list_tools()


@mcp_server.call_tool()
async def call_tool(name: str, arguments: dict):
    """调用工具"""
    return _get_handlers().tools.call_tool(name, arguments)


@mcp_server.list_resources()
async def list_resources():
    """列出所有可用的资源"""
    return _get_handlers().resources.list_resources()


@mcp_server.read_resource()
async def read_resource(uri: str):
    """读取资源"""
    content = _get_handlers().resources.read_resource(uri)
    return content


@mcp_server.list_prompts()
async def list_prompts():
    """列出所有可用的提示词"""
    return _get_handlers().prompts.list_prompts()


@mcp_server.get_prompt()
async def get_prompt(name: str, arguments: dict):
    """获取提示词"""
    return _get_handlers().prompts.get_prompt(name, arguments)


# FastAPI路由
@app.get("/")
async def root():
    """根路由 - 服务器信息"""
    return {
        "name": "RAG MCP Server",
        "version": "1.1.0",
        "protocol": "MCP over HTTP Stream",
        "status": "running",
        "endpoints": {
            "mcp": "/mcp",
            "health": "/health",
            "docs": "/docs"
        },
        "tools": [
            "build_index",
            "query_knowledge",
            "retrieve_documents",
            "get_index_status",
            "get_config",
            "delete_index",
            "batch_retrieve",
            "insert_documents",
            "update_documents",
            "delete_document",
            "list_documents",
            "get_document"
        ],
        "resources": [
            "config://system",
            "stats://index",
            "stats://index/verbose"
        ]
    }


@app.get("/health")
async def health():
    """健康检查"""
    return {"status": "healthy", "server": "rag-mcp"}


async def stream_response(data: dict) -> AsyncIterator[str]:
    """流式响应生成器"""
    yield json.dumps(data) + "\n"


@app.post("/mcp")
async def handle_mcp(request: Request):
    """MCP消息处理端点"""
    handlers = _get_handlers()
    
    try:
        # 解析请求体
        body = await request.json()
        method = body.get("method")
        params = body.get("params", {})
        
        # 处理不同的MCP方法
        if method == "tools/list":
            tools = handlers.tools.list_tools()
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "result": {
                    "tools": [
                        {
                            "name": tool.name,
                            "description": tool.description,
                            "inputSchema": tool.inputSchema
                        }
                        for tool in tools
                    ]
                }
            }
            
        elif method == "tools/call":
            name = params.get("name")
            arguments = params.get("arguments", {})
            result = handlers.tools.call_tool(name, arguments)
            # 提取TextContent对象中的文本
            content = []
            for item in result:
                if hasattr(item, 'text'):
                    content.append({
                        "type": "text",
                        "text": item.text
                    })
                else:
                    content.append(item)
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "result": {
                    "content": content
                }
            }
            
        elif method == "resources/list":
            resources = handlers.resources.list_resources()
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "result": {
                    "resources": [
                        {
                            "uri": resource.uri,
                            "name": resource.name,
                            "description": resource.description,
                            "mimeType": resource.mimeType
                        }
                        for resource in resources
                    ]
                }
            }
            
        elif method == "resources/read":
            uri = params.get("uri")
            content = handlers.resources.read_resource(uri)
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "result": {
                    "contents": [
                        {
                            "uri": uri,
                            "mimeType": "application/json",
                            "text": content
                        }
                    ]
                }
            }
            
        elif method == "prompts/list":
            prompts = handlers.prompts.list_prompts()
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "result": {
                    "prompts": [
                        {
                            "name": prompt.name,
                            "description": prompt.description,
                            "arguments": prompt.arguments if hasattr(prompt, 'arguments') else []
                        }
                        for prompt in prompts
                    ]
                }
            }
            
        elif method == "prompts/get":
            name = params.get("name")
            arguments = params.get("arguments", {})
            result = handlers.prompts.get_prompt(name, arguments)
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "result": result
            }
            
        elif method == "initialize":
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "result": {
                    "protocolVersion": "2024-11-05",
                    "capabilities": {
                        "tools": {},
                        "resources": {},
                        "prompts": {}
                    },
                    "serverInfo": {
                        "name": "rag-knowledge-base",
                        "version": "1.1.0"
                    }
                }
            }
            
        else:
            response = {
                "jsonrpc": "2.0",
                "id": body.get("id"),
                "error": {
                    "code": -32601,
                    "message": f"Method not found: {method}"
                }
            }
        
        return JSONResponse(content=response)
        
    except Exception as e:
        return JSONResponse(
            status_code=500,
            content={
                "jsonrpc": "2.0",
                "id": body.get("id") if "body" in locals() else None,
                "error": {
                    "code": -32603,
                    "message": f"Internal error: {str(e)}"
                }
            }
        )


def main():
    """主函数 - 启动MCP HTTP服务器"""
    # 从配置获取主机和端口
    host = config.settings.api_host
    port = config.settings.api_port
    
    print("=" * 60)
    print("🚀 RAG MCP HTTP Stream 服务器")
    print("=" * 60)
    print(f"📍 服务地址: http://{host}:{port}")
    print(f"📡 MCP端点: http://{host}:{port}/mcp")
    print(f"📚 API文档: http://{host}:{port}/docs")
    print(f"💚 健康检查: http://{host}:{port}/health")
    print("=" * 60)
    print("\n✨ 可用工具：")
    print("  - build_index          构建知识库索引")
    print("  - query_knowledge      查询知识库（使用LLM）")
    print("  - retrieve_documents   检索文档片段（不使用LLM）")
    print("  - get_index_status     获取索引状态")
    print("  - get_config           获取系统配置")
    print("  - delete_index         删除索引")
    print("  - batch_retrieve       批量异步检索多个问题")
    print("  - insert_documents     插入新文档到索引（增量添加）")
    print("  - update_documents     更新索引中的已有文档")
    print("  - delete_document      从索引中删除指定文档")
    print("  - list_documents       列出索引中的所有文档")
    print("  - get_document         获取索引中指定文档的详细信息")
    print("\n📦 可用资源：")
    print("  - config://system           系统配置信息")
    print("  - stats://index             索引统计信息")
    print("  - stats://index/verbose     索引详细统计")
    print("\n按 Ctrl+C 停止服务器\n")
    
    uvicorn.run(
        app,
        host=host,
        port=port,
        log_level="info"
    )


if __name__ == "__main__":
    main()
