src/agents/worker_agents.py

"""
Worker agents for specific tasks in the agentic AI system.
"""

import time
import concurrent.futures
from typing import Dict, Any, List, Optional

from .base_agent import BaseAgent, AgentResult, AgentState
from src.parsers.keyword_extractor import KeywordExtractor
from src.api_clients.api_matcher import APIMatcher
from src.api_clients.cdms_client import CDMSClient
from src.utils.logging_config import logger


class QueryParserAgent(BaseAgent):
    """Agent responsible for parsing natural language queries."""
    
    def _initialize(self) -> None:
        """Initialize the query parser."""
        try:
            self.keyword_extractor = KeywordExtractor(self.config.get('extractor', {}))
            logger.info("Query parser agent initialized")
        except Exception as e:
            logger.error(f"Failed to initialize query parser: {e}")
            raise
    
    def execute(self, input_data: Dict[str, Any]) -> AgentResult:
        """
        Parse a natural language query.
        
        Args:
            input_data: Dictionary containing 'query' key
            
        Returns:
            AgentResult with parsed query data
        """
        start_time = time.time()
        self.set_state(AgentState.RUNNING)
        
        try:
            query = input_data.get('query', '')
            if not query:
                raise ValueError("Query is required")
            
            # Extract keywords and parse query
            parsed_result = self.keyword_extractor.create_enhanced_parsed_query(query)
            
            # Convert to dictionary for serialization
            result_data = {
                'original_query': parsed_result.original_query,
                'keywords': parsed_result.keywords,
                'entities': parsed_result.entities,
                'intent': parsed_result.intent,
                'confidence': parsed_result.confidence,
                'metadata': parsed_result.metadata
            }
            
            execution_time = time.time() - start_time
            self.set_state(AgentState.COMPLETED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=True,
                data=result_data,
                execution_time=execution_time,
                metadata={
                    'keyword_count': len(parsed_result.keywords),
                    'entity_count': len(parsed_result.entities),
                    'extraction_method': parsed_result.metadata.get('extraction_method', 'unknown')
                }
            )
            
            self.log_execution(result)
            return result
            
        except Exception as e:
            execution_time = time.time() - start_time
            self.set_state(AgentState.FAILED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=False,
                data=None,
                error_message=str(e),
                execution_time=execution_time
            )
            
            self.log_execution(result)
            return result
    
    def get_capabilities(self) -> List[str]:
        """Get parser capabilities."""
        return [
            "keyword_extraction",
            "entity_recognition",
            "intent_classification",
            "semantic_analysis",
            "query_preprocessing"
        ]


class APIMatcherAgent(BaseAgent):
    """Agent responsible for matching keywords to API operations."""
    
    def _initialize(self) -> None:
        """Initialize the API matcher."""
        try:
            self.api_matcher = APIMatcher(self.config.get('matcher', {}))
            
            # Load API operations if provided in config
            openapi_spec = self.config.get('openapi_spec')
            if openapi_spec:
                self.api_matcher.load_openapi_spec(openapi_spec)
            
            logger.info("API matcher agent initialized")
        except Exception as e:
            logger.error(f"Failed to initialize API matcher: {e}")
            raise
    
    def execute(self, input_data: Dict[str, Any]) -> AgentResult:
        """
        Match keywords to API operations.
        
        Args:
            input_data: Dictionary containing 'keywords' and optional 'intent'
            
        Returns:
            AgentResult with API matches
        """
        start_time = time.time()
        self.set_state(AgentState.RUNNING)
        
        try:
            keywords = input_data.get('keywords', [])
            if not keywords:
                logger.warning("No keywords provided for API matching")
                keywords = []
            
            # Match keywords to API operations
            matches = self.api_matcher.match_keywords(keywords)
            
            # Convert matches to serializable format
            match_data = []
            for match in matches:
                match_dict = {
                    'keyword': match.keyword,
                    'operation': {
                        'name': match.operation.name,
                        'method': match.operation.method,
                        'endpoint': match.operation.endpoint,
                        'description': match.operation.description,
                        'tags': match.operation.tags
                    },
                    'match_type': match.match_type.value,
                    'similarity_score': match.similarity_score,
                    'confidence': match.confidence,
                    'reasoning': match.reasoning
                }
                match_data.append(match_dict)
            
            execution_time = time.time() - start_time
            self.set_state(AgentState.COMPLETED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=True,
                data={
                    'matches': match_data,
                    'match_count': len(match_data),
                    'keywords_processed': len(keywords)
                },
                execution_time=execution_time,
                metadata={
                    'api_operations_available': len(self.api_matcher.operations),
                    'matching_strategy': 'hybrid'
                }
            )
            
            self.log_execution(result)
            return result
            
        except Exception as e:
            execution_time = time.time() - start_time
            self.set_state(AgentState.FAILED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=False,
                data=None,
                error_message=str(e),
                execution_time=execution_time
            )
            
            self.log_execution(result)
            return result
    
    def load_api_spec(self, openapi_spec: Dict[str, Any]) -> None:
        """Load OpenAPI specification."""
        self.api_matcher.load_openapi_spec(openapi_spec)
        logger.info("Loaded new OpenAPI specification")
    
    def get_capabilities(self) -> List[str]:
        """Get matcher capabilities."""
        return [
            "api_matching",
            "fuzzy_matching",
            "semantic_matching",
            "pattern_matching",
            "openapi_integration"
        ]


class APIExecutorAgent(BaseAgent):
    """Agent responsible for executing API calls."""
    
    def _initialize(self) -> None:
        """Initialize the API executor."""
        try:
            # Initialize API clients
            self.api_clients = {}
            
            # Initialize CDMS client if configured
            cdms_config = self.config.get('cdms', {})
            if cdms_config.get('enabled', False):
                self.api_clients['cdms'] = CDMSClient(
                    api_url=cdms_config.get('api_url'),
                    api_key=cdms_config.get('api_key')
                )
            
            logger.info("API executor agent initialized")
        except Exception as e:
            logger.error(f"Failed to initialize API executor: {e}")
            raise
    
    def execute(self, input_data: Dict[str, Any]) -> AgentResult:
        """
        Execute API calls based on matches.
        
        Args:
            input_data: Dictionary containing 'matches' and optional 'query_context'
            
        Returns:
            AgentResult with API execution results
        """
        start_time = time.time()
        self.set_state(AgentState.RUNNING)
        
        try:
            matches = input_data.get('matches', [])
            query_context = input_data.get('query_context', {})
            
            if not matches:
                logger.info("No API matches to execute")
                execution_time = time.time() - start_time
                self.set_state(AgentState.COMPLETED)
                
                return AgentResult(
                    agent_id=self.agent_id,
                    success=True,
                    data={'results': [], 'executed_count': 0},
                    execution_time=execution_time
                )
            
            # Execute API calls in parallel
            results = []
            executed_count = 0
            top_matches = matches[:5]  # Limit to top 5 matches
            
            with concurrent.futures.ThreadPoolExecutor() as executor:
                future_to_match = {
                    executor.submit(self._execute_single_api, match, query_context): match
                    for match in top_matches
                }
                for future in concurrent.futures.as_completed(future_to_match):
                    match = future_to_match[future]
                    try:
                        api_result = future.result()
                        results.append(api_result)
                        executed_count += 1
                    except Exception as e:
                        logger.error(f"Failed to execute API {match.get('operation', {}).get('name', 'unknown')}: {e}")
                        results.append({
                            'operation': match.get('operation', {}),
                            'success': False,
                            'error': str(e)
                        })
            
            execution_time = time.time() - start_time
            self.set_state(AgentState.COMPLETED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=True,
                data={
                    'results': results,
                    'executed_count': executed_count,
                    'total_matches': len(matches)
                },
                execution_time=execution_time,
                metadata={
                    'api_clients_available': list(self.api_clients.keys())
                }
            )
            
            self.log_execution(result)
            return result
            
        except Exception as e:
            execution_time = time.time() - start_time
            self.set_state(AgentState.FAILED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=False,
                data=None,
                error_message=str(e),
                execution_time=execution_time
            )
            
            self.log_execution(result)
            return result
    
    def _execute_single_api(self, match: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """Execute a single API call."""
        operation = match.get('operation', {})
        operation_name = operation.get('name', 'unknown')
        
        # For now, simulate API execution
        # In a real implementation, this would make actual HTTP requests
        
        logger.info(f"Executing API: {operation_name}")
        
        # Check if this is a CDMS operation
        if 'cdms' in operation_name.lower() or 'label' in operation.get('description', '').lower():
            return self._execute_cdms_operation(operation, context)
        
        # Generic API execution simulation
        return {
            'operation': operation,
            'success': True,
            'data': {
                'message': f"Successfully executed {operation_name}",
                'method': operation.get('method', 'GET'),
                'endpoint': operation.get('endpoint', '/'),
                'simulated': True
            },
            'response_time': 0.1
        }
    
    def _execute_cdms_operation(self, operation: Dict[str, Any], context: Dict[str, Any]) -> Dict[str, Any]:
        """Execute CDMS-specific operation."""
        cdms_client = self.api_clients.get('cdms')
        
        if not cdms_client:
            return {
                'operation': operation,
                'success': False,
                'error': 'CDMS client not configured'
            }
        
        try:
            # Determine operation type based on method and endpoint
            method = operation.get('method', 'GET').upper()
            endpoint = operation.get('endpoint', '')
            
            if method == 'GET' and 'labels' in endpoint:
                # Get labels operation
                response = cdms_client.get_labels(limit=10)
                return {
                    'operation': operation,
                    'success': response.success,
                    'data': response.data if response.success else None,
                    'error': response.error_message if not response.success else None,
                    'response_time': response.response_time
                }
            
            # Default to health check
            healthy = cdms_client.health_check()
            return {
                'operation': operation,
                'success': healthy,
                'data': {'status': 'healthy' if healthy else 'unhealthy'},
                'response_time': 0.1
            }
            
        except Exception as e:
            return {
                'operation': operation,
                'success': False,
                'error': str(e)
            }
    
    def get_capabilities(self) -> List[str]:
        """Get executor capabilities."""
        return [
            "api_execution",
            "http_requests",
            "cdms_integration",
            "error_handling",
            "response_processing"
        ]


class ResultFormatterAgent(BaseAgent):
    """Agent responsible for formatting results for output."""
    
    def _initialize(self) -> None:
        """Initialize the result formatter."""
        self.output_format = self.config.get('output_format', 'structured')
        logger.info("Result formatter agent initialized")
    
    def execute(self, input_data: Dict[str, Any]) -> AgentResult:
        """
        Format results for output.
        
        Args:
            input_data: Dictionary containing query results
            
        Returns:
            AgentResult with formatted output
        """
        start_time = time.time()
        self.set_state(AgentState.RUNNING)
        
        try:
            query = input_data.get('query', '')
            parsed_query = input_data.get('parsed_query', {})
            api_matches = input_data.get('api_matches', {})
            api_results = input_data.get('api_results', {})
            metadata = input_data.get('metadata', {})
            
            # Format the output
            formatted_output = {
                'query': {
                    'original': query,
                    'keywords': parsed_query.get('keywords', []),
                    'intent': parsed_query.get('intent'),
                    'confidence': parsed_query.get('confidence', 0.0)
                },
                'api_matches': {
                    'count': api_matches.get('match_count', 0),
                    'matches': api_matches.get('matches', [])[:3]  # Top 3 matches
                },
                'results': {
                    'executed_count': api_results.get('executed_count', 0),
                    'successful_calls': len([r for r in api_results.get('results', []) if r.get('success', False)]),
                    'data': api_results.get('results', [])
                },
                'summary': self._generate_summary(query, api_results),
                'metadata': {
                    'processing_timestamp': time.time(),
                    'format_version': '1.0',
                    **metadata
                }
            }
            
            execution_time = time.time() - start_time
            self.set_state(AgentState.COMPLETED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=True,
                data=formatted_output,
                execution_time=execution_time,
                metadata={
                    'output_format': self.output_format,
                    'sections_included': list(formatted_output.keys())
                }
            )
            
            self.log_execution(result)
            return result
            
        except Exception as e:
            execution_time = time.time() - start_time
            self.set_state(AgentState.FAILED)
            
            result = AgentResult(
                agent_id=self.agent_id,
                success=False,
                data=None,
                error_message=str(e),
                execution_time=execution_time
            )
            
            self.log_execution(result)
            return result
    
    def _generate_summary(self, query: str, api_results: Dict[str, Any]) -> str:
        """Generate a human-readable summary."""
        executed_count = api_results.get('executed_count', 0)
        results = api_results.get('results', [])
        successful_count = len([r for r in results if r.get('success', False)])
        
        if executed_count == 0:
            return f"No API operations were executed for the query: '{query}'"
        
        if successful_count == executed_count:
            return f"Successfully executed {executed_count} API operation(s) for the query: '{query}'"
        elif successful_count > 0:
            return f"Executed {executed_count} API operation(s) with {successful_count} successful for the query: '{query}'"
        else:
            return f"Executed {executed_count} API operation(s) but none were successful for the query: '{query}'"
    
    def get_capabilities(self) -> List[str]:
        """Get formatter capabilities."""
        return [
            "result_formatting",
            "summary_generation",
            "output_structuring",
            "metadata_enrichment"
        ]