"""MCP tools hosted once inside the user-started web service."""

from __future__ import annotations

import asyncio
from functools import wraps
import hmac
import inspect
import ipaddress
import logging
import threading
from typing import get_type_hints

from starlette.responses import JSONResponse

from artemis.mcp.tool_server import ArtemisFastMCP
from artemis.service_registry import TOKEN_HEADER


class ToolExecutor:
    """One shared tool loop; synchronous legacy HTTP calls cannot block the UI.

    Device clients keep their loop affinity across MCP connections. Cancellation
    propagates to async tools; dispatched synchronous writes may finish in their
    worker thread, so callers must resolve the original request rather than retry it.
    """

    def __init__(self):
        self.loop = None
        self.thread = None

    def start(self):
        ready = threading.Event()

        def run():
            self.loop = asyncio.new_event_loop()
            asyncio.set_event_loop(self.loop)
            self.loop.call_soon(ready.set)
            self.loop.run_forever()
            self.loop.close()

        self.thread = threading.Thread(target=run, name="artemis-shared-mcp", daemon=True)
        self.thread.start()
        ready.wait()

    async def invoke(self, fn, *args, **kwargs):
        if self.loop is None or not self.loop.is_running():
            raise RuntimeError("ARTEMIS 主程序正在启动或关闭，请稍后重试。")

        async def dispatch():
            if inspect.iscoroutinefunction(fn):
                return await fn(*args, **kwargs)
            # Legacy tools use synchronous HTTP to call this same web service.
            # Execute them off both event loops, and never call them speculatively
            # to discover whether their return value is awaitable.
            result = await asyncio.to_thread(fn, *args, **kwargs)
            return await result if inspect.isawaitable(result) else result

        future = asyncio.run_coroutine_threadsafe(dispatch(), self.loop)
        return await asyncio.wrap_future(future)

    async def stop(self):
        if self.loop is None or not self.loop.is_running():
            return

        async def cancel_pending():
            pending = [task for task in asyncio.all_tasks() if task is not asyncio.current_task()]
            for task in pending:
                task.cancel()
            if pending:
                await asyncio.gather(*pending, return_exceptions=True)

        future = asyncio.run_coroutine_threadsafe(cancel_pending(), self.loop)
        try:
            await asyncio.wait_for(asyncio.wrap_future(future), timeout=5)
        except TimeoutError:
            logging.getLogger(__name__).warning("Timed out draining shared MCP tools on shutdown")
        finally:
            self.loop.call_soon_threadsafe(self.loop.stop)
            await asyncio.to_thread(self.thread.join, 2)


def create_shared_mcp():
    from mcp_server.base import mcp as definitions
    from mcp_server import tools

    executor = ToolExecutor()
    shared = ArtemisFastMCP(
        "artemis",
        strict_inputs=True,
        instructions=definitions.instructions,
        streamable_http_path="/",
        json_response=True,
    )

    def wrap(fn):
        @wraps(fn)
        async def call(*args, **kwargs):
            return await executor.invoke(fn, *args, **kwargs)

        call.__annotations__ = get_type_hints(fn)
        call.__signature__ = inspect.signature(fn, eval_str=True)
        return call

    for name in tools.__all__:
        shared.add_tool(wrap(getattr(tools, name)))
    return shared, executor


class LocalMcpBoundary:
    """Only authenticated loopback bridges may enter the shared tool runtime."""

    def __init__(self, app, token: str):
        self.app = app
        self.token = token

    async def __call__(self, scope, receive, send):
        if scope["type"] != "http":
            await self.app(scope, receive, send)
            return
        headers = dict(scope.get("headers", []))
        supplied = headers.get(TOKEN_HEADER.lower().encode(), b"").decode("latin-1")
        try:
            local = ipaddress.ip_address((scope.get("client") or ("",))[0]).is_loopback
        except ValueError:
            local = False
        if not local or not hmac.compare_digest(supplied, self.token):
            await JSONResponse({"detail": "MCP 仅接受本机 ARTEMIS 连接。"}, status_code=403)(
                scope, receive, send
            )
            return
        await self.app(scope, receive, send)