import asyncio
import json
import logging
from contextlib import asynccontextmanager

import click
from mcp import ClientSession, StdioServerParameters, stdio_client
from mcp.client.sse import sse_client
from mcp.types import CallToolResult, ReadResourceResult

logger = logging.getLogger(__name__)


@asynccontextmanager
async def init_session(host, port, transport):
    """Initializes and manages an MCP ClientSession based on the specified transport.

    This asynchronous context manager establishes a connection to an MCP server
    using either Server-Sent Events (SSE) or Standard I/O (STDIO) transport.
    It handles the setup and teardown of the connection and yields an active
    `ClientSession` object ready for communication.

    Args:
        host: The hostname or IP address of the MCP server (used for SSE).
        port: The port number of the MCP server (used for SSE).
        transport: The communication transport to use ('sse' or 'stdio').

    Yields:
        ClientSession: An initialized and ready-to-use MCP client session.

    Raises:
        ValueError: If an unsupported transport type is provided (implicitly,
                    as it won't match 'sse' or 'stdio').
        Exception: Other potential exceptions during client initialization or
                   session setup.
    """
    if transport == "sse":
        url = f"http://{host}:{port}/sse"
        async with sse_client(url) as (read_stream, write_stream):
            async with ClientSession(
                read_stream=read_stream, write_stream=write_stream
            ) as session:
                logger.debug("SSE ClientSession created, initializing...")
                await session.initialize()
                logger.info("SSE ClientSession initialized successfully...")
                yield session
    elif transport == "stdio":
        stdio_params = StdioServerParameters(
            command="uv",
            args=["run", "a2a-mcp_servers"],
        )
        async with stdio_client(stdio_params) as (read_stream, write_stream):
            async with ClientSession(
                read_stream=read_stream,
                write_stream=write_stream,
            ) as session:
                logger.debug("STDIO ClientSession created, initializing...")
                await session.initialize()
                logger.info("STDIO ClientSession initialized successfully.")
                yield session
    else:
        logger.error(f"Unsupported transport type: {transport}")
        raise ValueError(
            f"Unsupported transport type: {transport}. Must be 'sse' or 'stdio'."
        )


async def find_agent(session: ClientSession, query) -> CallToolResult:
    """Calls the 'find_agent' tool on the connected MCP server.

    Args:
        session: The active ClientSession,
        query: The natural language query to send to the 'find_agent' tool.

    Returns:
        The result of the tool call.
    """
    logger.info(f"Calling 'find_agent' tool with query: '{query[:50]}'")
    return await session.call_tool(
        name="find_agent",
        arguments={
            "query": query,
        },
    )


async def find_resource(session: ClientSession, resource) -> ReadResourceResult:
    """Reads a resource from the connected MCP server.

    Args:
        session: The active ClientSession.
        resource: The URI of the resource to read (e.g., 'resource://agent_cards/list').

    Returns:
        The result of the resource read operation.
    """
    logger.info(f"Reading resource: {resource}")
    return await session.read_resource(uri=resource)


# Test util
async def main(host, port, transport, query, resource, tool):
    """Main asynchronous function to connect to the MCP server and execute commands.

    Used for local testing.

    Args:
        host: Server hostname.
        port: Server port.
        transport: Connection transport ('sse' or 'stdio').
        query: Optional query string for the 'find_agent' tool.
        resource: Optional resource URI to read.
    """
    logger.info("Starting Client to connect to MCP")
    async with init_session(host, port, transport) as session:
        if query:
            result = await find_agent(session, query)
            data = json.loads(result.content[0].text)
            logger.info(json.dumps(data, indent=2))
        if resource:
            result = await find_resource(session, resource)
            logger.info(result)
            data = json.loads(result.contents[0].text)
            logger.info(json.dumps(data, indent=2))


# Command line tester
@click.command()
@click.option("--host", default="localhost", help="SSE Host")
@click.option("--port", default="10100", help="SSE Port")
@click.option("--transport", default="stdio", help="MCP Transport")
@click.option("--find_agent", help="Query to find an agent")
@click.option("--resource", help="URI of the resource to locate")
def cli(host, port, transport, find_agent, resource, tool_name):
    """A command-line client to interact with the Agent Cards MCP server."""
    asyncio.run(main(host, port, transport, find_agent, resource, tool_name))


if __name__ == "__main__":
    cli()
