#!/usr/bin/env python3
"""Inference Cooperative docs MCP server (stdio, stdlib-only).

Exposes a `search_docs` tool over the co-op documentation (co-op/docs on
git.inference.coop). Members opt in per conversation via the tools palette.

Design notes:
- stdio transport: newline-delimited JSON-RPC 2.0 on stdin/stdout (the MCP
  stdio protocol). All logging goes to stderr — stdout is protocol-only.
- No third-party dependencies: the corpus is fetched from the public Gitea
  raw API, chunked by heading, and ranked with a compact BM25.
- The index refreshes lazily if older than DOCS_TTL seconds (default 30 min).
"""

import json
import math
import re
import sys
import time
import urllib.request
from concurrent.futures import ThreadPoolExecutor

DOCS_REPO = "co-op/docs"
GITEA_API = "https://git.inference.coop/api/v1"
DOCS_TTL = 30 * 60
MAX_CHUNKS_RETURNED = 5

SERVER_NAME = "inference-cooperative-docs"
SERVER_VERSION = "1.0.0"

_index = {"chunks": [], "built_at": 0.0}
_lock_held = False

TOKEN_RE = re.compile(r"[a-z0-9]+")


def log(msg):
    sys.stderr.write(f"[inference-docs-mcp] {msg}\n")
    sys.stderr.flush()


def fetch_json(url):
    with urllib.request.urlopen(url, timeout=30) as r:
        return json.loads(r.read().decode("utf-8"))


def fetch_text(url):
    with urllib.request.urlopen(url, timeout=30) as r:
        return r.read().decode("utf-8")


def list_doc_files():
    files = fetch_json(f"{GITEA_API}/repos/{DOCS_REPO}/contents")
    return [
        {"path": f["path"], "name": f["name"]}
        for f in files
        if f["type"] == "file" and f["name"].endswith(".md")
    ]


def chunk_markdown(path, text):
    """Split a doc into section chunks keyed by its top title + heading."""
    lines = text.splitlines()
    title = lines[0].lstrip("# ").strip() if lines and lines[0].startswith("#") else path
    sections = []
    current_head = "Overview"
    current = []

    def flush():
        body = "\n".join(current).strip()
        if len(body) < 40:
            return
        sections.append({"doc": title, "path": path, "section": current_head, "text": body})

    for line in lines:
        if line.startswith("## "):
            flush()
            current_head = line.lstrip("# ").strip()
            current = []
        else:
            current.append(line)
    flush()
    return sections


def tokenize(text):
    return TOKEN_RE.findall(text.lower())


def build_index():
    chunks = []
    docs = list_doc_files()

    def load(doc):
        try:
            raw = fetch_text(f"{GITEA_API}/repos/{DOCS_REPO}/raw/{doc['path']}")
            return chunk_markdown(doc["path"], raw)
        except Exception as exc:  # noqa: BLE001 - one bad doc shouldn't kill the index
            log(f"skipping {doc['path']}: {exc}")
            return []

    with ThreadPoolExecutor(max_workers=6) as pool:
        for parts in pool.map(load, docs):
            chunks.extend(parts)

    # BM25 statistics
    doc_terms = [tokenize(c["text"]) for c in chunks]
    n = len(chunks)
    avg_len = sum(len(t) for t in doc_terms) / max(n, 1)
    df = {}
    for terms in doc_terms:
        for term in set(terms):
            df[term] = df.get(term, 0) + 1
    for c, terms in zip(chunks, doc_terms):
        c["_tf"] = {}
        for t in terms:
            c["_tf"][t] = c["_tf"].get(t, 0) + 1
        c["_len"] = len(terms)
    _index["chunks"] = chunks
    _index["df"] = df
    _index["n"] = n
    _index["avg_len"] = avg_len
    _index["built_at"] = time.time()
    log(f"index built: {n} chunks from {len(docs)} docs")


def ensure_index():
    if not _index["chunks"] or time.time() - _index["built_at"] > DOCS_TTL:
        build_index()


def search(query, limit=MAX_CHUNKS_RETURNED):
    ensure_index()
    k1, b = 1.5, 0.75
    q_terms = tokenize(query)
    scored = []
    for c in _index["chunks"]:
        score = 0.0
        for t in q_terms:
            tf = c["_tf"].get(t, 0)
            if not tf:
                continue
            idf = math.log(1 + (_index["n"] - _index["df"].get(t, 0) + 0.5) / (_index["df"].get(t, 0) + 0.5))
            score += idf * tf * (k1 + 1) / (tf + k1 * (1 - b + b * c["_len"] / _index["avg_len"]))
        if score > 0:
            scored.append((score, c))
    scored.sort(key=lambda pair: -pair[0])
    return scored[:limit]


def tool_search_docs(args):
    query = str(args.get("query") or "").strip()
    if not query:
        return {"content": [{"type": "text", "text": "Provide a query."}]}
    limit = min(int(args.get("limit") or MAX_CHUNKS_RETURNED), 10)
    results = search(query, limit)
    if not results:
        return {"content": [{"type": "text", "text": "No matching documentation found."}]}
    parts = []
    for score, c in results:
        parts.append(f"[{c['doc']} — {c['section']}] (relevance {score:.2f})\n{c['text']}")
    return {"content": [{"type": "text", "text": "\n\n---\n\n".join(parts)}]}


TOOLS = [
    {
        "name": "search_docs",
        "description": (
            "Search the Inference Cooperative documentation (getting started, models, "
            "API access, charter, policies). Returns the most relevant doc sections."
        ),
        "inputSchema": {
            "type": "object",
            "properties": {
                "query": {"type": "string", "description": "What to look for in the docs"},
                "limit": {"type": "integer", "description": "Max sections to return (default 5)"},
            },
            "required": ["query"],
        },
    }
]

TOOLS_HANDLERS = {"search_docs": tool_search_docs}


def handle(msg):
    method = msg.get("method")
    msg_id = msg.get("id")

    if method == "initialize":
        return {
            "jsonrpc": "2.0",
            "id": msg_id,
            "result": {
                "protocolVersion": "2024-11-05",
                "capabilities": {"tools": {}},
                "serverInfo": {"name": SERVER_NAME, "version": SERVER_VERSION},
            },
        }
    if method == "ping":
        return {"jsonrpc": "2.0", "id": msg_id, "result": {}}
    if method == "tools/list":
        return {"jsonrpc": "2.0", "id": msg_id, "result": {"tools": TOOLS}}
    if method == "tools/call":
        params = msg.get("params") or {}
        name = params.get("name")
        handler = TOOLS_HANDLERS.get(name)
        if not handler:
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "error": {"code": -32602, "message": f"Unknown tool: {name}"},
            }
        try:
            result = handler(params.get("arguments") or {})
            return {"jsonrpc": "2.0", "id": msg_id, "result": result}
        except Exception as exc:  # noqa: BLE001 - report tool errors in-band
            log(f"tool error: {exc}")
            return {
                "jsonrpc": "2.0",
                "id": msg_id,
                "result": {"content": [{"type": "text", "text": f"Search failed: {exc}"}], "isError": True},
            }
    if method and method.startswith("notifications/"):
        return None
    if msg_id is not None:
        return {"jsonrpc": "2.0", "id": msg_id, "error": {"code": -32601, "message": "Method not found"}}
    return None


def main():
    log("starting (stdio)")
    for line in sys.stdin:
        line = line.strip()
        if not line:
            continue
        try:
            msg = json.loads(line)
        except json.JSONDecodeError as exc:
            log(f"bad json: {exc}")
            continue
        try:
            resp = handle(msg)
        except Exception as exc:  # noqa: BLE001 - never crash the protocol loop
            log(f"handler crash: {exc}")
            if msg.get("id") is not None:
                resp = {"jsonrpc": "2.0", "id": msg["id"], "error": {"code": -32603, "message": "Internal error"}}
            else:
                resp = None
        if resp is not None:
            sys.stdout.write(json.dumps(resp) + "\n")
            sys.stdout.flush()


if __name__ == "__main__":
    main()