#!/usr/bin/env python3
"""AAE Lab 10: conceptual MCP 2026-07-28 envelope and authorization boundary.

This is a teaching simulator, not a complete MCP implementation. Use an official MCP SDK
for production wire handling.
"""

PROTOCOL = "2026-07-28"
TOOLS = {
    "get_service_status": {"protected": False},
    "restart_service": {"protected": True},
}


def validate_envelope(headers: dict, body: dict):
    if headers.get("MCP-Protocol-Version") != PROTOCOL:
        raise ValueError("wrong or missing modern protocol version")
    if headers.get("Mcp-Method") != body.get("method"):
        raise ValueError("Mcp-Method does not match JSON-RPC method")
    if body.get("method") == "tools/call":
        name = body.get("params", {}).get("name")
        if headers.get("Mcp-Name") != name:
            raise ValueError("Mcp-Name does not match tool name")
        if name not in TOOLS:
            raise ValueError("unknown tool")
    meta = body.get("params", {}).get("_meta", {})
    if "io.modelcontextprotocol/protocolVersion" not in meta:
        raise ValueError("per-request protocol metadata missing")


def handle(headers: dict, body: dict):
    validate_envelope(headers, body)
    method = body["method"]
    if method == "server/discover":
        return {"protocolVersion": PROTOCOL, "capabilities": {"tools": {}}}
    if method == "tools/call":
        name = body["params"]["name"]
        if TOOLS[name]["protected"] and not headers.get("Authorization", "").startswith("Bearer "):
            raise PermissionError("authorization required for protected tool")
        return {"resultType": "complete", "tool": name, "content": "fixture result"}
    raise ValueError("unsupported method")


def meta():
    return {"io.modelcontextprotocol/protocolVersion": PROTOCOL,
            "io.modelcontextprotocol/clientInfo": {"name": "aae-lab", "version": "1.0"}}

print("AAE Lab 10 — MCP 2026-07-28 concepts\n")
discover_body = {"jsonrpc": "2.0", "id": 1, "method": "server/discover", "params": {"_meta": meta()}}
discover_headers = {"MCP-Protocol-Version": PROTOCOL, "Mcp-Method": "server/discover"}
print("DISCOVER:", handle(discover_headers, discover_body))

read_body = {"jsonrpc": "2.0", "id": 2, "method": "tools/call",
             "params": {"name": "get_service_status", "arguments": {"service": "nginx"}, "_meta": meta()}}
read_headers = {"MCP-Protocol-Version": PROTOCOL, "Mcp-Method": "tools/call", "Mcp-Name": "get_service_status"}
print("READ TOOL:", handle(read_headers, read_body))

change_body = {"jsonrpc": "2.0", "id": 3, "method": "tools/call",
               "params": {"name": "restart_service", "arguments": {"service": "nginx"}, "_meta": meta()}}
change_headers = {"MCP-Protocol-Version": PROTOCOL, "Mcp-Method": "tools/call", "Mcp-Name": "restart_service"}
try:
    handle(change_headers, change_body)
except PermissionError as exc:
    print("✓ Protected tool rejected without auth:", exc)
else:
    raise SystemExit("FAIL: protected tool executed without authorization")

change_headers["Authorization"] = "Bearer demo-identity"
print("PROTECTED TOOL WITH AUTH:", handle(change_headers, change_body))

assert "Mcp-Session-Id" not in discover_headers
assert discover_body["method"] == "server/discover"
print("✓ Modern teaching path uses per-request metadata and no protocol session ID.")
print("✓ Discovery did not replace authorization.")