from __future__ import annotations

import base64
import contextlib
import json
import signal
import ssl
import threading
import time
from urllib import error, request


@contextlib.contextmanager
def _hard_timeout(seconds: int):
    if seconds <= 0 or threading.current_thread() is not threading.main_thread() or not hasattr(signal, "SIGALRM"):
        yield
        return
    previous_handler = signal.getsignal(signal.SIGALRM)
    previous_timer = signal.getitimer(signal.ITIMER_REAL)
    started = time.monotonic()

    def _raise_timeout(_signum, _frame):
        raise TimeoutError(f"mcp_call_timeout_{seconds}s")

    signal.signal(signal.SIGALRM, _raise_timeout)
    signal.setitimer(signal.ITIMER_REAL, seconds)
    try:
        yield
    finally:
        elapsed = time.monotonic() - started
        remaining = max(0.0, previous_timer[0] - elapsed) if previous_timer[0] else 0.0
        signal.setitimer(signal.ITIMER_REAL, remaining, previous_timer[1])
        signal.signal(signal.SIGALRM, previous_handler)


class GraylogMcpClient:
    """Small Streamable HTTP MCP client used for Graylog connection checks."""

    def __init__(self, url: str, token: str, *, timeout: int = 15, verify_tls: bool = True, use_proxy: bool = False) -> None:
        self.url = url.rstrip("/")
        self.token = token
        self.timeout = timeout
        self.verify_tls = verify_tls
        self.use_proxy = use_proxy
        self.session_id: str | None = None

    def _open(self, req: request.Request, context: ssl.SSLContext | None):
        if self.use_proxy:
            return request.urlopen(req, timeout=self.timeout, context=context)
        handlers = [request.ProxyHandler({})]
        if context is not None:
            handlers.append(request.HTTPSHandler(context=context))
        return request.build_opener(*handlers).open(req, timeout=self.timeout)

    def _call(self, method: str, params: dict[str, object] | None = None, *, notification: bool = False) -> dict[str, object]:
        payload: dict[str, object] = {"jsonrpc": "2.0", "method": method}
        if not notification:
            payload["id"] = 1
        if params is not None:
            payload["params"] = params
        authorization = self._authorization_header()
        headers = {
            "Authorization": authorization,
            "Content-Type": "application/json",
            "Accept": "application/json, text/event-stream",
        }
        if self.session_id:
            headers["Mcp-Session-Id"] = self.session_id
        req = request.Request(self.url, data=json.dumps(payload).encode("utf-8"), method="POST", headers=headers)
        try:
            context = None if self.verify_tls else ssl._create_unverified_context()
            with _hard_timeout(self.timeout + 1):
                with self._open(req, context) as response:
                    self.session_id = response.headers.get("Mcp-Session-Id", self.session_id)
                    body = response.read().decode("utf-8")
        except error.HTTPError as exc:
            raise RuntimeError(f"http_{exc.code}") from exc
        except Exception as exc:
            raise RuntimeError(f"connection_error: {exc}") from exc
        if notification:
            return {}
        if body.startswith("data:"):
            body = next((line[5:].strip() for line in body.splitlines() if line.startswith("data:")), "")
        try:
            response_data = json.loads(body)
        except json.JSONDecodeError as exc:
            raise RuntimeError("invalid_mcp_response") from exc
        if "error" in response_data:
            raise RuntimeError(f"mcp_error: {response_data['error']}")
        return response_data

    def _authorization_header(self) -> str:
        """Accept a raw token or the Basic/Base64 form commonly copied from MCP clients."""
        value = self.token.strip()
        if value.lower().startswith("basic "):
            return value
        try:
            decoded = base64.b64decode(value, validate=True).decode("utf-8")
        except (ValueError, UnicodeDecodeError):
            decoded = ""
        if decoded.endswith(":token"):
            return f"Basic {value}"
        credentials = base64.b64encode(f"{value}:token".encode("utf-8")).decode("ascii")
        return f"Basic {credentials}"

    def probe(self, *, force_new_session: bool = False) -> dict[str, object]:
        if force_new_session:
            self.session_id = None
        initialized = self._call(
            "initialize",
            {"protocolVersion": "2025-06-18", "capabilities": {}, "clientInfo": {"name": "fgAI", "version": "0.1"}},
        )
        self._call("notifications/initialized", notification=True)
        tools = self._call("tools/list")
        tool_list = tools.get("result", {}).get("tools", []) if isinstance(tools.get("result"), dict) else []
        names = [str(item.get("name")) for item in tool_list if isinstance(item, dict) and item.get("name")]
        schemas = {
            str(item.get("name")): item.get("inputSchema")
            for item in tool_list
            if isinstance(item, dict) and item.get("name") and isinstance(item.get("inputSchema"), dict)
        }
        version = initialized.get("result", {}).get("serverInfo", {}).get("version", "") if isinstance(initialized.get("result"), dict) else ""
        return {"status": "connected", "tools": names, "tool_schemas": schemas, "server_version": version}

    def call_tool(self, name: str, arguments: dict[str, object]) -> dict[str, object]:
        try:
            return self._call("tools/call", {"name": name, "arguments": arguments})
        except RuntimeError as exc:
            if not self._should_reconnect(str(exc)):
                raise
            self.probe(force_new_session=True)
            return self._call("tools/call", {"name": name, "arguments": arguments})

    @staticmethod
    def _should_reconnect(error_text: str) -> bool:
        return any(token in error_text for token in ("http_400", "http_404", "http_409", "connection_error"))
