from collections.abc import Sequence
from typing import Any, Literal

import mcp.types as types
from mcp import Client
from mcp.client import ClaimContext, ClientExtension, ResultClaim
from mcp.server.context import CallNext, HandlerResult, ServerRequestContext
from mcp.server.extension import Extension
from mcp.server.mcpserver import MCPServer, require_client_extension

EXTENSION_ID = "com.example/receipts"


class ReceiptResult(types.Result):
    """The claimed result shape; `result_type` pins the wire tag."""

    result_type: Literal["receipt"] = "receipt"
    receipt_token: str


class ReceiptIssuer(Extension):
    """Server half: answers `buy` with a receipt instead of a final result."""

    identifier = EXTENSION_ID

    async def intercept_tool_call(
        self,
        params: types.CallToolRequestParams,
        ctx: ServerRequestContext[Any, Any],
        call_next: CallNext,
    ) -> HandlerResult:
        if params.name != "buy":
            return await call_next(ctx)
        require_client_extension(ctx, EXTENSION_ID)
        return {"resultType": "receipt", "receiptToken": "r-117"}


class Receipts(ClientExtension):
    """Client half: claims the `receipt` shape and supplies the code that finishes it."""

    identifier = EXTENSION_ID

    def claims(self) -> Sequence[ResultClaim[Any]]:
        return [ResultClaim(result_type="receipt", model=ReceiptResult, resolve=self._redeem)]

    async def _redeem(self, claimed: ReceiptResult, ctx: ClaimContext) -> types.CallToolResult:
        return await ctx.session.call_tool("redeem", {"token": claimed.receipt_token})


mcp = MCPServer("shop", extensions=[ReceiptIssuer()])


@mcp.tool()
def buy(item: str) -> types.CallToolResult:
    """Buy an item."""
    raise NotImplementedError  # ReceiptIssuer answers `buy` before the tool runs


@mcp.tool()
def redeem(token: str) -> str:
    """Exchange a receipt token for the goods."""
    return f"goods for {token}"


async def main() -> None:
    async with Client(mcp, extensions=[Receipts()]) as client:
        result = await client.call_tool("buy", {"item": "lamp"})
        print(result.content)
        # [TextContent(text='goods for r-117')]