#!/usr/bin/env python3
"""Wikipedia MCP server (JSON-RPC 2.0 over stdio) using the free Wikipedia API.

Capabilities (callable as `mcp__wikipedia__<tool>`):
  status       — self-check: confirms Wikipedia API is reachable
  search       — search Wikipedia articles by keyword
  readArticle  — get the full text content of a Wikipedia article

Data source (free, no API key required):
  - Wikipedia API: https://en.wikipedia.org/w/api.php

Run with:
  python3 wikipedia_mcp_server.py
"""

from __future__ import annotations

import json
import sys
from typing import Any

import httpx

# ── Logging ─────────────────────────────────────────────────────────────────────

def log(msg: str) -> None:
    print(f"[wikipedia_mcp] {msg}", file=sys.stderr, flush=True)


# ── Constants ───────────────────────────────────────────────────────────────────

WIKI_API = "https://en.wikipedia.org/w/api.php"
USER_AGENT = "SkaldWikipediaMCP/1.0"
HTTP_TIMEOUT = 15.0


# ── Tool implementations ────────────────────────────────────────────────────────

def _status(args: dict[str, Any]) -> str:
    """Self-check: confirm Wikipedia API is reachable."""
    try:
        params = {
            "action": "query",
            "meta": "siteinfo",
            "format": "json",
        }
        with httpx.Client(timeout=HTTP_TIMEOUT) as client:
            resp = client.get(WIKI_API, params=params,
                              headers={"User-Agent": USER_AGENT})
            resp.raise_for_status()
            data = resp.json()
        if "query" in data:
            return ("OK: Wikipedia API is reachable. "
                    "All tools (search, readArticle) are operational.")
        return "Error: Wikipedia API responded but returned unexpected data."
    except httpx.TimeoutException:
        return "Error: Wikipedia API request timed out."
    except httpx.HTTPStatusError as e:
        return f"Error: Wikipedia API returned HTTP {e.response.status_code}."
    except httpx.HTTPError as e:
        return f"Error: Wikipedia API request failed: {e}"


def _search(args: dict[str, Any]) -> str:
    """Search Wikipedia articles by keyword, returning titles and summaries."""
    query = args.get("query", "").strip()
    if not query:
        return "Error: Missing required parameter 'query'."

    limit = args.get("limit", 8)
    try:
        limit = int(limit)
    except (TypeError, ValueError):
        limit = 8
    if limit < 1:
        limit = 1
    if limit > 50:
        limit = 50

    try:
        params = {
            "action": "opensearch",
            "search": query,
            "limit": limit,
            "namespace": 0,
            "format": "json",
        }
        with httpx.Client(timeout=HTTP_TIMEOUT) as client:
            resp = client.get(WIKI_API, params=params,
                              headers={"User-Agent": USER_AGENT})
            resp.raise_for_status()
            data = resp.json()
    except httpx.TimeoutException:
        return "Error: Wikipedia search request timed out."
    except httpx.HTTPStatusError as e:
        return f"Error: Wikipedia API returned HTTP {e.response.status_code}."
    except httpx.HTTPError as e:
        return f"Error: Wikipedia API request failed: {e}"

    # opensearch returns [query, [titles], [descriptions], [urls]]
    titles = data[1] if len(data) > 1 else []
    descriptions = data[2] if len(data) > 2 else []
    urls = data[3] if len(data) > 3 else []

    if not titles:
        return f"No results found for '{query}'."

    lines = [f"📚 Wikipedia search results for '{query}'"]
    lines.append("")
    for i, title in enumerate(titles):
        desc = descriptions[i] if i < len(descriptions) else ""
        url = urls[i] if i < len(urls) else ""
        lines.append(f"  {i+1}. **{title}**")
        if desc:
            lines.append(f"     {desc}")
        if url:
            lines.append(f"     {url}")
        lines.append("")

    return "\n".join(lines)


def _read_article(args: dict[str, Any]) -> str:
    """Get the full text content of a Wikipedia article by title."""
    title = args.get("title", "").strip()
    if not title:
        return "Error: Missing required parameter 'title'."

    try:
        # First, get the page ID / exact title
        params = {
            "action": "query",
            "titles": title,
            "redirects": 1,
            "format": "json",
        }
        with httpx.Client(timeout=HTTP_TIMEOUT) as client:
            resp = client.get(WIKI_API, params=params,
                              headers={"User-Agent": USER_AGENT})
            resp.raise_for_status()
            data = resp.json()

        pages = data.get("query", {}).get("pages", {})
        if not pages:
            return f"Error: Article '{title}' not found."

        page_id = next(iter(pages))
        page_info = pages[page_id]

        if "missing" in page_info:
            return f"Error: Article '{title}' not found on Wikipedia."

        actual_title = page_info.get("title", title)

        # Now get the extract (full text)
        params = {
            "action": "query",
            "pageids": page_id,
            "prop": "extracts",
            "explaintext": 1,
            "exlimit": 1,
            "format": "json",
        }
        with httpx.Client(timeout=HTTP_TIMEOUT) as client:
            resp = client.get(WIKI_API, params=params,
                              headers={"User-Agent": USER_AGENT})
            resp.raise_for_status()
            data = resp.json()

        pages = data.get("query", {}).get("pages", {})
        extract = pages.get(page_id, {}).get("extract", "")

        if not extract:
            return f"**{actual_title}**\n\n_(No content available.)_"

        # Get the URL
        article_url = f"https://en.wikipedia.org/wiki/{actual_title.replace(' ', '_')}"

        # Truncate very long articles to avoid blowing up context
        max_chars = 15000
        if len(extract) > max_chars:
            extract = extract[:max_chars] + "\n\n_...(article truncated, see full version on Wikipedia)_"

        return (
            f"# {actual_title}\n\n"
            f"_{article_url}_\n\n"
            f"{extract}"
        )

    except httpx.TimeoutException:
        return "Error: Wikipedia API request timed out."
    except httpx.HTTPStatusError as e:
        return f"Error: Wikipedia API returned HTTP {e.response.status_code}."
    except httpx.HTTPError as e:
        return f"Error: Wikipedia API request failed: {e}"


# ── Tool manifest ────────────────────────────────────────────────────────────────

TOOLS = [
    {
        "name": "status",
        "title": "Status",
        "description": (
            "Self-check that the Wikipedia integration is operational: confirms "
            "the Wikipedia API is reachable by issuing a cheap siteinfo query. "
            "Call this first if any Wikipedia tool fails."
        ),
        "inputSchema": {"type": "object", "properties": {}},
    },
    {
        "name": "search",
        "title": "Search Wikipedia",
        "description": (
            "Search Wikipedia articles by keyword. Returns a list of matching "
            "article titles with short descriptions and URLs. Use this to find "
            "the correct article title before calling readArticle."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {
                    "type": "string",
                    "description": "Search query (keywords or phrase).",
                },
                "limit": {
                    "type": "integer",
                    "description": "Maximum number of results to return (1-50). Default: 8.",
                },
            },
            "required": ["query"],
        },
    },
    {
        "name": "readArticle",
        "description": (
            "Get the full text content of a Wikipedia article by title. "
            "Returns the article body as Markdown-like text, with a link to the "
            "live Wikipedia page. Long articles are truncated at 15,000 characters. "
            "First use search() to find the exact article title."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "title": {
                    "type": "string",
                    "description": "Exact Wikipedia article title (e.g. 'Alan Turing', 'Rome').",
                },
            },
            "required": ["title"],
        },
    },
]

TOOL_DISPATCH = {
    "status":      _status,
    "search":      _search,
    "readArticle": _read_article,
}


# ── JSON-RPC dispatch ────────────────────────────────────────────────────────────

def _ok(req_id: Any, result: Any) -> str:
    return json.dumps({"jsonrpc": "2.0", "id": req_id, "result": result})


def _text_result(req_id: Any, text: str, is_error: bool = False) -> str:
    payload: dict = {
        "jsonrpc": "2.0",
        "id": req_id,
        "result": {"content": [{"type": "text", "text": text}]},
    }
    if is_error:
        payload["result"]["isError"] = True
    return json.dumps(payload)


def _error(req_id: Any, code: int, message: str) -> str:
    return json.dumps({
        "jsonrpc": "2.0",
        "id": req_id,
        "error": {"code": code, "message": message},
    })


def handle_request(msg: dict) -> str | None:
    method = msg.get("method", "")
    req_id = msg.get("id")

    if method == "initialize":
        return _ok(req_id, {
            "protocolVersion": "2024-11-05",
            "capabilities": {"tools": {}},
            "serverInfo": {
                "name": "wikipedia",
                "version": "1.0.0",
            },
        })

    if method == "notifications/initialized":
        return None

    if method == "ping":
        return _ok(req_id, {})

    if method == "tools/list":
        return _ok(req_id, {"tools": TOOLS})

    if method == "tools/call":
        params = msg.get("params", {})
        tool_name = params.get("name", "")
        tool_args = params.get("arguments", {})

        handler = TOOL_DISPATCH.get(tool_name)
        if handler is None:
            return _text_result(req_id, f"Error: Unknown tool: {tool_name}", is_error=True)

        try:
            result = handler(tool_args)
            is_err = result.startswith("Error:")
            return _text_result(req_id, result, is_error=is_err)
        except Exception as e:
            log(f"Unhandled exception in tool '{tool_name}': {e}")
            return _text_result(req_id, f"Error: Internal error in tool '{tool_name}': {e}", is_error=True)

    return _error(req_id, -32601, f"Method not found: {method}")


# ── Main loop ────────────────────────────────────────────────────────────────────

def main() -> None:
    log("Starting Wikipedia MCP server")
    try:
        for line in sys.stdin:
            line = line.strip()
            if not line:
                continue
            try:
                msg = json.loads(line)
            except json.JSONDecodeError as e:
                log(f"Invalid JSON input: {e}")
                continue

            resp = handle_request(msg)
            if resp is not None:
                sys.stdout.write(resp + "\n")
                sys.stdout.flush()
    except KeyboardInterrupt:
        pass


if __name__ == "__main__":
    main()