"""Factory for creating different types of agents."""

import os
import logging
from typing import List, Any, Optional, Dict
from enum import Enum
from langchain.chat_models import init_chat_model

# Import deep research components
from .deep_research_config import DeepResearchConfig
from .deep_research_workflow import DeepResearchWorkflow
from .manual_react_agent import create_manual_react_agent


class AgentType(Enum):
    """Types of agents that can be created."""
    REACT = "react"
    MANUAL_REACT = "react"
    DEEP_RESEARCH = "deep_research"
    REACT_WITH_BACKTRACKING = "react_with_backtracking"
    REACT_WITH_THINK = "react_with_think"
    REACT_WITH_EXPLORE = "react_with_explore"
    REACT_WITH_EXPLORE_REVISIT = "react_with_explore_revisit"


class AgentFactory:
    """Factory for creating evaluation agents."""
    
    @staticmethod
    def create_agent(
        model_config: dict,
        system_prompt: str,
        search_tools: List[Any],
        agent_type: AgentType = AgentType.REACT,
        search_engine: Optional[Any] = None,
        search_checker: Optional[Any] = None,
        search_engines: Optional[List[Any]] = None,
        logger: Optional[logging.Logger] = None,
        dr_config: Optional[DeepResearchConfig] = None,
        agent_stop_type: str = "default",
        agent_stop_kwargs: Optional[Dict[str, Any]] = None
    ) -> Any:
        """Create an agent with specified configuration.
        
        Args:
            model_config: Model configuration dictionary
            system_prompt: System prompt for the agent
            search_tools: List of search tools available to the agent
            agent_type: Type of agent to create
            search_engine: Search engine instance (required for DEEP_RESEARCH)
            search_checker: Search checker instance (for DEEP_RESEARCH)
            search_engines: List of search engines for multi-agent scenarios
            logger: Logger instance
            dr_config: Deep research configuration
            agent_stop_type: Type of stopping condition for React agents
            agent_stop_kwargs: Additional configuration for stopping conditions
        Returns:
            Configured agent instance
        """
        logger = logger or logging.getLogger(__name__)
        
        # Set up API key if needed
        if model_config.get('api_key_env'):
            api_key_env = model_config['api_key_env']
            if api_key_env not in os.environ:
                raise ValueError(f"API key environment variable {api_key_env} not set")
        
        # Initialize the model
        try:            
            # Build model initialization parameters
            model_init_params = {
                'model': model_config['name'],
                'model_provider': model_config['provider'],
                'temperature': model_config.get('temperature', 0.7),
                'max_tokens': model_config.get('max_tokens', 2048)
            }
            
            # Add random seed if specified for reproducible inference
            # Note: Some providers/models may not support seed parameter
            if model_config.get('random_seed') is not None:
                model_init_params['seed'] = model_config['random_seed']
                logger.info(f"Setting random seed for reproducible inference: {model_config['random_seed']}")
            
            model = init_chat_model(**model_init_params)
            logger.info(f"Initialized model: {model_config['name']} from {model_config['provider']}")
        except Exception as e:
            logger.error(f"Failed to initialize model: {e}")
            raise
        
        # Create the agent based on type
        if agent_type == AgentType.REACT:
            return AgentFactory._create_manual_react_agent(model, system_prompt, search_tools, logger, model_config['name'], agent_stop_type, agent_stop_kwargs)
        elif agent_type == AgentType.REACT_WITH_BACKTRACKING:
            # Use manual React agent for backtracking as well
            return AgentFactory._create_manual_react_agent(model, system_prompt, search_tools, logger, model_config['name'], agent_stop_type, agent_stop_kwargs)
        elif agent_type == AgentType.REACT_WITH_THINK:
            return AgentFactory._create_manual_react_agent(model, system_prompt, search_tools, logger, model_config['name'], agent_stop_type, agent_stop_kwargs)
        elif agent_type == AgentType.REACT_WITH_EXPLORE:
            return AgentFactory._create_manual_react_agent(model, system_prompt, search_tools, logger, model_config['name'], agent_stop_type, agent_stop_kwargs)
        elif agent_type == AgentType.REACT_WITH_EXPLORE_REVISIT:
            return AgentFactory._create_manual_react_agent(model, system_prompt, search_tools, logger, model_config['name'], agent_stop_type, agent_stop_kwargs)
        elif agent_type == AgentType.DEEP_RESEARCH:
            return AgentFactory._create_deep_research_agent(model_config, search_engine, search_tools, search_checker, search_engines, logger, dr_config)
        else:
            raise ValueError(f"Unsupported agent type: {agent_type}")
    
    @staticmethod
    def _create_manual_react_agent(
        model: Any, 
        system_prompt: str, 
        search_tools: List[Any], 
        logger: logging.Logger, 
        model_name: str,
        agent_stop_type: str = "default",
        agent_stop_kwargs: Optional[Dict[str, Any]] = None
    ) -> Any:
        """Create a ReAct agent with fine-grained interaction control.
        
        Args:
            model: Language model instance
            system_prompt: System prompt
            search_tools: List of tools
            logger: Logger instance
            model_name: Model name for caching decisions
            agent_stop_type: Type of stopping condition
            agent_stop_kwargs: Additional configuration for stopping conditions
            
        Returns:
            ReAct agent instance
        """
        try:
            # Create the ReAct agent
            agent = create_manual_react_agent(
                model=model,
                system_prompt=system_prompt,
                search_tools=search_tools,
                logger=logger,
                model_name=model_name,
                agent_stop_type=agent_stop_type,
                agent_stop_kwargs=agent_stop_kwargs or {}
            )
            logger.info(f"Created ReAct agent with {len(search_tools)} tools and stop type: {agent_stop_type}")
            return agent
        except Exception as e:
            logger.error(f"Failed to create ReAct agent: {e}")
            raise
    
    @staticmethod
    def _create_deep_research_agent(
        model_config: dict, 
        search_engine: Any, 
        search_tools: List[Any], 
        search_checker: Any, 
        search_engines: Optional[List[Any]], 
        logger: logging.Logger,
        dr_config: Optional[DeepResearchConfig] = None
    ) -> Any:
        """Create a Deep Research agent.
        
        Args:
            model_config: Model configuration dictionary
            search_engine: Search engine instance
            search_tools: List of search tools to use
            search_checker: Search checker instance
            search_engines: List of search engines for multi-agent scenarios
            logger: Logger instance
            
        Returns:
            Deep Research workflow instance
        """
        if search_engine is None:
            raise ValueError("Search engine is required for Deep Research agent")
        
        try:
            # Create deep research configuration
            dr_config = DeepResearchConfig.from_model_config(model_config, dr_config)
            
            # Create deep research workflow
            workflow = DeepResearchWorkflow(
                search_engine=search_engine,
                config=dr_config,
                search_tools=search_tools,
                search_checker=search_checker,
                search_engines=search_engines,
                logger=logger
            )
            
            logger.info(f"Created Deep Research agent with model: {model_config.get('name', 'unknown')}")
            return workflow
            
        except Exception as e:
            logger.error(f"Failed to create Deep Research agent: {e}")
            raise
