"""A connection-only stdio MCP bridge to a user-started Artemis web service."""

from __future__ import annotations

import asyncio
from contextlib import asynccontextmanager
import json
import logging
from pathlib import Path
import sys

import httpx
from mcp import ClientSession, McpError, types
from mcp.client.streamable_http import streamable_http_client
from mcp.server.lowlevel import Server
from mcp.server.models import InitializationOptions
from mcp.server.stdio import stdio_server

from artemis.service_registry import MCP_PATH, TOKEN_HEADER, discover_service


def load_catalog():
    return json.loads(Path(__file__).with_name("catalog.json").read_text(encoding="utf-8"))


@asynccontextmanager
async def backend_session():
    info = await asyncio.to_thread(discover_service)
    async with httpx.AsyncClient(
        headers={TOKEN_HEADER: info["token"]},
        timeout=httpx.Timeout(3600, connect=5),
        trust_env=False,
    ) as client:
        async with streamable_http_client(
            f"http://127.0.0.1:{info['port']}{MCP_PATH}", http_client=client
        ) as (reader, writer, _):
            async with ClientSession(reader, writer) as session:
                initialized = await session.initialize()
                yield session, initialized


def create_bridge() -> Server:
    server = Server("artemis")
    pairs = (
        (types.ListToolsRequest, types.ListToolsResult),
        (types.CallToolRequest, types.CallToolResult),
        (types.ListResourcesRequest, types.ListResourcesResult),
        (types.ListResourceTemplatesRequest, types.ListResourceTemplatesResult),
        (types.ReadResourceRequest, types.ReadResourceResult),
        (types.ListPromptsRequest, types.ListPromptsResult),
        (types.GetPromptRequest, types.GetPromptResult),
    )

    def forward(result_type):
        async def handler(request):
            # A fresh HTTP session per request makes the next request reconnect
            # after a backend restart. Never replay a possibly applied tool call.
            try:
                async with backend_session() as (session, _):
                    # Incoming SDK requests retain the outer JSON-RPC id/jsonrpc
                    # as extras. The upstream session must allocate its own id.
                    forwarded = type(request).model_validate(
                        request.model_dump(include={"method", "params"}, exclude_none=True)
                    )
                    result = await session.send_request(types.ClientRequest(forwarded), result_type)
                    return types.ServerResult(result)
            except (OSError, RuntimeError, httpx.HTTPError, ExceptionGroup, McpError) as exc:
                logging.getLogger(__name__).debug("Shared MCP request failed", exc_info=True)
                if isinstance(request, types.ListToolsRequest):
                    return types.ServerResult(types.ListToolsResult(tools=load_catalog()["tools"]))
                if isinstance(request, types.CallToolRequest):
                    return types.ServerResult(
                        types.CallToolResult(
                            isError=True,
                            content=[
                                types.TextContent(
                                    type="text",
                                    text=(
                                        "ARTEMIS 主程序未启动、版本尚未重载或连接已断开。"
                                        "请由用户运行 start.bat / start-cn.bat 或 artemis ui。"
                                        "MCP 不会自动启动主程序。本次调用不会自动重试；"
                                        "若已提交任务，请先查询状态，避免重复执行。"
                                    ),
                                )
                            ],
                        )
                    )
                raise RuntimeError(
                    "ARTEMIS 主程序连接失败或已断开。请确认网页服务正在运行。"
                    "本次调用不会自动重试；若提交过任务，请先查询任务状态，避免重复执行。"
                ) from exc

        return handler

    for request_type, result_type in pairs:
        server.request_handlers[request_type] = forward(result_type)
    return server


async def run() -> None:
    from artemis import __version__

    # An offline backend must never prevent the host agent session from opening.
    # The shipped catalog is schema-checked against the real tools by CI.
    options = InitializationOptions(
        server_name="artemis",
        server_version=__version__,
        capabilities=types.ServerCapabilities(tools=types.ToolsCapability(listChanged=False)),
        instructions=load_catalog()["instructions"],
    )
    server = create_bridge()
    async with stdio_server() as (reader, writer):
        await server.run(reader, writer, options)


def main() -> None:
    try:
        asyncio.run(run())
    except KeyboardInterrupt:
        return
    except (OSError, ValueError, RuntimeError, httpx.HTTPError, ExceptionGroup, McpError) as exc:
        # stdout is exclusively MCP JSON-RPC. No agent/device runtime fallback.
        print(f"ARTEMIS MCP 无法连接：{exc}", file=sys.stderr)
        raise SystemExit(1) from exc


if __name__ == "__main__":
    main()