"""
Semantic layer tools for agent discovery and coordination.

This module provides LangChain tools that agents can use to interact with
the semantic layer for discovering other agents, checking capabilities, and
monitoring health.
"""

import json
import os
from typing import Optional, List, Dict, Any

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: str) -> Optional[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:
        abi_logging(f"[🔍] Searching for agent matching query {query}")
        context = build_semantic_context_from_card(
            AGENT_CARD_PATH,
            tool_name="find_agent",
            query=query
        )

        mcp_response = await client.find_agent(mcp_session, query, context)
        if hasattr(mcp_response, '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(f'[X] Error parsing agent card {e}')
        else:
            abi_logging(f'No Agents found')
            return None

@tool
async def tool_list_agents(query: str) -> 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
        )
        resource = f'resource://agent_cards/list/{query}'
        abi_logging(f"[🔍] Listing agents matching query {query}")
        mcp_response = await client.find_resource(mcp_session, resource)
        agents = []
        if hasattr(mcp_response, '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(f'[X] Error parsing agent cards {e}')
        else:
            abi_logging(f'No Agents found')
            return agents


@tool
async def tool_recommend_agents(
    task_description: str,
    max_agents: int = 3
) -> List[Dict[str, Any]]:
    """Recommend multiple agents for a complex task.
    
    Args:
        task_description: Description of the task requiring multiple agents
        max_agents: Maximum number of agents to recommend (default: 3)
    
    Returns:
        List of recommended agents with relevance scores and confidence levels
    """
    async with client.init_session(
        _mcp_config.host,
        _mcp_config.port,
        _mcp_config.transport
    ) as mcp_session:
        abi_logging(f"[🔍] Recommending agents for: {task_description}")
        
        context = build_semantic_context_from_card(
            AGENT_CARD_PATH,
            tool_name="recommend_agents",
            query=task_description
        )
        
        mcp_response = await client.recommend_agents(
            mcp_session,
            task_description,
            max_agents,
            context
        )
        
        if hasattr(mcp_response, 'content') and mcp_response.content:
            try:
                if isinstance(mcp_response.content, list) and mcp_response.content:
                    recommendations = json.loads(mcp_response.content[0].text)
                else:
                    recommendations = mcp_response.content
                
                abi_logging(f"[✅] Found {len(recommendations)} recommendations")
                return recommendations if recommendations else []
            except Exception as e:
                abi_logging(f'[X] Error parsing recommendations: {e}')
                return []
        else:
            abi_logging(f'No recommendations found')
            return []


@tool
async def tool_check_agent_capability(
    agent_name: str,
    required_tasks: List[str]
) -> Dict[str, Any]:
    """Check if an agent has specific capabilities.
    
    Args:
        agent_name: Name of the agent to check
        required_tasks: List of required task names
    
    Returns:
        Capability check result with supported/missing tasks
    """
    async with client.init_session(
        _mcp_config.host,
        _mcp_config.port,
        _mcp_config.transport
    ) as mcp_session:
        abi_logging(f"[🔍] Checking capabilities for: {agent_name}")
        
        context = build_semantic_context_from_card(
            AGENT_CARD_PATH,
            tool_name="check_agent_capability",
            query=agent_name
        )
        
        mcp_response = await client.check_agent_capability(
            mcp_session,
            agent_name,
            required_tasks,
            context
        )
        
        if hasattr(mcp_response, 'content') and mcp_response.content:
            try:
                if isinstance(mcp_response.content, list) and mcp_response.content:
                    result = json.loads(mcp_response.content[0].text)
                else:
                    result = mcp_response.content
                
                abi_logging(f"[✅] Capability check complete")
                return result if result else {}
            except Exception as e:
                abi_logging(f'[X] Error parsing capability check: {e}')
                return {"error": str(e)}
        else:
            abi_logging(f'No capability data found')
            return {}


@tool
async def tool_check_agent_health(agent_name: str) -> Dict[str, Any]:
    """Check if an agent is online and responding.
    
    Args:
        agent_name: Name of the agent to check
    
    Returns:
        Health status with response time and status code
    """
    async with client.init_session(
        _mcp_config.host,
        _mcp_config.port,
        _mcp_config.transport
    ) as mcp_session:
        abi_logging(f"[🏥] Checking health for: {agent_name}")
        
        context = build_semantic_context_from_card(
            AGENT_CARD_PATH,
            tool_name="check_agent_health",
            query=agent_name
        )
        
        mcp_response = await client.check_agent_health(
            mcp_session,
            agent_name,
            context
        )
        
        if hasattr(mcp_response, 'content') and mcp_response.content:
            try:
                if isinstance(mcp_response.content, list) and mcp_response.content:
                    result = json.loads(mcp_response.content[0].text)
                else:
                    result = mcp_response.content
                
                abi_logging(f"[✅] Health check complete: {result.get('status', 'unknown')}")
                return result if result else {}
            except Exception as e:
                abi_logging(f'[X] Error parsing health check: {e}')
                return {"error": str(e)}
        else:
            abi_logging(f'No health data found')
            return {}


@tool
async def tool_register_agent(agent_card_dict: Dict[str, Any]) -> Dict[str, Any]:
    """Register a new agent in the semantic layer.
    
    Args:
        agent_card_dict: Complete agent card dictionary with auth credentials.
                        Must include: id, name, auth (with method, key_id, shared_secret),
                        description, supportedTasks, skills, etc.
    
    Returns:
        Registration result with success status and agent info
    
    Security:
        - Requires HMAC authentication via agent_card.auth
        - Requires authorization via OPA policy
        - Only trusted agents (orchestrator, planner, observer) or agents with
          'register_agents' permission can register new agents
    
    Example:
        agent_card = {
            "id": "agent://new_agent",
            "name": "new_agent",
            "description": "New agent description",
            "auth": {
                "method": "hmac_sha256",
                "key_id": "agent://new_agent-default",
                "shared_secret": "generated_secret_key"
            },
            "supportedTasks": ["task1", "task2"],
            "skills": [...]
        }
        result = await tool_register_agent(agent_card)
    """
    async with client.init_session(
        _mcp_config.host,
        _mcp_config.port,
        _mcp_config.transport
    ) as mcp_session:
        agent_name = agent_card_dict.get('name', 'unknown')
        abi_logging(f"[📝] Registering new agent: {agent_name}")
        
        context = build_semantic_context_from_card(
            AGENT_CARD_PATH,
            tool_name="register_agent",
            query=f"register {agent_name}"
        )
        
        mcp_response = await client.register_agent(
            mcp_session,
            agent_card_dict,
            context
        )
        
        if hasattr(mcp_response, 'content') and mcp_response.content:
            try:
                if isinstance(mcp_response.content, list) and mcp_response.content:
                    result = json.loads(mcp_response.content[0].text)
                else:
                    result = mcp_response.content
                
                if result.get('success'):
                    abi_logging(f"[✅] Agent registered: {result.get('agent_name')}")
                else:
                    abi_logging(f"[❌] Registration failed: {result.get('error')}")
                
                return result if result else {"success": False, "error": "Empty response"}
            except Exception as e:
                abi_logging(f'[X] Error parsing registration result: {e}')
                return {"success": False, "error": str(e)}
        else:
            abi_logging(f'No registration response')
            return {"success": False, "error": "No response from semantic layer"}