#!/usr/bin/env python3
"""
RAG检索MCP服务器
调用本地RAG API进行文档检索
"""

import asyncio
import hashlib
import json
import sys
import os
from typing import Any, Dict, List

try:
    import requests
except ImportError:
    print("Error: requests module is required. Please install it with: pip install requests")
    sys.exit(1)

from mcp_common import (
    create_error_response,
    create_success_response,
    create_initialize_response,
    create_ping_response,
    create_tools_list_response,
    load_tools_from_json,
    handle_mcp_streaming
)
BACKEND_HOST = os.getenv("BACKEND_HOST", "https://api-dev.gptbase.ai")
MASTERKEY = os.getenv("MASTERKEY", "master")

def rag_retrieve(query: str, top_k: int = 100) -> Dict[str, Any]:
    """调用RAG检索API"""
    try:
        bot_id = ""
        if len(sys.argv) > 1:
            bot_id = sys.argv[1]

        url = f"{BACKEND_HOST}/v1/rag_retrieve/{bot_id}"
        if not url:
            return {
                "content": [
                    {
                        "type": "text",
                        "text": "Error: RAG API URL not provided. Please provide URL as command line argument."
                    }
                ]
            }

        # 获取masterkey并生成认证token
        masterkey = MASTERKEY
        token_input = f"{masterkey}:{bot_id}"
        auth_token = hashlib.md5(token_input.encode()).hexdigest()

        headers = {
            "content-type": "application/json",
            "authorization": f"Bearer {auth_token}"
        }
        data = {
            "query": query,
            "top_k": top_k
        }
        
        # 发送POST请求
        response = requests.post(url, json=data, headers=headers, timeout=30)
        
        if response.status_code != 200:
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"Error: RAG API returned status code {response.status_code}. Response: {response.text}"
                    }
                ]
            }
        
        # 解析响应
        try:
            response_data = response.json()
        except json.JSONDecodeError as e:
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"Error: Failed to parse API response as JSON. Error: {str(e)}, Raw response: {response.text}"
                    }
                ]
            }
        
        # 提取markdown字段
        if "markdown" in response_data:
            markdown_content = response_data["markdown"]
            return {
                "content": [
                    {
                        "type": "text",
                        "text": markdown_content
                    }
                ]
            }
        else:
            return {
                "content": [
                    {
                        "type": "text",
                        "text": f"Error: 'markdown' field not found in API response. Response: {json.dumps(response_data, indent=2, ensure_ascii=False)}"
                    }
                ]
            }
            
    except requests.exceptions.RequestException as e:
        return {
            "content": [
                {
                    "type": "text",
                    "text": f"Error: Failed to connect to RAG API. {str(e)}"
                }
            ]
        }
    except Exception as e:
        return {
            "content": [
                {
                    "type": "text",
                    "text": f"Error: {str(e)}"
                }
            ]
        }


async def handle_request(request: Dict[str, Any]) -> Dict[str, Any]:
    """Handle MCP request"""
    try:
        method = request.get("method")
        params = request.get("params", {})
        request_id = request.get("id")
        
        if method == "initialize":
            return create_initialize_response(request_id, "rag-retrieve")
        
        elif method == "ping":
            return create_ping_response(request_id)
        
        elif method == "tools/list":
            # 从 JSON 文件加载工具定义
            tools = load_tools_from_json("rag_retrieve_tools.json")
            if not tools:
                # 如果 JSON 文件不存在，使用默认定义
                tools = [
                    {
                        "name": "rag_retrieve",
                        "description": "调用RAG检索API，根据查询内容检索相关文档。返回包含相关内容的markdown格式结果。",
                        "inputSchema": {
                            "type": "object",
                            "properties": {
                                "query": {
                                    "type": "string",
                                    "description": "检索查询内容"
                                }
                            },
                            "required": ["query"]
                        }
                    }
                ]
            return create_tools_list_response(request_id, tools)
        
        elif method == "tools/call":
            tool_name = params.get("name")
            arguments = params.get("arguments", {})
            
            if tool_name == "rag_retrieve":
                query = arguments.get("query", "")
                top_k = arguments.get("top_k", 100)

                if not query:
                    return create_error_response(request_id, -32602, "Missing required parameter: query")

                result = rag_retrieve(query, top_k)

                return {
                    "jsonrpc": "2.0",
                    "id": request_id,
                    "result": result
                }
            
            else:
                return create_error_response(request_id, -32601, f"Unknown tool: {tool_name}")
        
        else:
            return create_error_response(request_id, -32601, f"Unknown method: {method}")
    
    except Exception as e:
        return create_error_response(request.get("id"), -32603, f"Internal error: {str(e)}")


async def main():
    """Main entry point."""
    await handle_mcp_streaming(handle_request)


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