import asyncio
import json
import os
import click

from contextlib import asynccontextmanager
from mcp import ClientSession
from mcp.client.sse import sse_client
from mcp.client.streamable_http import streamable_http_client
from mcp.types import CallToolResult, ReadResourceResult
from abi_core.common.utils import abi_logging

from abi_core.common.utils import abi_logging


@asynccontextmanager
async def init_session(host, port, transport='streamable-http'):
    """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 Streamable HTTP 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.
        port: The port number of the MCP server.
        transport: The communication transport to use ('sse' or 'streamable-http').
                  Defaults to 'streamable-http'.

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

    Raises:
        ValueError: If an unsupported transport type is provided.
        Exception: Other potential exceptions during client initialization or
                   session setup.
    """

    if transport not in ['sse', 'streamable-http']:
        abi_logging(f'[❌] Unsupported Transport type {transport}')

        raise ValueError(
            f"Unsupported transport type: {transport}. Must be 'sse' or 'streamable-http'"
        )
    
    if transport == 'sse':
        url = f'http://{host}:{port}/sse'
        abi_logging(f'[🔌] Connecting to MCP server via SSE at {url}')
        
        try:
            async with sse_client(url) as (read_stream, write_stream):
                abi_logging('[✅] SSE connection established')
                try:
                    async with ClientSession(
                        read_stream=read_stream,
                        write_stream=write_stream,
                    ) as session:
                        abi_logging('[⚙️] SSE Client Session Initializing...')
                        await session.initialize()
                        abi_logging('[✅] SSE Client Session Initialized Successfully')
                        yield session
                except Exception as e:
                    abi_logging(f'[❌] Error initializing ClientSession: {e}')
                    raise
        except Exception as e:
            abi_logging(f'[❌] Error connecting to SSE server at {url}: {e}')
            raise
    
    elif transport == 'streamable-http':
        url = f'http://{host}:{port}/mcp'
        abi_logging(f'[🔌] Connecting to MCP server via Streamable HTTP at {url}')
        
        session = None
        read_stream = None
        write_stream = None
        
        try:
            # streamable_http_client returns 3 elements: (read_stream, write_stream, connection_metadata)
            # The third element contains connection metadata and is typically ignored
            async with streamable_http_client(url) as (read_stream, write_stream, _):
                abi_logging('[✅] Streamable HTTP connection established')

                try:
                    async with ClientSession(
                        read_stream=read_stream,
                        write_stream=write_stream,
                    ) as session:
                        abi_logging('[⚙️] Streamable HTTP Client Session Initializing...')
                        await session.initialize()
                        abi_logging('[✅] Streamable HTTP Client Session Initialized Successfully')
                        yield session
                except Exception as e:
                    abi_logging(f'[❌] Error initializing ClientSession: {e}')
                    raise
                finally:
                    # Ensure proper cleanup of session
                    if session:
                        try:
                            abi_logging('[🧹] Cleaning up Streamable HTTP session')
                        except Exception as cleanup_error:
                            abi_logging(f'[⚠️] Error during session cleanup: {cleanup_error}')
        except Exception as e:
            abi_logging(f'[❌] Error connecting to Streamable HTTP server at {url}: {e}')
            raise
        finally:
            # Ensure streams are properly closed
            abi_logging('[🧹] Streamable HTTP connection cleanup complete')

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

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

    Returns:
        The result of the tool call.
    """

    return await session.call_tool(
        name='find_agent',
        arguments={
            'query': query,
            '_request_context':ctx
        },
    )

async def find_resource(session: ClientSession, resource: str) -> 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.
    """
    abi_logging(f'[📖] Reading resource: {resource}')
    return await session.read_resource(resource)


async def recommend_agents(
    session: ClientSession,
    task_description: str,
    max_agents: int,
    ctx: dict
) -> CallToolResult:
    """Call the 'recommend_agents' tool on the connected MCP server.

    Args:
        session: The active ClientSession.
        task_description: Description of the task requiring multiple agents.
        max_agents: Maximum number of agents to recommend.
        ctx: Request context for authentication.

    Returns:
        The result of the tool call.
    """
    return await session.call_tool(
        name='recommend_agents',
        arguments={
            'task_description': task_description,
            'max_agents': max_agents,
            '_request_context': ctx
        },
    )


async def check_agent_capability(
    session: ClientSession,
    agent_name: str,
    required_tasks: list,
    ctx: dict
) -> CallToolResult:
    """Call the 'check_agent_capability' tool on the connected MCP server.

    Args:
        session: The active ClientSession.
        agent_name: Name of the agent to check.
        required_tasks: List of required task names.
        ctx: Request context for authentication.

    Returns:
        The result of the tool call.
    """
    return await session.call_tool(
        name='check_agent_capability',
        arguments={
            'agent_name': agent_name,
            'required_tasks': required_tasks,
            '_request_context': ctx
        },
    )


async def check_agent_health(
    session: ClientSession,
    agent_name: str,
    ctx: dict
) -> CallToolResult:
    """Call the 'check_agent_health' tool on the connected MCP server.

    Args:
        session: The active ClientSession.
        agent_name: Name of the agent to check.
        ctx: Request context for authentication.

    Returns:
        The result of the tool call.
    """
    return await session.call_tool(
        name='check_agent_health',
        arguments={
            'agent_name': agent_name,
            '_request_context': ctx
        },
    )


async def register_agent(
    session: ClientSession,
    agent_card: dict,
    ctx: dict
) -> CallToolResult:
    """Call the 'register_agent' tool on the connected MCP server.

    Args:
        session: The active ClientSession.
        agent_card: Complete agent card dictionary with auth credentials.
        ctx: Request context for authentication.

    Returns:
        The result of the tool call with registration status.
    """
    return await session.call_tool(
        name='register_agent',
        arguments={
            'agent_card': agent_card,
            '_request_context': ctx
        },
    )

async def custom_tool(
    session: ClientSession,
    tool_name: str,
    ctx: dict,
    payload: dict = None
) -> CallToolResult:
    """Call custom tool on the connected MCP server.
    
    Args:
        session: The active ClientSession.
        tool_name: The name of the tool to call.
        ctx: Request context for authentication.
        payload: Optional payload to be used by the tool (default: None).

    Returns:
        The result of the tool call.
    """
    if payload is None:
        payload = {}
    
    # Merge payload with request context
    arguments = {
        '_request_context': ctx,
        **payload  # Unpack payload directly into arguments
    }
    
    abi_logging(f'[🔧] Calling custom tool: {tool_name}')
    return await session.call_tool(
        name=tool_name,
        arguments=arguments,
    )
