#!/usr/bin/env python3
"""
Arthas MCP Server

通用的 Arthas 代理 MCP Server，支持多实例管理。
"""

import json
import os
import sys
from typing import Any, Dict, List, Optional
import requests
from mcp.server import Server
from mcp.types import Tool, TextContent
import mcp.server.stdio

# 配置文件路径
CONFIG_FILE = os.path.join(os.path.dirname(__file__), "config.json")

class ArthasConfig:
    """Arthas 配置管理"""

    def __init__(self, config_file: str = CONFIG_FILE):
        self.config_file = config_file
        self.instances: List[Dict[str, str]] = []
        self.load_config()

    def load_config(self):
        """加载配置文件"""
        try:
            with open(self.config_file, 'r', encoding='utf-8') as f:
                config = json.load(f)
                self.instances = config.get('instances', [])

            # 替换环境变量
            for instance in self.instances:
                if 'token' in instance and instance['token'].startswith('${'):
                    env_var = instance['token'][2:-1]  # 去掉 ${ 和 }
                    instance['token'] = os.getenv(env_var, instance['token'])

        except FileNotFoundError:
            print(f"Warning: Config file not found: {self.config_file}", file=sys.stderr)
            self.instances = []
        except json.JSONDecodeError as e:
            print(f"Error: Invalid JSON in config file: {e}", file=sys.stderr)
            self.instances = []

    def get_instance(self, name: str) -> Optional[Dict[str, str]]:
        """获取指定实例配置"""
        for instance in self.instances:
            if instance.get('name') == name:
                return instance
        return None

    def list_instances(self) -> List[str]:
        """列出所有实例名称"""
        return [inst.get('name', 'unknown') for inst in self.instances]


class ArthasClient:
    """Arthas HTTP 客户端"""

    def __init__(self, base_url: str, token: str):
        self.base_url = base_url.rstrip('/')
        self.token = token
        self.session = requests.Session()
        self.session.headers.update({
            'X-Arthas-Token': token,
            'Content-Type': 'application/json'
        })

    def execute(self, command: str, operator: str = "mcp-user") -> Dict[str, Any]:
        """执行 Arthas 命令"""
        url = f"{self.base_url}/execute"
        payload = {
            "command": command,
            "operator": operator
        }

        try:
            response = self.session.post(url, json=payload, timeout=30)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {
                "success": False,
                "message": f"Request failed: {str(e)}"
            }

    def get_info(self) -> Dict[str, Any]:
        """获取应用信息"""
        url = f"{self.base_url}/info"

        try:
            response = self.session.get(url, timeout=10)
            response.raise_for_status()
            return response.json()
        except requests.exceptions.RequestException as e:
            return {
                "success": False,
                "message": f"Request failed: {str(e)}"
            }


# 初始化配置和 Server
config = ArthasConfig()
server = Server("arthas-mcp")


@server.list_tools()
async def list_tools() -> List[Tool]:
    """列出可用的工具"""
    return [
        Tool(
            name="arthas_execute",
            description="Execute Arthas command on target instance",
            inputSchema={
                "type": "object",
                "properties": {
                    "instance": {
                        "type": "string",
                        "description": "Target instance name (e.g., 'shop-frontend-daily')"
                    },
                    "command": {
                        "type": "string",
                        "description": "Arthas command to execute (e.g., 'dashboard', 'watch com.example.Service method')"
                    },
                    "operator": {
                        "type": "string",
                        "description": "Operator name for audit log",
                        "default": "mcp-user"
                    }
                },
                "required": ["instance", "command"]
            }
        ),
        Tool(
            name="arthas_info",
            description="Get application information from target instance",
            inputSchema={
                "type": "object",
                "properties": {
                    "instance": {
                        "type": "string",
                        "description": "Target instance name"
                    }
                },
                "required": ["instance"]
            }
        ),
        Tool(
            name="arthas_list_instances",
            description="List all available Arthas instances",
            inputSchema={
                "type": "object",
                "properties": {}
            }
        )
    ]


@server.call_tool()
async def call_tool(name: str, arguments: Any) -> List[TextContent]:
    """处理工具调用"""
    
    if name == "arthas_list_instances":
        instances = config.list_instances()
        if not instances:
            return [TextContent(
                type="text",
                text="No Arthas instances configured. Please check config.json"
            )]
        
        result = "Available Arthas instances:\n"
        for idx, inst_name in enumerate(instances, 1):
            result += f"{idx}. {inst_name}\n"
        
        return [TextContent(type="text", text=result)]
    
    elif name == "arthas_info":
        instance_name = arguments.get("instance")
        if not instance_name:
            return [TextContent(type="text", text="Error: instance parameter is required")]
        
        instance = config.get_instance(instance_name)
        if not instance:
            return [TextContent(
                type="text",
                text=f"Error: Instance '{instance_name}' not found. Use arthas_list_instances to see available instances."
            )]
        
        client = ArthasClient(instance['url'], instance['token'])
        result = client.get_info()
        
        return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))]
    
    elif name == "arthas_execute":
        instance_name = arguments.get("instance")
        command = arguments.get("command")
        operator = arguments.get("operator", "mcp-user")
        
        if not instance_name or not command:
            return [TextContent(type="text", text="Error: instance and command parameters are required")]
        
        instance = config.get_instance(instance_name)
        if not instance:
            return [TextContent(
                type="text",
                text=f"Error: Instance '{instance_name}' not found. Use arthas_list_instances to see available instances."
            )]
        
        client = ArthasClient(instance['url'], instance['token'])
        result = client.execute(command, operator)
        
        return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))]
    
    else:
        return [TextContent(type="text", text=f"Error: Unknown tool '{name}'")]


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


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