from __future__ import annotations

import base64
import json
from urllib import error, request


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

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

    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:
            with request.urlopen(req, timeout=self.timeout) 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) -> dict[str, object]:
        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")]
        version = initialized.get("result", {}).get("serverInfo", {}).get("version", "") if isinstance(initialized.get("result"), dict) else ""
        return {"status": "connected", "tools": names, "server_version": version}

    def call_tool(self, name: str, arguments: dict[str, object]) -> dict[str, object]:
        return self._call("tools/call", {"name": name, "arguments": arguments})
