# prompt_server.py
import asyncio
from mcp.server import Server
import mcp.types as types
from mcp.server.stdio import stdio_server
from typing import Any

PROMPTS = {
    "simple_prompt": types.Prompt(
        name="simple_prompt",
        description="A basic prompt that asks the LLM to act like a pirate.",
    ),
    "explain_code": types.Prompt(
        name="explain-code",
        description="Explain a snippet of code.",
        arguments=[
            types.PromptArgument(
                name="code", description="The code snippet to explain", required=True
            ),
            types.PromptArgument(
                name="language", description="The programming language", required=False
            ),
        ],
    ),
}

app = Server("prompt-test-server")


@app.list_prompts()
async def list_prompts() -> list[types.Prompt]:
    return list(PROMPTS.values())


@app.get_prompt()
async def get_prompt(
    name: str, arguments: dict[str, Any] | None = None
) -> types.GetPromptResult:
    arguments = arguments or {}
    if name == "simple_prompt":
        return types.GetPromptResult(
            messages=[
                types.PromptMessage(
                    role="system",
                    content=types.TextContent(
                        text="You are a helpful pirate assistant. You always answer in pirate slang."
                    ),
                ),
                types.PromptMessage(
                    role="user",
                    content=types.TextContent(
                        text="What are the benefits of using Python?"
                    ),
                ),
            ]
        )
    if name == "explain-code":
        code = arguments.get("code", "")
        language = arguments.get("language", "unknown")
        return types.GetPromptResult(
            messages=[
                types.PromptMessage(
                    role="user",
                    content=types.TextContent(
                        text=f"Explain this {language} code snippet to a five-year-old:\n\n```\n{code}\n```"
                    ),
                )
            ]
        )
    raise ValueError("Prompt not found")


async def main():
    # This server only has prompts, no tools or resources.
    async with stdio_server() as (reader, writer):
        await app.run(reader, writer, app.create_initialization_options())


if __name__ == "__main__":
    asyncio.run(main())