import threading
import time
import asyncio
from typing import Literal
from datetime import timedelta

from mcp import ClientSession
from mcp.client.streamable_http import streamablehttp_client
from mcp.types import ImageContent as MCPImageContent


def start_comprehensive_mcp_server(transport: Literal["sse", "streamable-http"], port=int):
    from mcp.server import FastMCP

    mcp = FastMCP("Comprehensive MCP Server", port=port)

    @mcp.tool(description="Tool that will timeout")
    def timeout_tool() -> str:
        time.sleep(10)
        return "This tool has timed out"

    @mcp.tool(description="Calculator tool which performs calculations")
    def calculator(x: int, y: int) -> int:
        return x + y

    mcp.run(transport=transport)


# Start server
server_thread = threading.Thread(
    target=start_comprehensive_mcp_server, 
    kwargs={"transport": "streamable-http", "port": 8001}, 
    daemon=True
)
server_thread.start()
time.sleep(3)
print("SERVER STARTED", flush=True)


async def test_raw_mcp():
    print("CREATING RAW MCP CONNECTION...", flush=True)
    
    async with streamablehttp_client(
        url="http://127.0.0.1:8001/mcp",
        sse_read_timeout=2,  # 2 second timeout
    ) as (read_stream, write_stream, get_session_id):
        print("TRANSPORT CONNECTED", flush=True)
        
        async with ClientSession(read_stream, write_stream) as session:
            print("SESSION CREATED", flush=True)
            
            await session.initialize()
            print("SESSION INITIALIZED", flush=True)
            
            print("CALLING TOOL (with 2s timeout)...", flush=True)
            try:
                result = await session.call_tool(
                    "timeout_tool", 
                    arguments={},
                    # read_timeout_seconds=timedelta(seconds=5), # This needs to be provided or the client will hang forever
                )
                print(f"TOOL RESULT: {result}", flush=True)
            except Exception as e:
                print(f"TOOL CALL EXCEPTION: {type(e).__name__}: {e}", flush=True)
            
            print("DONE", flush=True)


print("RUNNING TEST...", flush=True)
asyncio.run(test_raw_mcp())
print("TEST COMPLETE", flush=True)