import asyncio
import logging
from contextlib import asynccontextmanager
from typing import AsyncGenerator, List

import anyio
import pytest

from mcp import types, ClientSession, Server
from mcp.shared.memory import create_connected_server_and_client_session
from mcp.types import CallToolResult, TextContent, ProgressToken

logger = logging.getLogger(__name__)
logging.basicConfig(level=logging.INFO)

# --- Server Setup ---
@pytest.fixture
def mcp_server() -> Server:
    server = Server(name="test_notification_server")

    @server.call_tool()
    async def tool_with_notification(name: str, arguments: dict) -> List[TextContent]:
        logger.info("Server: Tool 'tool_with_notification' started.")
        # Get session to send notification
        session = server.request_context.session
        progress_token: ProgressToken | None = None
        if server.request_context.meta:
            progress_token = server.request_context.meta.progressToken

        if progress_token is not None:
            logger.info(f"Server: Sending progress notification for token {progress_token}...")
            await session.send_progress_notification(
                progress_token=progress_token, progress=50.0, total=100.0
            )
            logger.info("Server: Progress notification sent.")
        else:
            logger.warning("Server: No progress token found in request meta.")

        # Simulate some work after notification
        await anyio.sleep(0.05)

        logger.info("Server: Tool 'tool_with_notification' finishing.")
        return [TextContent(type="text", text="Tool finished successfully")]

    return server

# --- Test ---
@pytest.mark.asyncio
async def test_call_tool_with_notification_hangs(mcp_server: Server):
    """
    This test demonstrates the hang when using the default memory stream utility.
    """
    logger.info("Starting test...")
    async with create_connected_server_and_client_session(
        mcp_server,
        raise_exceptions=True # Make test fail clearly on server errors
    ) as client_session:
        logger.info("Client session initialized.")

        tool_name = "tool_with_notification"
        tool_args = {}

        # Add progress token to trigger notification
        request_meta = types.RequestParams.Meta(progressToken="test-token-123")

        logger.info(f"Client: Calling tool '{tool_name}'...")
        try:
            # This await call hangs indefinitely with default unbuffered streams
            result: CallToolResult = await asyncio.wait_for(
                client_session.send_request( # Use send_request to add meta easily
                    types.ClientRequest(
                        types.CallToolRequest(
                            method="tools/call",
                            params=types.CallToolRequestParams(
                                name=tool_name,
                                arguments=tool_args,
                                meta=request_meta # Add meta here
                            )
                        )
                    ),
                    types.CallToolResult # Expected result type
                ),
                timeout=5.0 # Add a timeout to prevent infinite hang in test run
            )
            logger.info(f"Client: Received result: {result}")

            # Assertions (won't be reached if it hangs)
            assert not result.isError
            assert len(result.content) == 1
            assert isinstance(result.content[0], TextContent)
            assert result.content[0].text == "Tool finished successfully"

        except asyncio.TimeoutError:
            pytest.fail("Test timed out - call_tool likely hanged.")
        except Exception as e:
             logger.error(f"Client: Error during call_tool: {e}", exc_info=True)
             pytest.fail(f"call_tool raised an unexpected exception: {e}")