"""Discover an extension's capability entry, call its tool, then send its vendor method."""

from typing import Literal

import mcp.types as types
from mcp.client import Client, advertise
from mcp.types import TextContent
from stories._harness import Target, run_client

EXTENSION_ID = "com.example/catalog"


class SearchParams(types.RequestParams):
    query: str
    limit: int = 3


class SearchRequest(types.Request[SearchParams, Literal["com.example/search"]]):
    method: Literal["com.example/search"] = "com.example/search"
    params: SearchParams


class SearchResult(types.Result):
    items: list[str]


async def main(target: Target, *, mode: str = "auto") -> None:
    # Declare the extension client-side so the server's `require_client_extension`
    # gate on `com.example/search` passes.
    async with Client(target, mode=mode, extensions=[advertise(EXTENSION_ID)]) as client:
        # The extensions capability map rides `server/discover` (modern only). On a
        # legacy connection it is absent, so assert it only when present.
        if client.server_capabilities.extensions is not None:
            assert client.server_capabilities.extensions == {EXTENSION_ID: {"suggest": True}}, (
                client.server_capabilities.extensions
            )

        # The extension's tool is a regular tool: listed and callable like any other.
        listed = await client.list_tools()
        assert [tool.name for tool in listed.tools] == ["suggest"], listed
        result = await client.call_tool("suggest", {"prefix": "mcp"})
        assert isinstance(result.content[0], TextContent)
        assert result.content[0].text == "mcp-suggestion", result.content[0].text

        # Vendor methods drop one layer to `client.session` (see custom_methods/).
        request = SearchRequest(params=SearchParams(query="mcp", limit=3))
        found = await client.session.send_request(request, SearchResult)
        assert found.items == ["mcp-0", "mcp-1", "mcp-2"], found


if __name__ == "__main__":
    run_client(main)