"""Drive docs_server.py with the official mcp Python SDK client over stdio and streamable HTTP.
Measures: connect time, tools/list payload, tools/call payload sizes by k and snippet length, per-call latency."""
import asyncio, json, os, statistics as st, subprocess, sys, time
import httpx2
from mcp import Client
from mcp.client.stdio import StdioServerParameters

HERE = os.path.dirname(os.path.abspath(__file__))
SERVER = os.path.join(HERE, "docs_server.py")
QUERIES = ["background tasks after response", "OAuth2 password flow", "response_model exclude unset",
           "CORS middleware allow origins", "dependency injection with yield", "upload a file",
           "websocket", "custom exception handler", "path parameters with enum", "testing with TestClient"]

def payload(res):
    return sum(len(c.text.encode()) for c in res.content if hasattr(c, "text"))

def wire(res):
    # whole result object as JSON: text content blocks plus structured_content
    return len(res.model_dump_json(by_alias=True, exclude_none=True).encode())

async def run(label, target, mode='auto'):
    t0 = time.perf_counter()
    async with Client(target, mode=mode) as c:
        connect_ms = (time.perf_counter() - t0) * 1000
        tools = await c.list_tools()
        list_bytes = len(tools.model_dump_json(by_alias=True, exclude_none=True).encode())
        print(f"\n== {label} ==")
        print(f"connect: {connect_ms:.0f} ms; tools: {[t.name for t in tools.tools]}; tools/list payload: {list_bytes} bytes")
        for k, sn in [(3, 0), (3, 200), (5, 200), (10, 200), (10, 800), (20, 800)]:
            sizes = []; wires = []
            for q in QUERIES:
                r = await c.call_tool("search_docs", {"query": q, "k": k, "snippet_chars": sn})
                sizes.append(payload(r)); wires.append(wire(r))
            print(f"search_docs k={k:<2} snippet={sn:<3}: text blocks mean {st.mean(sizes):.0f} bytes, max {max(sizes)}; whole result JSON mean {st.mean(wires):.0f} bytes over {len(QUERIES)} queries")
        r2 = await c.call_tool("search_docs", {"query": QUERIES[0], "k": 3})
        print("content blocks for k=3:", len(r2.content), "| structured_content present:", r2.structured_content is not None)
        # get_section sizes for the top hit of each query
        gs = []
        for q in QUERIES:
            rr = await c.call_tool("search_docs", {"query": q, "k": 1})
            top = rr.structured_content
            top = top["result"][0] if isinstance(top, dict) and "result" in top else (top[0] if top else None)
            g = await c.call_tool("get_section", {"id": top["id"]})
            gs.append(payload(g))
        print(f"get_section on top hit: mean {st.mean(gs):.0f} bytes, max {max(gs)} bytes")
        e = await c.call_tool("get_section", {"id": "nope"})
        print("unknown id -> is_error:", e.is_error, "|", e.content[0].text[:120])
        lat = []
        for i in range(100):
            t = time.perf_counter()
            await c.call_tool("search_docs", {"query": QUERIES[i % 10], "k": 5})
            lat.append((time.perf_counter() - t) * 1000)
        lat.sort()
        print(f"100 search_docs calls (k=5): median {st.median(lat):.1f} ms, p95 {lat[94]:.1f} ms, max {lat[-1]:.1f} ms")

async def main():
    await run("stdio", StdioServerParameters(command=sys.executable, args=[SERVER]))
    p = subprocess.Popen([sys.executable, SERVER, "http", "8765"], stdout=subprocess.DEVNULL, stderr=subprocess.DEVNULL)
    try:
        for _ in range(100):
            try:
                httpx2.get("http://127.0.0.1:8765/mcp", timeout=1); break
            except Exception:
                time.sleep(0.2)
        await run("streamable HTTP (mode=auto)", "http://127.0.0.1:8765/mcp")
        await run("streamable HTTP (mode=legacy, initialize handshake)", "http://127.0.0.1:8765/mcp", mode="legacy")
    finally:
        p.terminate()

asyncio.run(main())