"""Self-contained MCP server demonstrating JSON Schema 2020-12 tool inputs,
enabled by SEP-2106 (tool inputSchema/outputSchema conform to JSON Schema
2020-12). Each tool in the shared ``schemas.json`` exercises a different 2020-12
keyword that was not permitted in a tool input before SEP-2106 (enum varieties,
top-level oneOf, discriminated unions, $ref/$defs, if/then/else, prefixItems
tuples).

This is the Python sibling of the TypeScript ``demo-server-tool-inputs`` server;
the two read the SAME single-source-of-truth ``../schemas.json`` and advertise
each raw JSON Schema verbatim in tools/list. The MCP Python SDK's low-level
``Server`` passes ``input_schema`` through untouched (no Pydantic remodeling) and
does NOT auto-validate tool-call inputs, so — exactly as the TS server needs an
explicit ``CfWorkerJsonSchemaValidator({ draft: "2020-12" })`` — this server runs
each tool's arguments through a per-tool ``jsonschema.Draft202012Validator``.
jsonschema>=4.20 defaults to the 2020-12 dialect, so ``prefixItems`` /
``items: false`` are honored (a valid ``[x, y]`` tuple is accepted, a 3-element
one rejected).

SEP-2106: https://github.com/modelcontextprotocol/modelcontextprotocol/blob/main/seps/2106-json-schema-2020-12.md

Run over stdio:  uv run python server.py   (then point MCP Inspector at it)

The MCP Python SDK low-level ``Server`` API (constructor-based ``on_list_tools`` /
``on_call_tool`` handlers, ``types.Tool``, ``streamable_http_app``) was confirmed
against mcp 2.0.0b1 (modelcontextprotocol/python-sdk). In 2.0 beta the wire types
moved out of ``mcp.types`` into the standalone ``mcp_types`` package, so this
module does ``import mcp_types as types`` (the SDK's own idiom); every ``types.X``
reference is otherwise unchanged from the 2.0 alpha.
"""

from __future__ import annotations

import json
import os
from pathlib import Path
from typing import Any

import anyio
import mcp_types as types
from jsonschema import Draft202012Validator
from mcp.server import Server, ServerRequestContext

# schemas.json is the single source of truth, one level up at the topic root:
# demo-server-py/ -> tool-schemas/. Loaded verbatim, mirroring the TS server's
# demo-tools.ts relative-path load, so the two consumers can never drift.
# SCHEMAS_PATH overrides the location so the Docker/HF-Space image can point at a
# vendored copy (the build context can't reach ../schemas.json at runtime); the
# default keeps stdio + selftest working unchanged from the repo tree.
_SCHEMAS_PATH = Path(os.environ.get("SCHEMAS_PATH") or Path(__file__).resolve().parent.parent / "schemas.json")
DEMO_TOOLS: list[dict[str, Any]] = json.loads(_SCHEMAS_PATH.read_text(encoding="utf-8"))["tools"]

# One compiled 2020-12 validator per tool, keyed by tool name.
_VALIDATORS: dict[str, Draft202012Validator] = {t["name"]: Draft202012Validator(t["jsonSchema"]) for t in DEMO_TOOLS}
# Validators for each tool's outputSchema (the result/structuredContent shape).
_OUTPUT_VALIDATORS: dict[str, Draft202012Validator] = {
    t["name"]: Draft202012Validator(t["outputSchema"]) for t in DEMO_TOOLS if t.get("outputSchema")
}

# A canned `structuredContent` per tool that CONFORMS to that tool's outputSchema —
# so the server demonstrates a real 2020-12 outputSchema round-trip (advertise +
# return + validate). selftest asserts each validates against its outputSchema.
_OUTPUTS: dict[str, Any] = {
    "get-enum-selections": {"result": "applied", "appliedCount": 3},
    "lookup-record": {"found": True, "name": "Acme Corp"},
    "create-payment": {"receiptId": "rcpt_001", "outcome": {"status": "settled", "settledAt": "2026-06-25T12:00:00Z"}},
    "create-shipment": {
        "shipmentId": "shp_001",
        "origin": {"city": "San Francisco", "eta": "2026-06-26"},
        "destination": {"city": "London", "eta": "2026-06-28"},
    },
    "register-address": {"status": "active", "validUntil": "2027-06-25"},
    "plot-point": {"plotted": [3, 7], "ok": True},
}


def _echo(name: str, arguments: dict[str, Any] | None) -> types.CallToolResult:
    """Shared handler: echo the validated arguments AND return a 2020-12-conforming
    ``structuredContent`` for the tool's outputSchema. Mirrors the TS server.
    """
    entries = list((arguments or {}).items())
    if entries:
        text = "You provided:\n" + "\n".join(f"- {k}: {json.dumps(v)}" for k, v in entries)
    else:
        text = "No arguments provided."
    return types.CallToolResult(
        content=[types.TextContent(type="text", text=text)],
        structured_content=_OUTPUTS.get(name),
    )


def _validation_error(name: str, errors: list[Any]) -> types.CallToolResult:
    """Message-rich isError result listing every failing 2020-12 keyword."""
    lines = [f"Invalid arguments for '{name}' (JSON Schema 2020-12 validation failed):"]
    for err in errors:
        location = "/".join(str(p) for p in err.absolute_path) or "(root)"
        lines.append(f"- {location}: failed keyword '{err.validator}' — {err.message}")
    return types.CallToolResult(
        content=[types.TextContent(type="text", text="\n".join(lines))],
        is_error=True,
    )


async def handle_list_tools(
    ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
    return types.ListToolsResult(
        tools=[
            types.Tool(
                name=tool["name"],
                title=tool["title"],
                description=tool["description"],
                # Advertised VERBATIM — the raw 2020-12 dicts from schemas.json.
                input_schema=tool["jsonSchema"],
                output_schema=tool.get("outputSchema"),
                annotations=types.ToolAnnotations(
                    read_only_hint=True,
                    destructive_hint=False,
                    idempotent_hint=True,
                    open_world_hint=False,
                ),
            )
            for tool in DEMO_TOOLS
        ]
    )


async def handle_call_tool(ctx: ServerRequestContext, params: types.CallToolRequestParams) -> types.CallToolResult:
    validator = _VALIDATORS.get(params.name)
    if validator is None:
        raise ValueError(f"Unknown tool: {params.name}")

    arguments = params.arguments or {}
    # The low-level server does NOT validate tool inputs, so do it explicitly
    # against the tool's own 2020-12 schema (parity with the TS server).
    errors = sorted(validator.iter_errors(arguments), key=lambda e: list(e.absolute_path))
    if errors:
        return _validation_error(params.name, errors)
    return _echo(params.name, arguments)


def create_server() -> Server[Any]:
    """Build the demo server (factory so it can be mounted over HTTP too)."""
    return Server(
        "tool-input-schemas-demo",
        version="0.1.0",
        on_list_tools=handle_list_tools,
        on_call_tool=handle_call_tool,
    )


async def _amain() -> None:
    from mcp.server.stdio import stdio_server

    server = create_server()
    async with stdio_server() as (read_stream, write_stream):
        await server.run(read_stream, write_stream, server.create_initialization_options())


def main() -> None:
    anyio.run(_amain)


if __name__ == "__main__":
    main()