modius_performance/data/api_client.py

"""
API client for interacting with the Modius Agent Performance API.
"""
import requests
import logging
from typing import List, Dict, Any, Optional

from ..config.constants import API_BASE_URL, DATA_CONFIG
from .models import (
    AgentTypeResponse, AttributeDefinitionResponse, 
    AgentsListResponse, AttributeValuesResponse
)

logger = logging.getLogger(__name__)


class ModiusAPIClient:
    """Client for interacting with the Modius Agent Performance API."""
    
    def __init__(self, base_url: str = API_BASE_URL):
        self.base_url = base_url
        self.session = requests.Session()
        # Set default timeout and headers
        self.session.timeout = 30
        self.session.headers.update({
            'User-Agent': 'Modius-Agent-Performance/1.0'
        })
    
    def _make_request(self, endpoint: str, params: Optional[Dict[str, Any]] = None) -> Dict[str, Any]:
        """Make a request to the API and return the response."""
        url = f"{self.base_url}{endpoint}"
        logger.debug(f"Making API request to: {url}")
        
        try:
            response = self.session.get(url, params=params)
            logger.debug(f"Response status: {response.status_code}")
            
            if response.status_code == 404:
                logger.warning(f"Resource not found: {url}")
                return {"error": "Not found", "status_code": 404}
            
            response.raise_for_status()
            return {"data": response.json(), "status_code": response.status_code}
            
        except requests.exceptions.RequestException as e:
            logger.error(f"API request error for {url}: {e}")
            return {"error": str(e), "status_code": getattr(e.response, 'status_code', 500)}
    
    def get_agent_type_by_name(self, type_name: str) -> AgentTypeResponse:
        """Get agent type by name."""
        endpoint = f"/api/agent-types/name/{type_name}"
        result = self._make_request(endpoint)
        
        if "error" in result:
            return AgentTypeResponse({}, result["status_code"])
        
        return AgentTypeResponse(result["data"], result["status_code"])
    
    def get_attribute_definition_by_name(self, attr_name: str) -> AttributeDefinitionResponse:
        """Get attribute definition by name."""
        endpoint = f"/api/attributes/name/{attr_name}"
        result = self._make_request(endpoint)
        
        if "error" in result:
            return AttributeDefinitionResponse({}, result["status_code"])
        
        return AttributeDefinitionResponse(result["data"], result["status_code"])
    
    def get_agents_by_type(self, type_id: int) -> AgentsListResponse:
        """Get all agents of a specific type."""
        endpoint = f"/api/agent-types/{type_id}/agents/"
        result = self._make_request(endpoint)
        
        if "error" in result:
            return AgentsListResponse([], result["status_code"])
        
        return AgentsListResponse(result["data"], result["status_code"])
    
    def get_agent_attributes(self, agent_id: int, limit: int = None) -> AttributeValuesResponse:
        """Get attributes for a specific agent."""
        endpoint = f"/api/agents/{agent_id}/attributes/"
        params = {"limit": limit or DATA_CONFIG["api_limit"]}
        result = self._make_request(endpoint, params)
        
        if "error" in result:
            return AttributeValuesResponse([], result["status_code"])
        
        return AttributeValuesResponse(result["data"], result["status_code"])
    
    def get_attribute_values_by_type_and_attr(
        self, 
        agents: List[Dict[str, Any]], 
        attr_def_id: int
    ) -> List[Dict[str, Any]]:
        """Get all attribute values for a specific attribute definition across all agents."""
        all_attributes = []
        logger.debug(f"Getting attributes for {len(agents)} agents with attr_def_id: {attr_def_id}")
        
        for agent in agents:
            agent_id = agent["agent_id"]
            
            # Get agent attributes
            response = self.get_agent_attributes(agent_id)
            
            if not response.is_success():
                logger.error(f"Failed to get attributes for agent ID {agent_id}")
                continue
            
            agent_attrs = response.get_attribute_values()
            logger.debug(f"Agent {agent_id} has {len(agent_attrs)} attributes")
            
            # Filter for the specific attribute definition ID
            filtered_attrs = [
                attr for attr in agent_attrs 
                if attr.get("attr_def_id") == attr_def_id
            ]
            logger.debug(f"Agent {agent_id} has {len(filtered_attrs)} matching attributes")
            
            if filtered_attrs:
                logger.debug(f"Sample attribute for agent {agent_id}: {filtered_attrs[0]}")
            
            all_attributes.extend(filtered_attrs)
        
        logger.info(f"Total matching attributes found across all agents: {len(all_attributes)}")
        return all_attributes
    
    def close(self):
        """Close the session."""
        self.session.close()
    
    def __enter__(self):
        return self
    
    def __exit__(self, exc_type, exc_val, exc_tb):
        self.close()


# Utility functions for backward compatibility
def get_agent_name(agent_id: int, agents: List[Dict[str, Any]]) -> str:
    """Get agent name from agent ID."""
    for agent in agents:
        if agent["agent_id"] == agent_id:
            return agent["agent_name"]
    return "Unknown"


# Global client instance (can be replaced with dependency injection later)
_client_instance = None


def get_api_client() -> ModiusAPIClient:
    """Get the global API client instance."""
    global _client_instance
    if _client_instance is None:
        _client_instance = ModiusAPIClient()
    return _client_instance


def close_api_client():
    """Close the global API client instance."""
    global _client_instance
    if _client_instance is not None:
        _client_instance.close()
        _client_instance = None