Prompts

from typing import Annotated, Literal

from pydantic import BaseModel, Field

from mcp.server import MCPServer
from mcp.server.mcpserver.prompts.base import AssistantMessage, Message, UserMessage
from mcp.server.mcpserver.exceptions import ResourceNotFoundError, ToolError
from mcp.types import ToolAnnotations

mcp = MCPServer("Shop support")

ORDERS = {
    "A17": {"item": "blue mug", "status": "shipped", "total": 12.50},
    "B42": {"item": "desk lamp", "status": "processing", "total": 48.00},
}

ARTICLES = {
    "Where is my order?": "orders",
    "Changing an order": "orders",
    "How refunds work": "refunds",
    "Refunds for damaged items": "refunds",
    "Resetting your password": "account",
}


class Order(BaseModel):
    id: str
    item: str
    status: Literal["processing", "shipped", "delivered"]
    total: float


@mcp.tool(annotations=ToolAnnotations(read_only_hint=True))
def lookup_order(order_id: str) -> Order:
    """Look up an order by its id."""
    if order_id not in ORDERS:
        raise ToolError(f"No order with id {order_id!r}. Order ids look like A17.")
    return Order(id=order_id, **ORDERS[order_id])


@mcp.tool()
def search_help(
    query: Annotated[str, Field(description="Words to look for in the help articles.")],
    topic: Literal["orders", "refunds", "account"] | None = None,
    limit: Annotated[int, Field(ge=1, le=5)] = 3,
) -> str:
    """Search the help centre and return matching article titles."""
    found = [title for title, t in ARTICLES.items() if query.lower() in title.lower() and topic in (None, t)]
    return "; ".join(found[:limit]) or "No articles found."


@mcp.tool(
    title="Refund an order",
    annotations=ToolAnnotations(read_only_hint=False, destructive_hint=False, idempotent_hint=False),
)
def refund_order(order_id: str, reason: str) -> str:
    """Refund the full total of an order to the customer's card."""
    if order_id not in ORDERS:
        raise ToolError(f"No order with id {order_id!r}.")
    return f"Refunded {ORDERS[order_id]['total']:.2f} for order {order_id}: {reason}."


@mcp.resource("policy://refunds", mime_type="text/markdown")
def refund_policy() -> str:
    """The shop's refund policy."""
    return "# Refunds\n\nFull refund within 30 days of delivery. Damaged items: refund or replacement."


@mcp.resource("orders://{order_id}", mime_type="application/json")
def order_record(order_id: str) -> Order:
    """The full record for one order."""
    if order_id not in ORDERS:
        raise ResourceNotFoundError(f"No order with id {order_id!r}.")
    return Order(id=order_id, **ORDERS[order_id])


@mcp.prompt(title="Reply to a customer")
def reply_to_customer(ticket: str, tone: str = "friendly") -> list[Message]:
    """Draft a reply to a support ticket."""
    return [
        UserMessage(f"Write a {tone} reply to this support ticket:\n\n{ticket}"),
        AssistantMessage("Hello, and thank you for getting in touch."),
    ]