import json
import sys
from typing import Any, Callable

from pydantic import ValidationError

from cli._version import __version__
from common.capabilities import (
    ACTIONS_REQUIRING_EXTRA_VALUE,
    CONTROL_PLANES,
    EXECUTION_MODES,
    SUPPORTED_ACTIONS,
    SUPPORTED_PLATFORMS,
)
from common.tool_protocol import ToolRequest

JSONRPC_VERSION = "2.0"
MCP_SERVER_NAME = "screenforge"
MCP_SERVER_VERSION = __version__
MCP_TOOL_CAPABILITIES = "ui_agent_capabilities"
MCP_TOOL_EXECUTE = "ui_agent_execute"
MCP_TOOL_LOAD_RUN = "ui_agent_load_run"
MCP_TOOL_INSPECT_UI = "ui_agent_inspect_ui"
MCP_TOOL_LOAD_CASE_MEMORY = "ui_agent_load_case_memory"
SUPPORTED_MCP_PROTOCOL_VERSIONS = (
    "2025-11-25",
    "2025-06-18",
    "2025-03-26",
    "2024-11-05",
)


def _jsonrpc_response(request_id: Any, result: dict) -> dict:
    return {
        "jsonrpc": JSONRPC_VERSION,
        "id": request_id,
        "result": result,
    }


def _jsonrpc_error(request_id: Any, code: int, message: str, data: Any | None = None) -> dict:
    error = {
        "code": code,
        "message": message,
    }
    if data is not None:
        error["data"] = data
    return {
        "jsonrpc": JSONRPC_VERSION,
        "id": request_id,
        "error": error,
    }


def _negotiate_protocol_version(requested_version: str) -> str:
    requested = str(requested_version).strip()
    if requested in SUPPORTED_MCP_PROTOCOL_VERSIONS:
        return requested
    return SUPPORTED_MCP_PROTOCOL_VERSIONS[0]


def _build_capabilities_tool_schema() -> dict:
    return {
        "type": "object",
        "properties": {},
        "additionalProperties": False,
    }


def _build_execute_tool_schema() -> dict:
    return {
        "type": "object",
        "properties": {
            "mode": {
                "type": "string",
                "enum": list(EXECUTION_MODES),
                "description": "执行模式：run / doctor / plan_only / dry_run",
            },
            "platform": {
                "type": "string",
                "enum": list(SUPPORTED_PLATFORMS),
                "description": "目标平台",
            },
            "env": {
                "type": "string",
                "description": "环境名称，默认 dev",
            },
            "vision": {
                "type": "boolean",
                "description": "是否启用视觉模式",
            },
            "context": {
                "type": "string",
                "description": "上下文文件路径",
            },
            "output": {
                "type": "string",
                "description": "输出脚本路径",
            },
            "resume_run_id": {
                "type": "string",
                "description": "恢复上下文的 run_id",
            },
            "workflow": {
                "type": "object",
                "properties": {
                    "path": {
                        "type": "string",
                        "description": "workflow YAML 路径",
                    },
                    "vars": {
                        "type": "object",
                        "description": "workflow 变量覆盖",
                        "additionalProperties": {"type": "string"},
                    },
                },
                "required": ["path"],
                "additionalProperties": False,
            },
            "action": {
                "type": "object",
                "properties": {
                    "action": {
                        "type": "string",
                        "enum": list(SUPPORTED_ACTIONS),
                        "description": "即时动作类型",
                    },
                    "action_name": {
                        "type": "string",
                        "description": "即时动作名称",
                    },
                    "locator_type": {
                        "type": "string",
                        "description": "定位器类型",
                    },
                    "locator_value": {
                        "type": "string",
                        "description": "定位器值",
                    },
                    "extra_value": {
                        "type": "string",
                        "description": "动作附加值",
                    },
                },
                "required": ["action"],
                "additionalProperties": False,
            },
        },
        "additionalProperties": False,
    }


def _build_inspect_ui_tool_schema() -> dict:
    return {
        "type": "object",
        "properties": {
            "platform": {
                "type": "string",
                "enum": list(SUPPORTED_PLATFORMS),
                "description": "目标平台",
            },
            "env": {
                "type": "string",
                "description": "环境名称，默认 dev",
            },
            "vision": {
                "type": "boolean",
                "description": "是否同时返回截图 base64",
            },
        },
        "additionalProperties": False,
    }


def _build_load_case_memory_tool_schema() -> dict:
    return {
        "type": "object",
        "properties": {
            "platform": {
                "type": "string",
                "enum": list(SUPPORTED_PLATFORMS),
                "description": "可选的平台过滤条件",
            },
            "control_kind": {
                "type": "string",
                "enum": ["goal", "workflow", "action", "doctor"],
                "description": "可选的控制面过滤条件",
            },
            "source_ref": {
                "type": "string",
                "description": "可选的来源引用过滤条件",
            },
            "query": {
                "type": "string",
                "description": "按 control_label / source_ref / successful_actions 模糊查询",
            },
            "limit": {
                "type": "integer",
                "description": "最多返回条数，默认 20",
                "minimum": 1,
            },
        },
        "additionalProperties": False,
    }


def _build_load_run_tool_schema() -> dict:
    return {
        "type": "object",
        "properties": {
            "run_id": {
                "type": "string",
                "description": "需要加载的历史运行 run_id",
            }
        },
        "required": ["run_id"],
        "additionalProperties": False,
    }


def build_mcp_tools() -> list[dict]:
    return [
        {
            "name": MCP_TOOL_CAPABILITIES,
            "title": "ScreenForge Capabilities",
            "description": "Returns a capability snapshot of supported platforms, modes, control planes, and actions.",
            "inputSchema": _build_capabilities_tool_schema(),
        },
        {
            "name": MCP_TOOL_EXECUTE,
            "title": "ScreenForge Execute",
            "description": (
                "Execute or dry-run a ScreenForge request. Agent entry supports workflow, action, and doctor control planes. "
                f"Execution modes: {', '.join(EXECUTION_MODES)}. "
                f"Platforms: {', '.join(SUPPORTED_PLATFORMS)}. "
                f"Control planes: {', '.join(CONTROL_PLANES)}. "
                f"Actions requiring extra_value: {', '.join(sorted(ACTIONS_REQUIRING_EXTRA_VALUE))}."
            ),
            "inputSchema": _build_execute_tool_schema(),
        },
        {
            "name": MCP_TOOL_INSPECT_UI,
            "title": "ScreenForge Inspect UI",
            "description": "Connect to the target platform and return a cleaned XML/DOM tree for the Agent to analyze and locate elements.",
            "inputSchema": _build_inspect_ui_tool_schema(),
        },
        {
            "name": MCP_TOOL_LOAD_CASE_MEMORY,
            "title": "ScreenForge Load Case Memory",
            "description": "Load cross-run test memory for the Agent to reuse successful actions, locator strategies, and pytest assets.",
            "inputSchema": _build_load_case_memory_tool_schema(),
        },
        {
            "name": MCP_TOOL_LOAD_RUN,
            "title": "ScreenForge Load Run",
            "description": "Load a historical run by run_id, returning its summary, resume_context, and replay assets.",
            "inputSchema": _build_load_run_tool_schema(),
        },
    ]


def _build_initialize_result(protocol_version: str) -> dict:
    return {
        "protocolVersion": _negotiate_protocol_version(protocol_version),
        "capabilities": {
            "tools": {
                "listChanged": False,
            }
        },
        "serverInfo": {
            "name": MCP_SERVER_NAME,
            "version": MCP_SERVER_VERSION,
        },
        "instructions": (
            "Use ui_agent_capabilities to inspect ScreenForge support, ui_agent_inspect_ui "
            "to capture the current UI tree, ui_agent_load_case_memory to reuse existing test knowledge, "
            "and ui_agent_execute to run workflow/action/doctor requests against the execution engine."
        ),
    }


def _build_tool_result(payload: dict) -> dict:
    return {
        "content": [
            {
                "type": "text",
                "text": json.dumps(payload, ensure_ascii=False, indent=2),
            }
        ],
        "structuredContent": payload,
        "isError": not bool(payload.get("ok", False)),
    }


class McpServerSession:
    def __init__(
        self,
        execute_tool_request: Callable[[ToolRequest], dict],
        load_run_payload: Callable[[str], dict] | None = None,
        inspect_ui_payload: Callable[[ToolRequest], dict] | None = None,
        load_case_memory_payload: Callable[[ToolRequest], dict] | None = None,
    ):
        self._execute_tool_request = execute_tool_request
        self._load_run_payload = load_run_payload or (
            lambda run_id: {
                "ok": False,
                "operation": "load_run",
                "exit_code": 2,
                "error": f"未配置 load_run 处理器: {run_id}",
            }
        )
        self._inspect_ui_payload = inspect_ui_payload or (
            lambda request: {
                "ok": False,
                "operation": "inspect_ui",
                "exit_code": 2,
                "error": f"未配置 inspect_ui 处理器: {request.platform}",
            }
        )
        self._load_case_memory_payload = load_case_memory_payload or (
            lambda request: {
                "ok": False,
                "operation": "load_case_memory",
                "exit_code": 2,
                "error": "load_case_memory handler not configured",
            }
        )
        self._initialized = False

    def handle_message(self, message: dict) -> dict | None:
        if not isinstance(message, dict):
            return _jsonrpc_error(None, -32600, "Invalid JSON-RPC request")

        method = str(message.get("method", "")).strip()
        request_id = message.get("id")
        params = message.get("params", {}) or {}

        if method == "notifications/initialized":
            self._initialized = True
            return None

        if method == "initialize":
            self._initialized = True
            return _jsonrpc_response(
                request_id,
                _build_initialize_result(params.get("protocolVersion", "")),
            )

        if method == "ping":
            return _jsonrpc_response(request_id, {})

        if method == "tools/list":
            return _jsonrpc_response(request_id, {"tools": build_mcp_tools()})

        if method == "tools/call":
            return self._handle_tool_call(request_id, params)

        if request_id is None:
            return None

        return _jsonrpc_error(request_id, -32601, f"不支持的方法: {method}")

    def _handle_tool_call(self, request_id: Any, params: dict) -> dict:
        tool_name = str(params.get("name", "")).strip()
        arguments = params.get("arguments", {}) or {}

        if tool_name == MCP_TOOL_CAPABILITIES:
            request = ToolRequest(operation="capabilities")
        elif tool_name == MCP_TOOL_INSPECT_UI:
            try:
                payload = self._inspect_ui_payload(
                    ToolRequest.model_validate(
                        {
                            "operation": "inspect_ui",
                            **arguments,
                        }
                    )
                )
            except ValidationError as exc:
                return _jsonrpc_error(request_id, -32602, "tool 参数校验失败", str(exc))
            except Exception as exc:
                return _jsonrpc_error(request_id, -32603, "工具执行失败", str(exc))
            return _jsonrpc_response(request_id, _build_tool_result(payload))
        elif tool_name == MCP_TOOL_LOAD_CASE_MEMORY:
            try:
                payload = self._load_case_memory_payload(
                    ToolRequest.model_validate(
                        {
                            "operation": "load_case_memory",
                            **arguments,
                        }
                    )
                )
            except ValidationError as exc:
                return _jsonrpc_error(request_id, -32602, "tool 参数校验失败", str(exc))
            except Exception as exc:
                return _jsonrpc_error(request_id, -32603, "工具执行失败", str(exc))
            return _jsonrpc_response(request_id, _build_tool_result(payload))
        elif tool_name == MCP_TOOL_LOAD_RUN:
            run_id = str(arguments.get("run_id", "")).strip()
            if not run_id:
                return _jsonrpc_error(request_id, -32602, "tool 参数校验失败", "run_id 不能为空")
            try:
                payload = self._load_run_payload(run_id)
            except Exception as exc:
                return _jsonrpc_error(request_id, -32603, "工具执行失败", str(exc))
            return _jsonrpc_response(request_id, _build_tool_result(payload))
        elif tool_name == MCP_TOOL_EXECUTE:
            try:
                request = ToolRequest.model_validate(
                    {
                        "operation": "execute",
                        **arguments,
                    }
                )
            except ValidationError as exc:
                return _jsonrpc_error(request_id, -32602, "tool 参数校验失败", str(exc))
        else:
            return _jsonrpc_error(request_id, -32602, f"不支持的工具: {tool_name}")

        try:
            payload = self._execute_tool_request(request)
        except Exception as exc:
            return _jsonrpc_error(request_id, -32603, "工具执行失败", str(exc))

        return _jsonrpc_response(request_id, _build_tool_result(payload))


def run_stdio_mcp_server(
    execute_tool_request: Callable[[ToolRequest], dict],
    load_run_payload: Callable[[str], dict],
    inspect_ui_payload: Callable[[ToolRequest], dict],
    load_case_memory_payload: Callable[[ToolRequest], dict],
    stdin=None,
    stdout=None,
) -> int:
    input_stream = stdin or sys.stdin
    output_stream = stdout or sys.stdout
    session = McpServerSession(
        execute_tool_request=execute_tool_request,
        load_run_payload=load_run_payload,
        inspect_ui_payload=inspect_ui_payload,
        load_case_memory_payload=load_case_memory_payload,
    )

    for raw_line in input_stream:
        line = str(raw_line).strip()
        if not line:
            continue

        try:
            message = json.loads(line)
        except json.JSONDecodeError as exc:
            response = _jsonrpc_error(None, -32700, "JSON parse error", str(exc))
        else:
            if isinstance(message, list):
                response = _jsonrpc_error(None, -32600, "Batch requests are not supported")
            else:
                response = session.handle_message(message)

        if response is None:
            continue

        output_stream.write(json.dumps(response, ensure_ascii=False) + "\n")
        output_stream.flush()

    return 0
