"""
Configuration classes for Multi-Swarm components.
"""

from typing import Dict, Any, Optional, List
from pydantic import BaseModel, Field
import yaml
from pathlib import Path

class ModelConfig(BaseModel):
    """Configuration for a specific model."""
    name: str
    provider: str
    api_version: str
    max_tokens: int = 4096
    temperature: float = 0.7
    extra_params: Dict[str, Any] = Field(default_factory=dict)

class ProviderConfig(BaseModel):
    """Configuration for a provider."""
    name: str
    models: Dict[str, ModelConfig]
    default_model: str
    task_preferences: Dict[str, str] = Field(default_factory=dict)

class ModelRegistry(BaseModel):
    """Central registry for all model configurations."""
    providers: Dict[str, ProviderConfig]
    task_categories: Dict[str, List[str]]
    default_provider: str = "claude"
    
    @classmethod
    def from_yaml(cls, path: Path) -> "ModelRegistry":
        """Load configuration from YAML file."""
        with open(path, 'r') as f:
            data = yaml.safe_load(f)
        return cls(**data)
    
    def to_yaml(self, path: Path) -> None:
        """Save configuration to YAML file."""
        with open(path, 'w') as f:
            yaml.dump(self.dict(), f)
    
    def get_model_for_task(self, task: str, provider: Optional[str] = None) -> ModelConfig:
        """Get the appropriate model configuration for a task."""
        if provider:
            provider_config = self.providers[provider]
        else:
            # Find provider based on task preferences
            for p in self.providers.values():
                if task in p.task_preferences:
                    provider_config = p
                    break
            else:
                # Use default provider if no specific preference
                provider_config = self.providers[self.default_provider]
        
        # Get model based on task or default
        model_name = provider_config.task_preferences.get(task, provider_config.default_model)
        return provider_config.models[model_name]

class AgentConfig(BaseModel):
    """Configuration for an agent."""
    name: str = Field(..., description="Name of the agent")
    description: str = Field(..., description="Description of what the agent does")
    instructions_path: str = Field(..., description="Path to the instructions file for the agent")
    tools_folder: str = Field(..., description="Path to the folder containing the agent's tools")
    llm_provider: str = Field(default="claude", description="LLM provider to use")
    provider_config: Dict[str, Any] = Field(default_factory=dict, description="Provider-specific configuration")
    temperature: float = Field(default=0.5, description="Temperature for LLM responses")
    max_prompt_tokens: int = Field(default=25000, description="Maximum number of tokens in the prompt")
    
    class Config:
        """Pydantic model configuration."""
        arbitrary_types_allowed = True 