import json
import os

from abi_core.common.utils import get_mcp_server_config, abi_logging
from abi_core.abi_mcp import client
from abi_core.security.agent_auth import build_semantic_context_from_card

from langchain.tools import tool
from a2a.types import AgentCard

AGENT_CARD_PATH = os.getenv('AGENT_CARD')
_mcp_config = get_mcp_server_config()

@tool
async def tool_find_agent(query) -> AgentCard:
        """Find an Angent to complete especific task"""
        
        async with client.init_session(
            _mcp_config.host,
            _mcp_config.port,
            _mcp_config.transport
        ) as mcp_session:
            
            context = build_semantic_context_from_card(
                AGENT_CARD_PATH,
                tool_name="find_agent",
                query=query
            )

            mcp_response = await client.find_agent(mcp_session, query, _request_context=context)
            if hasattr(mcp_session, 'content') and mcp_response.content:
                try:
                    if isinstance(mcp_response.content, list) and mcp_response.content:
                        agent_card_json = json.loads(mcp_response.content[0].text)
                    else:
                        agent_card_json = mcp_response.content
                    
                    if agent_card_json:
                        return AgentCard(**agent_card_json)
                    else:
                        return None
                except Exception as e:
                    abi_logging('[X] Error parsing agent card {e}')
            else:
                abi_logging(f'No Agents found')
                return None

@tool
async def tool_list_agents(query) -> list[AgentCard]:
        """List Agents that can complete especific task"""
        
        async with client.init_session(
            _mcp_config.host,
            _mcp_config.port,
            _mcp_config.transport
        ) as mcp_session:
            
            context = build_semantic_context_from_card(
                AGENT_CARD_PATH,
                tool_name="list_agents",
                query=query
            )

            mcp_response = await client.list_agents(mcp_session, query, _request_context=context)
            agents = []
            if hasattr(mcp_session, 'content') and mcp_response.content:
                try:
                    for content in mcp_response.content:
                        agent_card_json = json.loads(content.text)
                        agents.append(AgentCard(**agent_card_json))
                    return agents
                except Exception as e:
                    abi_logging('[X] Error parsing agent cards {e}')
            else:
                abi_logging(f'No Agents found')
                return agents