"""MCP server over the FastAPI docs corpus: BM25 search + fetch-by-id tools.
Run: python docs_server.py            (stdio)
     python docs_server.py http 8765  (streamable HTTP at /mcp)
"""
import glob, math, os, re, sys, collections
from mcp.server.mcpserver import MCPServer
from mcp.server.mcpserver.exceptions import ToolError

CORPUS = os.path.join(os.path.dirname(__file__), "..", "corpus")
TOK = re.compile(r"[a-z0-9_]+")
HEADING = re.compile(r"^(#{1,3}) (.+?)\s*(\{.*\})?$", re.M)

def build():
    chunks = []
    for path in sorted(glob.glob(os.path.join(CORPUS, "*.md"))):
        text = open(path, encoding="utf-8").read()
        name = os.path.basename(path)
        ms = list(HEADING.finditer(text))
        for i, m in enumerate(ms):
            end = ms[i + 1].start() if i + 1 < len(ms) else len(text)
            body = text[m.end():end].strip()
            if body:
                chunks.append({"id": f"{name}#{len(chunks)}", "file": name, "heading": m.group(2).strip(), "text": body})
    return chunks

CHUNKS = build()
BY_ID = {c["id"]: c for c in CHUNKS}
DOCS = [TOK.findall((c["heading"] + " " + c["text"]).lower()) for c in CHUNKS]
TF = [collections.Counter(d) for d in DOCS]
DF = collections.Counter(t for d in DOCS for t in set(d))
AVG = sum(map(len, DOCS)) / len(DOCS)

def bm25(query, k1=1.5, b=0.75):
    q = TOK.findall(query.lower()); N = len(DOCS); out = []
    for i, tf in enumerate(TF):
        s = 0.0
        for t in q:
            if t in tf:
                idf = math.log(1 + (N - DF[t] + .5) / (DF[t] + .5))
                s += idf * tf[t] * (k1 + 1) / (tf[t] + k1 * (1 - b + b * len(DOCS[i]) / AVG))
        out.append(s)
    return out

mcp = MCPServer("fastapi-docs")

@mcp.tool()
def search_docs(query: str, k: int = 5, snippet_chars: int = 200) -> list[dict]:
    """Search the FastAPI docs. Returns up to k hits with id, file, heading, score and a short snippet.
    Call get_section(id) for the full text of a hit."""
    k = max(1, min(k, 20))
    scores = bm25(query)
    top = sorted(range(len(scores)), key=lambda i: -scores[i])[:k]
    return [{"id": CHUNKS[i]["id"], "file": CHUNKS[i]["file"], "heading": CHUNKS[i]["heading"],
             "score": round(scores[i], 3), "snippet": CHUNKS[i]["text"][:snippet_chars]}
            for i in top if scores[i] > 0]

@mcp.tool()
def get_section(id: str) -> dict:
    """Return the full text of one docs section by the id returned from search_docs."""
    c = BY_ID.get(id)
    if c is None:
        raise ToolError(f"unknown id {id!r}; use an id from search_docs")
    return c

if __name__ == "__main__":
    if len(sys.argv) > 1 and sys.argv[1] == "http":
        mcp.run(transport="streamable-http", port=int(sys.argv[2]))
    else:
        mcp.run()