"""Simple interactive task server demonstrating elicitation and sampling.

This example shows the simplified task API where:
- server.experimental.enable_tasks() sets up all infrastructure
- ctx.experimental.run_task() handles task lifecycle automatically
- ServerTaskContext.elicit() and ServerTaskContext.create_message() queue requests properly
"""

from typing import Any

import click
import uvicorn
from mcp import types
from mcp.server import Server, ServerRequestContext
from mcp.server.experimental.task_context import ServerTaskContext


async def handle_list_tools(
    ctx: ServerRequestContext, params: types.PaginatedRequestParams | None
) -> types.ListToolsResult:
    return types.ListToolsResult(
        tools=[
            types.Tool(
                name="confirm_delete",
                description="Asks for confirmation before deleting (demonstrates elicitation)",
                input_schema={
                    "type": "object",
                    "properties": {"filename": {"type": "string"}},
                },
                execution=types.ToolExecution(task_support=types.TASK_REQUIRED),
            ),
            types.Tool(
                name="write_haiku",
                description="Asks LLM to write a haiku (demonstrates sampling)",
                input_schema={"type": "object", "properties": {"topic": {"type": "string"}}},
                execution=types.ToolExecution(task_support=types.TASK_REQUIRED),
            ),
        ]
    )


async def handle_confirm_delete(ctx: ServerRequestContext, arguments: dict[str, Any]) -> types.CreateTaskResult:
    """Handle the confirm_delete tool - demonstrates elicitation."""
    ctx.experimental.validate_task_mode(types.TASK_REQUIRED)

    filename = arguments.get("filename", "unknown.txt")
    print(f"\n[Server] confirm_delete called for '{filename}'")

    async def work(task: ServerTaskContext) -> types.CallToolResult:
        print(f"[Server] Task {task.task_id} starting elicitation...")

        result = await task.elicit(
            message=f"Are you sure you want to delete '{filename}'?",
            requested_schema={
                "type": "object",
                "properties": {"confirm": {"type": "boolean"}},
                "required": ["confirm"],
            },
        )

        print(f"[Server] Received elicitation response: action={result.action}, content={result.content}")

        if result.action == "accept" and result.content:
            confirmed = result.content.get("confirm", False)
            text = f"Deleted '{filename}'" if confirmed else "Deletion cancelled"
        else:
            text = "Deletion cancelled"

        print(f"[Server] Completing task with result: {text}")
        return types.CallToolResult(content=[types.TextContent(type="text", text=text)])

    return await ctx.experimental.run_task(work)


async def handle_write_haiku(ctx: ServerRequestContext, arguments: dict[str, Any]) -> types.CreateTaskResult:
    """Handle the write_haiku tool - demonstrates sampling."""
    ctx.experimental.validate_task_mode(types.TASK_REQUIRED)

    topic = arguments.get("topic", "nature")
    print(f"\n[Server] write_haiku called for topic '{topic}'")

    async def work(task: ServerTaskContext) -> types.CallToolResult:
        print(f"[Server] Task {task.task_id} starting sampling...")

        result = await task.create_message(
            messages=[
                types.SamplingMessage(
                    role="user",
                    content=types.TextContent(type="text", text=f"Write a haiku about {topic}"),
                )
            ],
            max_tokens=50,
        )

        haiku = "No response"
        if isinstance(result.content, types.TextContent):
            haiku = result.content.text

        print(f"[Server] Received sampling response: {haiku[:50]}...")
        return types.CallToolResult(content=[types.TextContent(type="text", text=f"Haiku:\n{haiku}")])

    return await ctx.experimental.run_task(work)


async def handle_call_tool(
    ctx: ServerRequestContext, params: types.CallToolRequestParams
) -> types.CallToolResult | types.CreateTaskResult:
    """Dispatch tool calls to their handlers."""
    arguments = params.arguments or {}

    if params.name == "confirm_delete":
        return await handle_confirm_delete(ctx, arguments)
    elif params.name == "write_haiku":
        return await handle_write_haiku(ctx, arguments)

    return types.CallToolResult(
        content=[types.TextContent(type="text", text=f"Unknown tool: {params.name}")],
        is_error=True,
    )


server = Server(
    "simple-task-interactive",
    on_list_tools=handle_list_tools,
    on_call_tool=handle_call_tool,
)

# Enable task support - this auto-registers all handlers
server.experimental.enable_tasks()


@click.command()
@click.option("--port", default=8000, help="Port to listen on")
def main(port: int) -> int:
    starlette_app = server.streamable_http_app()
    print(f"Starting server on http://localhost:{port}/mcp")
    uvicorn.run(starlette_app, host="127.0.0.1", port=port)
    return 0