#!/usr/bin/env python3
"""
Arthas MCP Server - 简化版

直接通过 URL 调用 Arthas，无需配置文件。
适用于内部 daily 环境快速诊断。
"""

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

# 默认 Token（可通过环境变量覆盖）
DEFAULT_TOKEN = os.getenv("ARTHAS_DEFAULT_TOKEN", "daily-arthas-token")


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

    def __init__(self, base_url: str, token: str = DEFAULT_TOKEN):
        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 = "claude") -> 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)}"
            }


# 创建 MCP Server
server = Server("arthas-mcp-server")


@server.list_tools()
async def list_tools() -> list[Tool]:
    """列出可用工具"""
    return [
        Tool(
            name="arthas_execute",
            description=(
                "在指定的 Arthas 服务上执行诊断命令。"
                "用于排查 Java 应用问题，如线程死锁、内存泄漏、方法调用追踪等。"
                "\n\n支持的命令："
                "\n- dashboard: 系统实时数据面板"
                "\n- thread: 查看线程信息，可加 -b 查看死锁"
                "\n- jvm: JVM 详细信息"
                "\n- memory: 内存使用情况"
                "\n- watch: 观测方法执行，如 watch com.example.Service method '{params, returnObj}'"
                "\n- trace: 追踪方法调用路径，如 trace com.example.Service method"
                "\n- stack: 查看方法调用堆栈"
                "\n- tt: 时间隧道，记录方法调用"
                "\n- monitor: 监控方法执行统计"
                "\n- jad: 反编译类"
                "\n- sc: 查找类"
                "\n- sm: 查找方法"
            ),
            inputSchema={
                "type": "object",
                "properties": {
                    "url": {
                        "type": "string",
                        "description": "Arthas Proxy 的完整 URL，例如: https://daily.example.com/internal/arthas"
                    },
                    "command": {
                        "type": "string",
                        "description": "要执行的 Arthas 命令，例如: dashboard, thread, jvm 等"
                    },
                    "token": {
                        "type": "string",
                        "description": f"访问 Token（可选，默认使用: {DEFAULT_TOKEN}）"
                    },
                    "operator": {
                        "type": "string",
                        "description": "操作者标识（可选，默认: claude）"
                    }
                },
                "required": ["url", "command"]
            }
        ),
        Tool(
            name="arthas_get_info",
            description="获取 Arthas 服务的应用信息，包括应用名称、PID、启动时间等。",
            inputSchema={
                "type": "object",
                "properties": {
                    "url": {
                        "type": "string",
                        "description": "Arthas Proxy 的完整 URL"
                    },
                    "token": {
                        "type": "string",
                        "description": f"访问 Token（可选，默认使用: {DEFAULT_TOKEN}）"
                    }
                },
                "required": ["url"]
            }
        )
    ]


@server.call_tool()
async def call_tool(name: str, arguments: Any) -> list[TextContent]:
    """处理工具调用"""

    if name == "arthas_execute":
        url = arguments.get("url")
        command = arguments.get("command")
        token = arguments.get("token", DEFAULT_TOKEN)
        operator = arguments.get("operator", "claude")

        if not url or not command:
            return [TextContent(
                type="text",
                text="Error: 'url' and 'command' are required parameters."
            )]

        client = ArthasClient(url, token)
        result = client.execute(command, operator)

        return [TextContent(type="text", text=json.dumps(result, indent=2, ensure_ascii=False))]

    elif name == "arthas_get_info":
        url = arguments.get("url")
        token = arguments.get("token", DEFAULT_TOKEN)

        if not url:
            return [TextContent(
                type="text",
                text="Error: 'url' is a required parameter."
            )]

        client = ArthasClient(url, token)
        result = client.get_info()

        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())
