from typing import Optional, Dict, Any, List, Union, Literal, Tuple
from pathlib import Path
import os
from .base_agent import BaseAgent
from .thread import Thread
import importlib
import inspect
from multi_swarm.tools import BaseTool
from multi_swarm.core.config import AgentConfig
from .thread_manager import ThreadManager
from datetime import datetime, UTC, timezone
from pydantic import BaseModel, Field, ConfigDict
import yaml

from .rag import RAGSystem, Document, SearchResult

class MockClient:
    """Mock LLM client for testing."""
    def messages(self):
        return self
        
    def create(self, *args, **kwargs):
        return type('Response', (), {
            'content': [type('Content', (), {'text': 'Mock response'})()]
        })

class Agent(BaseModel):
    """Base class for all agents in the Multi-Swarm framework."""
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    name: str = Field(..., description="Name of the agent")
    description: str = Field(..., description="Description of what the agent does")
    storage_path: Optional[str] = None
    use_rag: bool = Field(default=True)
    use_code_interpreter: bool = Field(default=False)
    use_file_storage: bool = Field(default=False)
    llm_provider: str = Field(default="claude")
    provider_config: Dict[str, Any] = Field(
        default_factory=lambda: {
            "model": "claude-3-sonnet",
            "api_version": "2024-03",
            "max_tokens": 4096
        }
    )
    temperature: float = Field(default=0.7)
    max_prompt_tokens: int = Field(default=25000)
    rag_system: Optional[RAGSystem] = None
    knowledge_stats: Dict[str, Any] = Field(
        default_factory=lambda: {
            "total_documents": 0,
            "last_update": None,
            "quality_metrics": {}
        }
    )

    def __init__(self, **data):
        super().__init__(**data)
        self._initialize_components()
        
    def _initialize_components(self):
        """Initialize agent components based on configuration."""
        if self.use_rag:
            self._initialize_rag()
        if self.use_code_interpreter:
            self._initialize_code_interpreter()
        if self.use_file_storage:
            self._initialize_file_storage()
        self._initialize_llm()
        
    def _initialize_rag(self):
        """Initialize the RAG system."""
        if not self.storage_path:
            self.storage_path = str(Path.home() / ".multi-swarm" / "agents" / self.name)
        
        self.rag_system = RAGSystem(
            storage_path=Path(self.storage_path) / "knowledge"
        )
        
    def _initialize_code_interpreter(self):
        """Initialize the code interpreter."""
        # Code interpreter initialization will be implemented later
        pass
        
    def _initialize_file_storage(self):
        """Initialize file storage."""
        if not self.storage_path:
            self.storage_path = str(Path.home() / ".multi-swarm" / "agents" / self.name)
        os.makedirs(Path(self.storage_path) / "files", exist_ok=True)
        
    def _initialize_llm(self):
        """Initialize the LLM provider."""
        if self.llm_provider == "claude":
            self._initialize_claude()
        elif self.llm_provider == "gemini":
            self._initialize_gemini()
            
    def _initialize_claude(self):
        """Initialize Claude provider."""
        try:
            from anthropic import Anthropic
            self.llm = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
        except ImportError:
            raise ImportError("anthropic package is required for Claude support")
            
    def _initialize_gemini(self):
        """Initialize Gemini provider."""
        try:
            import google.generativeai as genai
            genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
            self.llm = genai.GenerativeModel(
                model_name=self.provider_config["model"]
            )
        except ImportError:
            raise ImportError("google-generativeai package is required for Gemini support")
            
    def add_knowledge(
        self,
        content: str,
        metadata: Optional[Dict] = None
    ) -> Document:
        """Add knowledge to the agent's RAG system."""
        if not self.rag_system:
            raise ValueError("RAG system is not initialized")
            
        if metadata is None:
            metadata = {}
        metadata["agent"] = self.name
        metadata["added_at"] = datetime.now(UTC).isoformat()
        
        document = self.rag_system.add_document(content, metadata)
        self.knowledge_stats["total_documents"] += 1
        self.knowledge_stats["last_update"] = datetime.now(UTC).isoformat()
        return document
        
    def search_knowledge(
        self,
        query: str,
        k: int = 5,
        threshold: Optional[float] = None
    ) -> List[SearchResult]:
        """Search the agent's knowledge base."""
        if not self.rag_system:
            raise ValueError("RAG system is not initialized")
            
        return self.rag_system.search(query, k=k, threshold=threshold)
        
    def search_with_reflection(
        self,
        query: str,
        k: int = 5
    ) -> List[SearchResult]:
        """Search with self-reflection on results."""
        if not self.rag_system:
            raise ValueError("RAG system is not initialized")
            
        return self.rag_system.search_with_reflection(query, k=k)
        
    def process_message(self, thread: Thread) -> str:
        """Process a message in a thread."""
        # Get the last message
        last_message = thread.messages[-1]
        
        # If RAG is enabled, augment with relevant knowledge
        if self.use_rag:
            results = self.search_knowledge(last_message.content, k=3)
            if results:
                context = "\n\n".join([
                    f"Relevant information:\n{r.document.content}"
                    for r in results
                ])
                last_message.content = f"{context}\n\nUser query: {last_message.content}"
        
        # Process with LLM
        response = self._process_with_llm(thread)
        return response
        
    def _process_with_llm(self, thread: Thread) -> str:
        """Process a thread with the LLM provider."""
        if self.llm_provider == "claude":
            return self._process_with_claude(thread)
        elif self.llm_provider == "gemini":
            return self._process_with_gemini(thread)
        else:
            raise ValueError(f"Unsupported LLM provider: {self.llm_provider}")
            
    def _process_with_claude(self, thread: Thread) -> str:
        """Process with Claude."""
        messages = []
        for msg in thread.messages:
            messages.append({
                "role": msg.role,
                "content": msg.content
            })
        
        response = self.llm.messages.create(
            messages=messages,
            model=self.provider_config["model"],
            max_tokens=self.provider_config.get("max_tokens", 4096),
            temperature=self.temperature
        )
        return response.content[0].text
        
    def _process_with_gemini(self, thread: Thread) -> str:
        """Process with Gemini."""
        # Combine messages into a single prompt
        prompt = "\n\n".join([
            f"{msg.role}: {msg.content}"
            for msg in thread.messages
        ])
        
        response = self.llm.generate_content(
            prompt,
            generation_config={
                "temperature": self.temperature,
                "max_output_tokens": self.provider_config.get("max_tokens", 4096)
            }
        )
        return response.text
        
    def get_knowledge_stats(self) -> Dict[str, Any]:
        """Get statistics about the agent's knowledge base."""
        if not self.rag_system:
            raise ValueError("RAG system is not initialized")
            
        self.knowledge_stats["quality_metrics"] = {
            "total_embeddings": len(self.rag_system.documents),
            "average_length": sum(len(d.content) for d in self.rag_system.documents.values()) / len(self.rag_system.documents) if self.rag_system.documents else 0
        }
        
        return self.knowledge_stats

class RAGAgent(Agent):
    """
    Specialized agent for knowledge retrieval and coordination.
    
    This agent:
    1. Manages the RAG system
    2. Coordinates with other agents for knowledge retrieval
    3. Provides self-reflective search capabilities
    4. Maintains knowledge quality and relevance
    """
    model_config = ConfigDict(arbitrary_types_allowed=True)
    
    model_name: str = Field(default=RAGSystem.DEFAULT_MODEL)
    use_rag: Literal[True] = True
    
    def __init__(
        self,
        name: str,
        description: str = "Research and knowledge gathering specialist",
        storage_path: Optional[str] = None,
        model_name: str = RAGSystem.DEFAULT_MODEL
    ):
        """Initialize the RAG agent."""
        super().__init__(
            name=name,
            description=description,
            storage_path=storage_path,
            use_rag=True
        )
        
        # Initialize RAG system
        rag_storage = Path(storage_path) / "rag" if storage_path else None
        self.model_name = model_name
        self.rag_system = RAGSystem(
            model_name=model_name,
            storage_path=rag_storage
        )
    
    def add_knowledge(
        self,
        content: str,
        metadata: Dict = None
    ):
        """Add knowledge to the RAG system."""
        doc = self.rag_system.add_document(content, metadata)
        self.knowledge_stats["total_documents"] += 1
        self.knowledge_stats["last_update"] = datetime.now(UTC)
        return doc
    
    def search_knowledge(
        self,
        query: str,
        k: int = 5,
        threshold: float = None
    ) -> List[SearchResult]:
        """Search the knowledge base using RAG."""
        return self.rag_system.search(query, k, threshold)
    
    def search_with_coordination(
        self,
        query: str,
        collaborating_agents: List[Any],
        k: int = 5
    ) -> Dict[str, Any]:
        """
        Search knowledge with coordination across multiple agents.
        
        Args:
            query: Search query
            collaborating_agents: List of agents to coordinate with
            k: Number of results to return per agent
            
        Returns:
            Dict containing:
            - primary_results: Results from this agent
            - collaborative_results: Results from all agents (including this one)
            - synthesis: Analysis of combined results including:
                - confidence_score
                - knowledge_gaps
                - recommendations
                - total_sources
        """
        # Get primary results from this agent
        primary_results = self.search_knowledge(query, k=k)
        
        # Get collaborative results from each collaborating agent
        collaborative_results = []
        collaborative_results.extend(primary_results)  # Include primary results in collaborative results
        
        for agent in collaborating_agents:
            if agent.use_rag:
                agent_results = agent.search_knowledge(query, k=k)
                collaborative_results.extend(agent_results)
        
        # Combine all results for analysis
        all_results = collaborative_results  # All results are in collaborative_results
        
        # Calculate confidence scores
        confidence_scores = [result.score for result in all_results]
        avg_confidence = sum(confidence_scores) / len(confidence_scores) if confidence_scores else 0
        
        # Check temporal coverage
        dates = []
        for result in all_results:
            if hasattr(result.document, 'created_at'):
                date = result.document.created_at
                if date.tzinfo is None:
                    date = date.replace(tzinfo=timezone.utc)
                dates.append(date)
            if isinstance(result.document.metadata, dict) and "created_at" in result.document.metadata:
                meta_date = result.document.metadata["created_at"]
                if isinstance(meta_date, datetime):
                    if meta_date.tzinfo is None:
                        meta_date = meta_date.replace(tzinfo=timezone.utc)
                    dates.append(meta_date)
        
        # Identify knowledge gaps
        knowledge_gaps = []
        
        # Check temporal gaps
        if dates:
            latest_date = max(dates)
            oldest_date = min(dates)
            current_time = datetime.now(timezone.utc)
            
            # Check if data is too old (older than 6 months)
            if (current_time - latest_date).days > 180:
                knowledge_gaps.append("Temporal gap: Most recent data is over 6 months old")
            
            # Check for large gaps between oldest and newest data
            if (latest_date - oldest_date).days > 365:
                knowledge_gaps.append("Temporal gap: Large time span between oldest and newest data")
        
        # Check coverage gaps based on query terms
        query_terms = set(query.lower().split())
        covered_terms = set()
        for result in all_results:
            covered_terms.update(result.document.content.lower().split())
        missing_terms = query_terms - covered_terms
        if missing_terms:
            knowledge_gaps.append(f"Coverage gap: Missing information about {', '.join(missing_terms)}")
        
        # Check confidence-based gaps
        low_confidence_results = [r for r in all_results if r.score < 0.5]
        if low_confidence_results:
            knowledge_gaps.append("Confidence gap: Some results have low confidence scores")
        
        # Generate recommendations
        recommendations = []
        if knowledge_gaps:
            if any("temporal" in gap.lower() for gap in knowledge_gaps):
                recommendations.append("Consider updating knowledge base with more recent information")
            if any("coverage" in gap.lower() for gap in knowledge_gaps):
                recommendations.append("Consider expanding knowledge base with more comprehensive content")
            if any("confidence" in gap.lower() for gap in knowledge_gaps):
                recommendations.append("Consider adding more authoritative sources")
        
        # Return combined results and analysis
        return {
            "primary_results": primary_results,
            "collaborative_results": collaborative_results,
            "synthesis": {
                "total_sources": len(all_results),
                "confidence_score": avg_confidence,
                "knowledge_gaps": knowledge_gaps,
                "recommendations": recommendations
            }
        }
    
    def _calculate_confidence(self, results: List[SearchResult]) -> float:
        """Calculate confidence score for search results."""
        if not results:
            return 0.0
        
        # Consider factors like:
        # - Document metadata (if confidence is specified)
        # - Search relevance scores
        # - Number of corroborating sources
        scores = []
        for result in results:
            base_score = result.score
            
            # Factor in metadata confidence if available
            confidence_boost = 0.0
            if result.document and result.document.metadata:
                confidence = result.document.metadata.get("confidence", "").lower()
                if confidence == "high":
                    confidence_boost = 0.3
                elif confidence == "medium":
                    confidence_boost = 0.15
                elif confidence == "low":
                    confidence_boost = 0.0
            
            scores.append(base_score + confidence_boost)
        
        return sum(scores) / len(scores)
    
    def _identify_gaps(self, query: str, results: List[SearchResult]) -> List[str]:
        """Identify knowledge gaps in search results."""
        gaps = []
        
        # Check temporal coverage
        dates = []
        for result in results:
            if result.document and result.document.metadata:
                created_at = result.document.metadata.get("created_at")
                if created_at:
                    # Ensure datetime is timezone-aware
                    if not created_at.tzinfo:
                        created_at = created_at.replace(tzinfo=UTC)
                    dates.append(created_at)
        
        if dates:
            latest_date = max(dates)
            current_time = datetime.now(UTC)
            if (current_time - latest_date).days > 180:  # 6 months
                gaps.append("Temporal gap: Most recent information is over 6 months old")
        
        # Check confidence distribution
        confidence_levels = []
        for result in results:
            if result.document and result.document.metadata:
                confidence = result.document.metadata.get("confidence")
                if confidence:
                    confidence_levels.append(confidence)
        
        if confidence_levels:
            high_confidence = sum(1 for c in confidence_levels if c == "high")
            if high_confidence < len(confidence_levels) * 0.3:  # Less than 30% high confidence
                gaps.append("Confidence gap: Limited high-confidence sources")
        
        # Check content coverage
        if len(results) < 3:
            gaps.append("Coverage gap: Limited number of sources")
        
        return gaps
    
    def _generate_recommendations(self, query: str, results: List[SearchResult]) -> List[str]:
        """Generate recommendations based on identified gaps."""
        recommendations = []
        
        # Analyze gaps and suggest improvements
        gaps = self._identify_gaps(query, results)
        for gap in gaps:
            if "temporal" in gap.lower():
                recommendations.append("Consider gathering more recent information")
            elif "confidence" in gap.lower():
                recommendations.append("Validate findings with additional authoritative sources")
            elif "coverage" in gap.lower():
                recommendations.append("Expand search to include more diverse sources")
        
        return recommendations
    
    def search_with_reflection(
        self,
        query: str,
        k: int = 5
    ) -> List[SearchResult]:
        """Search with self-reflection on results."""
        results = self.rag_system.search_with_reflection(query, k)
        return results
    
    def get_knowledge_stats(self) -> Dict:
        """Get statistics about the agent's knowledge base."""
        return {
            **self.knowledge_stats,
            "document_count": len(self.rag_system.documents),
            "storage_path": str(self.rag_system.storage_path) if self.rag_system.storage_path else None
        }

class Agent(BaseAgent):
    """
    Base class for all agents in the Multi-Swarm framework.
    
    This class provides core functionality for agents, including:
    - Thread management for conversations
    - File storage and retrieval
    - Code execution in a secure environment
    - RAG capabilities for knowledge retrieval
    - Tool management and execution
    - Automatic LLM provider selection
    """
    
    # Latest model configurations with task-specific settings
    CLAUDE_CONFIG = {
        "default": {
            "model": "claude-3-5-sonnet-latest",
            "max_tokens": 4096,
            "api_version": "2024-03"
        },
        "code": {
            "model": "claude-3-5-sonnet-latest",
            "max_tokens": 4096,
            "api_version": "2024-03",
            "temperature": 0.1  # Lower temperature for code tasks
        },
        "research": {
            "model": "claude-3-5-opus-latest",  # Use Opus for research
            "max_tokens": 4096,
            "api_version": "2024-03",
            "temperature": 0.7  # Higher temperature for research
        }
    }
    
    # Gemini model configurations
    GEMINI_CONFIG = {
        "default": {
            "model": "gemini-pro",
            "max_tokens": 4096,
            "api_version": "2024-01"
        },
        "world_knowledge": {
            "model": "gemini-2.0-pro-exp-02-05",
            "max_tokens": 8192,
            "api_version": "2024-02",
            "temperature": 0.7
        },
        "complex_reasoning": {
            "model": "gemini-2.0-flash-thinking-exp-01-21",
            "max_tokens": 4096,
            "api_version": "2024-01",
            "temperature": 0.8
        },
        "multimodal": {
            "model": "learnlm-1.5-pro-experimental",
            "max_tokens": 4096,
            "api_version": "2024-02",
            "temperature": 0.6
        }
    }
    
    # Task categories and their preferred providers/models
    TASK_PREFERENCES = {
        # Code Generation & Development
        "code_generation": ("claude", "code"),
        "code_review": ("claude", "code"),
        "debugging": ("claude", "code"),
        "refactoring": ("claude", "code"),
        "api_design": ("claude", "code"),
        
        # Research & Analysis
        "research": ("claude", "research"),
        "planning": ("claude", "research"),
        "architecture": ("claude", "research"),
        "documentation": ("claude", "research"),
        "security": ("claude", "research"),
        
        # Complex Reasoning & Problem Solving
        "complex_reasoning": ("gemini", "complex_reasoning"),
        "problem_solving": ("gemini", "complex_reasoning"),
        "strategy": ("gemini", "complex_reasoning"),
        "decision_making": ("gemini", "complex_reasoning"),
        
        # World Knowledge & Context
        "world_knowledge": ("gemini", "world_knowledge"),
        "context_analysis": ("gemini", "world_knowledge"),
        "long_context": ("gemini", "world_knowledge"),
        
        # Multimodal Processing
        "image_analysis": ("gemini", "multimodal"),
        "audio_processing": ("gemini", "multimodal"),
        "video_analysis": ("gemini", "multimodal"),
        "multimodal_fusion": ("gemini", "multimodal"),
        
        # Data Processing & Analysis
        "data_processing": ("gemini", "data"),
        "statistical_analysis": ("gemini", "data"),
        "data_pipeline": ("gemini", "data"),
        "data_validation": ("gemini", "data"),
        
        # Real-time Operations
        "monitoring": ("gemini", "realtime"),
        "alerting": ("gemini", "realtime"),
        "metrics": ("gemini", "realtime"),
        "logging": ("gemini", "realtime")
    }
    
    # Task categories grouped by provider and capability
    TASK_CATEGORIES = {
        "code": ["code_generation", "code_review", "debugging", "refactoring", "api_design"],
        "research": ["research", "planning", "architecture", "documentation", "security"],
        "complex_reasoning": ["complex_reasoning", "problem_solving", "strategy", "decision_making"],
        "world_knowledge": ["world_knowledge", "context_analysis", "long_context"],
        "multimodal": ["image_analysis", "audio_processing", "video_analysis", "multimodal_fusion"],
        "data": ["data_processing", "statistical_analysis", "data_pipeline", "data_validation"],
        "realtime": ["monitoring", "alerting", "metrics", "logging"]
    }
    
    # Task-related keywords for better matching
    TASK_KEYWORDS = {
        'code': ['programming', 'software', 'development', 'coding', 'debugging', 'refactoring'],
        'research': ['research', 'study', 'investigation', 'planning', 'architecture'],
        'complex_reasoning': ['reasoning', 'problem solving', 'strategy', 'decision making', 'thinking'],
        'world_knowledge': ['knowledge', 'context', 'information', 'understanding', 'long text'],
        'multimodal': ['image', 'audio', 'video', 'visual', 'sound', 'multimedia'],
        'data': ['statistics', 'analysis', 'processing', 'pipeline', 'validation'],
        'realtime': ['monitoring', 'alerts', 'live', 'metrics', 'logging']
    }
    
    def __init__(
        self,
        name: str,
        description: str,
        instructions: str,
        tools_folder: str,
        llm_provider: Optional[str] = None,
        provider_config: Optional[Dict[str, Any]] = None,
        temperature: float = 0.7,
        max_prompt_tokens: int = 25000,
    ):
        """Initialize an agent with the given parameters."""
        self.name = name
        self.description = description
        self.instructions = self._load_instructions(instructions)
        self.tools_folder = tools_folder
        self.tools = self._load_tools()
        self.temperature = temperature
        self.max_prompt_tokens = max_prompt_tokens
        
        # Load model configuration
        self.model_config = self._load_model_config()
        
        # Set provider and config based on configuration
        self.llm_provider = llm_provider or self.model_config.get("default_provider", "claude")
        self.provider_config = provider_config or self._get_default_provider_config()
        
        # Initialize the LLM client
        self.client = self._initialize_llm()
        
    def _load_model_config(self) -> Dict[str, Any]:
        """Load the model configuration from the YAML file."""
        config_path = Path(__file__).parent.parent / "config" / "default_models.yaml"
        try:
            with open(config_path, "r") as f:
                return yaml.safe_load(f)
        except Exception as e:
            print(f"Warning: Could not load model config: {e}")
            return {}
            
    def _get_default_provider_config(self) -> Dict[str, Any]:
        """Get the default configuration for the current provider."""
        if not self.model_config:
            return {}
            
        provider = self.model_config.get("providers", {}).get(self.llm_provider, {})
        default_model = provider.get("default_model")
        
        if not default_model:
            return {}
            
        model_config = provider.get("models", {}).get(default_model, {})
        return {
            "model": model_config.get("name"),
            "api_version": model_config.get("api_version"),
            "max_tokens": model_config.get("max_tokens", 4096),
            "temperature": model_config.get("temperature", 0.7)
        }
        
    def _get_model_for_task(self, task: str) -> Optional[Dict[str, Any]]:
        """Get the appropriate model configuration for a given task."""
        if not self.model_config:
            return None
            
        # Get the provider configuration
        provider = self.model_config.get("providers", {}).get(self.llm_provider, {})
        
        # Find the task category
        task_category = None
        for category, tasks in self.model_config.get("task_categories", {}).items():
            if task in tasks:
                task_category = category
                break
                
        # Get the preferred model for the task
        preferred_model = provider.get("task_preferences", {}).get(task_category)
        if not preferred_model:
            preferred_model = provider.get("default_model")
            
        # Get the model configuration
        if preferred_model:
            model_config = provider.get("models", {}).get(preferred_model, {})
            if model_config:
                return {
                    "model": model_config.get("name"),
                    "api_version": model_config.get("api_version"),
                    "max_tokens": model_config.get("max_tokens", 4096),
                    "temperature": model_config.get("temperature", 0.7)
                }
        return None
        
    def _initialize_llm(self):
        """Initialize the LLM client based on the provider."""
        if self.llm_provider == "claude":
            from multi_swarm.providers.claude_provider import initialize_claude_client
            return initialize_claude_client(self.provider_config)
        elif self.llm_provider == "gemini":
            from multi_swarm.providers.gemini_provider import initialize_gemini_client
            return initialize_gemini_client(self.provider_config)
        else:
            raise ValueError(f"Unsupported LLM provider: {self.llm_provider}")
            
    def _process_with_llm(self, messages: List[Dict[str, str]], task: Optional[str] = None) -> str:
        """Process messages with the LLM, optionally using a task-specific model."""
        try:
            # Get task-specific model configuration if available
            if task:
                task_config = self._get_model_for_task(task)
                if task_config:
                    self.provider_config.update(task_config)
                    
            # Process with the configured model
            if self.llm_provider == "claude":
                response = self.client.messages.create(
                    model=self.provider_config.get("model"),
                    messages=messages,
                    max_tokens=self.provider_config.get("max_tokens", 4096),
                    temperature=self.provider_config.get("temperature", 0.7)
                )
                return response.content[0].text
                
            elif self.llm_provider == "gemini":
                model = self.client.GenerativeModel(
                    model_name=self.provider_config.get("model"),
                    generation_config={
                        "max_output_tokens": self.provider_config.get("max_tokens", 4096),
                        "temperature": self.provider_config.get("temperature", 0.7)
                    }
                )
                response = model.generate_content(messages)
                return response.text
                
        except Exception as e:
            print(f"Error processing with LLM: {e}")
            # Fallback to default model if task-specific model fails
            if task and self.provider_config != self._get_default_provider_config():
                print("Falling back to default model...")
                self.provider_config = self._get_default_provider_config()
                return self._process_with_llm(messages)
            raise

    def _load_instructions(self, instructions: str) -> str:
        """Load instructions from a file."""
        with open(instructions, "r") as f:
            return f.read()

    def _load_tools(self) -> Dict[str, BaseTool]:
        """Load all tools from the tools directory."""
        tools = {}
        if not os.path.exists(self.tools_folder):
            return tools
            
        for file in os.listdir(self.tools_folder):
            if file.endswith('.py') and not file.startswith('__'):
                module_name = file[:-3]
                module_path = os.path.join(self.tools_folder, file)
                spec = importlib.util.spec_from_file_location(module_name, module_path)
                if spec and spec.loader:
                    module = importlib.util.module_from_spec(spec)
                    spec.loader.exec_module(module)
                    for name, obj in inspect.getmembers(module):
                        if (inspect.isclass(obj) and issubclass(obj, BaseTool) 
                            and obj != BaseTool):
                            tool = obj()
                            tools[tool.name or name.lower()] = tool
        return tools

    def _validate_environment(self) -> None:
        """Validate that all required environment variables are set."""
        # First validate the provider
        valid_providers = ["claude", "gemini", "mock"]
        if self.llm_provider not in valid_providers:
            raise ValueError(f"Invalid LLM provider: {self.llm_provider}. Must be one of {valid_providers}")
        
        # Validate API version compatibility
        if self.llm_provider == "gemini":
            api_version = self.provider_config.get("api_version", "")
            if api_version and api_version < "2024-01":
                raise ValueError(f"Incompatible API version: {api_version}. Must be 2024-01 or later for experimental models.")
        
        required_vars = {
            "claude": ["ANTHROPIC_API_KEY"],
            "gemini": ["GOOGLE_API_KEY"],
            "mock": []  # Mock provider doesn't need any environment variables
        }
        
        for var in required_vars[self.llm_provider]:
            if not os.getenv(var):
                raise ValueError(f"Missing required environment variable: {var}")
    
    def _determine_provider_from_description(self, description: str) -> Tuple[str, Dict[str, Any]]:
        """
        Automatically determine the best provider and configuration based on agent description.
        
        Returns:
            Tuple of (provider_name, config_dict)
        """
        description_lower = description.lower()
        
        # First try exact task matching
        for task, (provider, config_name) in self.TASK_PREFERENCES.items():
            task_name = task.replace("_", " ")
            if task_name in description_lower:
                if provider == "claude":
                    return provider, self.CLAUDE_CONFIG[config_name]
                else:  # gemini
                    return provider, self.GEMINI_CONFIG[config_name]
        
        # Try category matching using keywords
        for category, keywords in self.TASK_KEYWORDS.items():
            if any(keyword in description_lower for keyword in keywords):
                # Get the default provider and config for this category
                for task in self.TASK_CATEGORIES[category]:
                    if task in self.TASK_PREFERENCES:
                        provider, config_name = self.TASK_PREFERENCES[task]
                        if provider == "claude":
                            return provider, self.CLAUDE_CONFIG[config_name]
                        else:  # gemini
                            return provider, self.GEMINI_CONFIG[config_name]
        
        # Check for specific keywords that indicate provider preference
        world_knowledge_keywords = ["world", "knowledge", "theory", "explain", "describe"]
        complex_reasoning_keywords = ["analyze", "reason", "think", "solve", "strategy"]
        multimodal_keywords = ["image", "audio", "video", "multimodal", "process"]
        
        if any(keyword in description_lower for keyword in world_knowledge_keywords):
            return "gemini", self.GEMINI_CONFIG["world_knowledge"]
        elif any(keyword in description_lower for keyword in complex_reasoning_keywords):
            return "gemini", self.GEMINI_CONFIG["complex_reasoning"]
        elif any(keyword in description_lower for keyword in multimodal_keywords):
            return "gemini", self.GEMINI_CONFIG["multimodal"]
        
        # Use the specified provider if available, otherwise default to Claude
        return self.llm_provider or "claude", (
            self.GEMINI_CONFIG["default"] if self.llm_provider == "gemini" 
            else self.CLAUDE_CONFIG["default"]
        )
    
    def _get_provider_config(self, config: Optional[Dict[str, Any]] = None, description: Optional[str] = None) -> Dict[str, Any]:
        """Get provider configuration, using defaults if not specified."""
        # If config is provided, use it directly without modification
        if config is not None:
            return config.copy()
            
        # Get the default config first
        if self.llm_provider == "claude":
            provider_config = self.CLAUDE_CONFIG["default"].copy()
        else:  # gemini
            provider_config = self.GEMINI_CONFIG["default"].copy()
            
        # Determine the task category from description
        description_lower = description.lower() if description else ""
        task_category = "default"
        
        # First try to match exact tasks
        for task, (provider, config) in self.TASK_PREFERENCES.items():
            task_name = task.replace("_", " ")
            if task_name in description_lower and provider == self.llm_provider:
                # Find the category this task belongs to
                for category, tasks in self.TASK_CATEGORIES.items():
                    if task in tasks:
                        task_category = category
                        break
                break
        
        # If no exact task match, try category matching
        if task_category == "default":
            for category, keywords in self.TASK_KEYWORDS.items():
                if any(keyword in description_lower for keyword in keywords):
                    if self.TASK_PREFERENCES.get(category):
                        task_category = category
                        break
        
        # Update with task-specific config if available
        if task_category != "default":
            if self.llm_provider == "claude":
                task_config = self.CLAUDE_CONFIG.get(task_category)
            else:  # gemini
                task_config = self.GEMINI_CONFIG.get(task_category)
                
            if task_config:
                provider_config.update(task_config)
            
        return provider_config
        
    def _load_instructions(self, instructions: str) -> str:
        """Load instructions from a file."""
        with open(instructions, "r") as f:
            return f.read()

    def _validate_environment(self) -> None:
        """Validate that all required environment variables are set."""
        # First validate the provider
        valid_providers = ["claude", "gemini", "mock"]
        if self.llm_provider not in valid_providers:
            raise ValueError(f"Invalid LLM provider: {self.llm_provider}. Must be one of {valid_providers}")
        
        # Validate API version compatibility
        if self.llm_provider == "gemini":
            api_version = self.provider_config.get("api_version", "")
            if api_version and api_version < "2024-01":
                raise ValueError(f"Incompatible API version: {api_version}. Must be 2024-01 or later for experimental models.")
        
        required_vars = {
            "claude": ["ANTHROPIC_API_KEY"],
            "gemini": ["GOOGLE_API_KEY"],
            "mock": []  # Mock provider doesn't need any environment variables
        }
        
        for var in required_vars[self.llm_provider]:
            if not os.getenv(var):
                raise ValueError(f"Missing required environment variable: {var}")
    
    def _process_with_llm(self, thread: Thread) -> str:
        """Process a thread with the LLM provider."""
        try:
            # Get the last few messages from the thread
            messages = thread.get_recent_messages(self.max_prompt_tokens)
            
            # Format messages for the provider
            if self.llm_provider == "claude":
                formatted_messages = [
                    {
                        "role": msg.role,
                        "content": msg.content
                    }
                    for msg in messages
                ]
                
                response = self.client.messages.create(
                    model=self.provider_config["model"],
                    max_tokens=self.provider_config.get("max_tokens", 4096),
                    messages=formatted_messages,
                    system=self.instructions
                )
                return response.content[0].text
                
            elif self.llm_provider == "gemini":
                # For Gemini, include instructions in the system message
                formatted_messages = [
                    {
                        "role": "system",
                        "content": self.instructions
                    }
                ]
                formatted_messages.extend([
                    {
                        "role": msg.role,
                        "content": msg.content
                    }
                    for msg in messages
                ])

                # Get the model configuration
                model = self.client.GenerativeModel(self.provider_config["model"])
                
                # Configure the model
                model_config = {
                    "max_output_tokens": self.provider_config.get("max_tokens", 4096),
                    "temperature": self.provider_config.get("temperature", 0.7),
                }

                # Get the last user message
                last_user_message = None
                for msg in reversed(formatted_messages):
                    if msg["role"] == "user":
                        last_user_message = msg["content"]
                        break

                if not last_user_message:
                    raise ValueError("No user message found in thread")

                # Generate response using the model
                response = model.generate_content(
                    f"Instructions: {self.instructions}\n\nUser: {last_user_message}",
                    generation_config=model_config
                )
                
                if response.text:
                    return response.text
                else:
                    # Fallback to standard model if experimental model fails
                    model = self.client.GenerativeModel("gemini-pro")
                    response = model.generate_content(
                        f"Instructions: {self.instructions}\n\nUser: {last_user_message}",
                        generation_config=model_config
                    )
                    return response.text
                
            elif self.llm_provider == "mock":
                return self.client.create().content[0].text
                
            else:
                raise ValueError(f"Unsupported LLM provider: {self.llm_provider}")
                
        except Exception as e:
            # If using an experimental model and it fails, try falling back to the default model
            if (self.llm_provider == "gemini" and 
                any(model in self.provider_config.get("model", "") 
                    for model in ["exp", "experimental"])):
                # Fall back to default model
                self.provider_config = self.GEMINI_CONFIG["default"]
                return self._process_with_llm(thread)
            raise e

    def _initialize_llm(self) -> None:
        """Initialize the LLM client based on the provider."""
        if self.llm_provider == "claude":
            from anthropic import Anthropic
            self.client = Anthropic(api_key=os.getenv("ANTHROPIC_API_KEY"))
        elif self.llm_provider == "gemini":
            import google.generativeai as genai
            genai.configure(api_key=os.getenv("GOOGLE_API_KEY"))
            self.client = genai.GenerativeModel(
                model_name=self.provider_config["model"]
            )
        elif self.llm_provider == "mock":
            self.client = MockClient()
        else:
            raise ValueError(f"Unsupported LLM provider: {self.llm_provider}") 