#!/usr/bin/env python3
"""
Generated test file for MCP Tool Error Test
Framework: mcp
Execution mode: async
Streaming: no"""

import os
import sentry_sdk
import asyncio

import socket
import threading
import json
import anyio
from sentry_sdk.integrations.mcp import MCPIntegration
from mcp.client.session import ClientSession
from mcp.server.lowlevel import Server
from mcp.types import (
    Tool,
    TextContent,
    Resource,
    ResourceTemplate,
    Prompt,
    PromptMessage,
    GetPromptResult,
    PromptArgument,
)
from mcp.client.sse import sse_client


sentry_sdk.init(
    dsn=os.environ.get("SENTRY_DSN"),
    traces_sample_rate=1.0,
    send_default_pii=True,
    stream_gen_ai_spans=True,
    integrations=[MCPIntegration()],
)




mcp_server = Server("test-server")

@mcp_server.list_tools()
async def list_tools():
    return [
        Tool(
            name="failing_tool",
            description="A tool that always fails",
            inputSchema={
                "type": "object",
                "properties": {
                    "input": {"type": "string", "description": "Input value"},
                },
                "required": ["input"],
            },
        ),
    ]


@mcp_server.call_tool()
async def call_tool(name, arguments):
    if name == "failing_tool":
        raise Exception("Something went wrong in the tool")
    else:
        raise ValueError(f"Unknown tool: {name}")


def _wait_for_port(port, host="127.0.0.1", timeout=10):
    """Poll until the server is accepting connections on the given port."""
    import time as _time
    deadline = _time.monotonic() + timeout
    while _time.monotonic() < deadline:
        try:
            with socket.create_connection((host, port), timeout=0.5):
                return
        except OSError:
            _time.sleep(0.1)
    raise RuntimeError(f"Server on {host}:{port} did not start within {timeout}s")

SSE_PORT = int(os.environ.get("MCP_SSE_PORT", "8765"))

def run_sse_server():
    """Run the low-level MCP server with SSE transport in a background thread."""
    from mcp.server.sse import SseServerTransport
    from starlette.applications import Starlette
    from starlette.routing import Route, Mount
    import uvicorn

    sse = SseServerTransport("/messages/")

    async def handle_sse(request):
        async with sse.connect_sse(request.scope, request.receive, request._send) as streams:
            await mcp_server.run(
                streams[0], streams[1], mcp_server.create_initialization_options()
            )

    app = Starlette(routes=[
        Route("/sse", endpoint=handle_sse),
        Mount("/messages/", app=sse.handle_post_message),
    ])
    uvicorn.run(app, host="127.0.0.1", port=SSE_PORT)

async def main():
    # Start MCP server with SSE transport in background thread
    server_thread = threading.Thread(target=run_sse_server, daemon=True)
    server_thread.start()
    _wait_for_port(SSE_PORT)

    async with sse_client(f"http://127.0.0.1:{SSE_PORT}/sse") as (read_stream, write_stream):
        async with ClientSession(read_stream, write_stream) as session:
            await session.initialize()
            try:
                result = await session.call_tool("failing_tool", {"input":"test"})
                print(f"Tool result: {result}")
            except Exception as e:
                print(f"Tool error (expected): {type(e).__name__}: {e}")
    # Give the server thread time to finish processing and create spans
    await asyncio.sleep(0.5)

if __name__ == "__main__":
    with sentry_sdk.start_transaction(op="test", name="MCP Tool Error Test"):
        asyncio.run(main())
    sentry_sdk.flush(timeout=5)