from __future__ import annotations

import asyncio
import atexit
import json
import os
import signal
import threading
from typing import Any

from .config import iter_enabled_servers

PID_FILE = os.path.expanduser("~/.mugi/mcp_pids.json")


class SessionManager:
    """Maintain long-lived MCP sessions per enabled server."""

    def __init__(self, *, register_atexit: bool = True) -> None:
        self._loop: asyncio.AbstractEventLoop | None = None
        self._thread: threading.Thread | None = None
        self._lock = threading.Lock()
        self._loop_ready = threading.Event()
        self._started = False
        self._mcp = None
        self._stdio_client = None
        self._streamable_http_client = None
        self._servers_async_lock: asyncio.Lock | None = None
        self._sessions: dict[str, Any] = {}
        self._session_configs: dict[str, dict[str, Any]] = {}
        self._runner_tasks: dict[str, asyncio.Task[Any]] = {}
        self._stop_events: dict[str, asyncio.Event] = {}
        self._init_futures: dict[str, asyncio.Future[Any]] = {}
        self._sem_dict: dict[str, asyncio.Semaphore] = {}
        if register_atexit:
            atexit.register(self.shutdown)

    def start(self) -> None:
        with self._lock:
            if self._started:
                pass
            else:
                self._cleanup_orphan_pids()
                self._started = True
                self._loop_ready.clear()
                self._thread = threading.Thread(target=self._run_loop, name="mcp_loop", daemon=True)
                self._thread.start()
        # Wait briefly so first tool call does not race loop initialization.
        self._loop_ready.wait(timeout=2.0)

    def _run_loop(self) -> None:
        self._loop = asyncio.new_event_loop()
        self._loop_ready.set()
        asyncio.set_event_loop(self._loop)
        self._servers_async_lock = asyncio.Lock()
        self._loop.run_forever()

    def _ensure_sdk(self) -> bool:
        try:
            from mcp import ClientSession, StdioServerParameters  # type: ignore
            from mcp.client.stdio import stdio_client  # type: ignore
            from mcp.client.streamable_http import streamable_http_client  # type: ignore
            self._mcp = (ClientSession, StdioServerParameters)
            self._stdio_client = stdio_client
            self._streamable_http_client = streamable_http_client
            return True
        except Exception:
            return False

    def _semaphore_for(self, server_id: str, raw: dict[str, Any]) -> asyncio.Semaphore:
        if server_id not in self._sem_dict:
            max_concurrent = int(raw.get("max_concurrent", 1 if raw.get("command") else 3))
            self._sem_dict[server_id] = asyncio.Semaphore(max(1, max_concurrent))
        return self._sem_dict[server_id]

    async def _stdio_session_runner(
        self,
        server_id: str,
        raw: dict[str, Any],
        ready: asyncio.Future[Any],
        stop_event: asyncio.Event,
    ) -> None:
        ClientSession, StdioServerParameters = self._mcp
        params = StdioServerParameters(
            command=str(raw.get("command") or ""),
            args=list(raw.get("args") or []),
            env=dict(raw.get("env") or {}),
        )
        try:
            async with self._stdio_client(params) as (read, write):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    self._sessions[server_id] = session
                    self._session_configs[server_id] = raw
                    if not ready.done():
                        ready.set_result(True)
                    await stop_event.wait()
        except Exception as exc:
            if not ready.done():
                ready.set_exception(exc)
        finally:
            self._sessions.pop(server_id, None)
            self._session_configs.pop(server_id, None)
            self._runner_tasks.pop(server_id, None)
            self._stop_events.pop(server_id, None)
            self._sem_dict.pop(server_id, None)

    async def _http_session_runner(
        self,
        server_id: str,
        raw: dict[str, Any],
        ready: asyncio.Future[Any],
        stop_event: asyncio.Event,
    ) -> None:
        ClientSession = self._mcp[0]
        url = str(raw.get("url") or "").strip()
        headers = raw.get("headers")
        http_client = None
        try:
            if isinstance(headers, dict) and headers:
                import httpx  # type: ignore

                http_client = httpx.AsyncClient(
                    headers={str(k): str(v) for k, v in headers.items()},
                    timeout=60.0,
                )
            stream_kwargs: dict[str, Any] = {}
            if http_client is not None:
                stream_kwargs["http_client"] = http_client
            async with self._streamable_http_client(url, **stream_kwargs) as (read, write, _):
                async with ClientSession(read, write) as session:
                    await session.initialize()
                    self._sessions[server_id] = session
                    self._session_configs[server_id] = raw
                    if not ready.done():
                        ready.set_result(True)
                    await stop_event.wait()
        except Exception as exc:
            if not ready.done():
                ready.set_exception(exc)
        finally:
            if http_client is not None:
                try:
                    await http_client.aclose()
                except Exception:
                    pass
            self._sessions.pop(server_id, None)
            self._session_configs.pop(server_id, None)
            self._runner_tasks.pop(server_id, None)
            self._stop_events.pop(server_id, None)
            self._sem_dict.pop(server_id, None)

    async def _ensure_session(self, server_id: str, raw: dict[str, Any]) -> None:
        if self._servers_async_lock is None:
            raise RuntimeError("MCP loop not initialized.")
        if server_id in self._sessions:
            return

        created_here = False
        async with self._servers_async_lock:
            if server_id in self._sessions:
                return
            ready = self._init_futures.get(server_id)
            if ready is None:
                ready = self._loop.create_future()
                self._init_futures[server_id] = ready
                created_here = True
                stop_event = asyncio.Event()
                self._stop_events[server_id] = stop_event
                if str(raw.get("url") or "").strip():
                    task = asyncio.create_task(
                        self._http_session_runner(server_id, raw, ready, stop_event),
                        name=f"mcp_http:{server_id}",
                    )
                else:
                    task = asyncio.create_task(
                        self._stdio_session_runner(server_id, raw, ready, stop_event),
                        name=f"mcp_stdio:{server_id}",
                    )
                self._runner_tasks[server_id] = task

        try:
            await asyncio.wait_for(ready, timeout=120.0)
        finally:
            if created_here:
                self._init_futures.pop(server_id, None)

    async def _call_tool_async(self, server_id: str, tool_name: str, arguments: dict[str, Any]) -> tuple[bool, str]:
        raw = dict(iter_enabled_servers()).get(server_id)
        if not raw:
            return False, f"Tool error: MCP server '{server_id}' is not enabled."
        if not self._ensure_sdk():
            return False, "Tool error: MCP SDK is not installed."
        try:
            await self._ensure_session(server_id, raw)
        except Exception as exc:
            return False, f"Tool error: {exc}"

        session = self._sessions.get(server_id)
        if session is None:
            return False, f"Tool error: MCP server '{server_id}' session unavailable."

        sem = self._semaphore_for(server_id, self._session_configs.get(server_id, raw))
        async with sem:
            try:
                res = await session.call_tool(tool_name, arguments=arguments)
                is_err = bool(getattr(res, "isError", False))
                content = getattr(res, "content", []) or []
                text_parts: list[str] = []
                for item in content:
                    t = getattr(item, "text", None)
                    if t:
                        text_parts.append(str(t))
                out = "\n".join(text_parts).strip()
                if not out:
                    out = str(getattr(res, "structuredContent", "") or "")
                return (not is_err), out
            except Exception as exc:
                return False, f"Tool error: {exc}"

    def call_tool(self, server_id: str, tool_name: str, arguments: dict[str, Any], timeout_seconds: float = 30.0) -> tuple[bool, str]:
        self.start()
        if self._loop is None:
            return False, "Tool error: MCP loop unavailable."
        fut = asyncio.run_coroutine_threadsafe(self._call_tool_async(server_id, tool_name, arguments), self._loop)
        try:
            return fut.result(timeout=timeout_seconds)
        except Exception as exc:
            return False, f"Tool error: {exc}"

    def clear_server_cache(self) -> None:
        self.start()
        if self._loop is None:
            self._sessions.clear()
            self._session_configs.clear()
            self._runner_tasks.clear()
            self._stop_events.clear()
            self._init_futures.clear()
            self._sem_dict.clear()
            return

        async def _clear_async() -> None:
            for stop_event in list(self._stop_events.values()):
                stop_event.set()
            for task in list(self._runner_tasks.values()):
                try:
                    await asyncio.wait_for(task, timeout=15.0)
                except Exception:
                    task.cancel()
                    try:
                        await task
                    except Exception:
                        pass
            self._sessions.clear()
            self._session_configs.clear()
            self._runner_tasks.clear()
            self._stop_events.clear()
            self._init_futures.clear()
            self._sem_dict.clear()

        fut = asyncio.run_coroutine_threadsafe(_clear_async(), self._loop)
        try:
            fut.result(timeout=30.0)
        except Exception:
            pass

    async def _discover_all_async(
        self,
        servers: list[tuple[str, dict[str, Any]]] | None = None,
    ) -> tuple[list[dict[str, Any]], list[dict[str, str]]]:
        records: list[dict[str, Any]] = []
        errors: list[dict[str, str]] = []
        if not self._ensure_sdk():
            return [], [{"server_id": "_", "error": "MCP SDK is not installed."}]
        target_servers = servers if servers is not None else iter_enabled_servers()
        for server_id, raw in target_servers:
            try:
                await self._ensure_session(server_id, raw)
                session = self._sessions.get(server_id)
                if session is None:
                    errors.append({"server_id": server_id, "error": "session unavailable"})
                    continue
                tools_resp = await session.list_tools()
                for t in getattr(tools_resp, "tools", []) or []:
                    tname = getattr(t, "name", "") or ""
                    if not tname:
                        continue
                    allow = raw.get("tool_allowlist")
                    if isinstance(allow, list) and allow and tname not in allow:
                        continue
                    schema = getattr(t, "inputSchema", None) or getattr(t, "input_schema", None) or {}
                    desc = getattr(t, "description", "") or ""
                    records.append(
                        {
                            "server_id": server_id,
                            "tool_name": tname,
                            "description": desc,
                            "input_schema": schema if isinstance(schema, dict) else {},
                        }
                    )
            except Exception as exc:
                errors.append({"server_id": server_id, "error": str(exc)})
        return records, errors

    def discover_all_sync(
        self,
        timeout_seconds: float = 120.0,
        servers: list[tuple[str, dict[str, Any]]] | None = None,
    ) -> tuple[list[dict[str, Any]], list[dict[str, str]]]:
        self.start()
        if self._loop is None:
            return [], [{"server_id": "_", "error": "MCP loop unavailable."}]
        fut = asyncio.run_coroutine_threadsafe(self._discover_all_async(servers=servers), self._loop)
        try:
            return fut.result(timeout=timeout_seconds)
        except Exception as exc:
            return [], [{"server_id": "_", "error": str(exc)}]

    def shutdown(self) -> None:
        if self._loop and self._loop.is_running():
            self.clear_server_cache()
            self._loop.call_soon_threadsafe(self._loop.stop)
        self._cleanup_pid_file()

    def _cleanup_orphan_pids(self) -> None:
        try:
            if not os.path.isfile(PID_FILE):
                return
            with open(PID_FILE, "r", encoding="utf-8") as f:
                data = json.load(f)
            if not isinstance(data, dict):
                return
            for _, pid in data.items():
                try:
                    p = int(pid)
                    os.kill(p, 0)
                    os.kill(p, signal.SIGTERM)
                except Exception:
                    pass
        except Exception:
            pass
        self._cleanup_pid_file()

    def _cleanup_pid_file(self) -> None:
        try:
            if os.path.isfile(PID_FILE):
                os.remove(PID_FILE)
        except Exception:
            pass


SESSION_MANAGER = SessionManager()
