#!/usr/bin/env python3
"""
E2E Sampling Server - Powered by chuk-mcp
Demonstrates server-side sampling capability using chuk-mcp framework.
"""

import asyncio
import logging
import sys

from chuk_mcp.server import MCPServer
from chuk_mcp.protocol.types import ServerCapabilities
from server_helpers import run_stdio_server

# Configure logging to stderr
logging.basicConfig(level=logging.WARNING, stream=sys.stderr)


async def main():
    """Create and run sampling demo server."""
    capabilities = ServerCapabilities(sampling={})  # Empty dict = supported
    server = MCPServer(
        name="sampling-demo-server", version="1.0.0", capabilities=capabilities
    )

    # Register sampling handler
    async def handle_sampling_create_message(message, session_id):
        """
        Handle sampling/createMessage request.

        Server receives sampling request from client and returns simulated AI response.
        """
        # Access params - message is already a Pydantic model
        params = message.params if hasattr(message, "params") else {}
        messages = params.get("messages", []) if isinstance(params, dict) else []

        # Extract user message from first message
        user_message = ""
        if messages and len(messages) > 0:
            first_msg = messages[0]
            content = first_msg.get("content", {})
            user_message = content.get("text", "")

        # Simulate AI response (in production, call actual LLM here)
        response_text = f"This is a simulated AI response to: '{user_message}'. In a real application, this would be generated by an actual LLM."

        # Return properly formatted result
        result = {
            "role": "assistant",
            "content": {"type": "text", "text": response_text},
            "model": "simulated-model",
            "stopReason": "endTurn",
        }

        return server.protocol_handler.create_response(message.id, result), None

    # Register protocol handler
    server.protocol_handler.register_method(
        "sampling/createMessage", handle_sampling_create_message
    )

    # Run server using stdio transport helper
    await run_stdio_server(server)


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