from typing import Dict, Any, Optional

class AgentConfig:
    """
    Classe pour configurer les paramètres LLM par rôle d'agent.
    Cette classe permet d'harmoniser les paramètres entre différents LLM.
    """
    
    # Configurations par défaut pour chaque rôle d'agent
    DEFAULT_CONFIGS = {
        "json_analyser": {
            "temperature": 0.2,  # Besoin d'analyse précise
            "top_p": 0.9,
            "max_tokens": 1500,
            "system_prompt": "Tu es un assistant spécialisé dans l'analyse de tickets JSON. Extrais les informations pertinentes et structure ta réponse de manière claire.",
        },
        "image_sorter": {
            "temperature": 0.3,  # Décision de classification binaire
            "top_p": 0.9,
            "max_tokens": 200,
            "system_prompt": "Tu es un assistant spécialisé dans le tri d'images. Tu dois décider si une image est pertinente ou non pour BRG_Lab.",
        },
        "image_analyser": {
            "temperature": 0.5,  # Analyse créative mais factuelle
            "top_p": 0.95,
            "max_tokens": 1000,
            "system_prompt": "Tu es un assistant spécialisé dans l'analyse d'images. Décris ce que tu vois en tenant compte du contexte fourni.",
        },
        "report_generator": {
            "temperature": 0.7,  # Génération de rapport plus créative
            "top_p": 1.0,
            "max_tokens": 2000,
            "system_prompt": "Tu es un assistant spécialisé dans la génération de rapports. Synthétise les informations fournies de manière claire et professionnelle.",
        }
    }
    
    # Paramètres spécifiques à chaque type de modèle pour harmoniser les performances
    MODEL_ADJUSTMENTS = {
        "mistral": {},  # Pas d'ajustement nécessaire, modèle de référence
        "pixtral": {
            "temperature": -0.1,  # Légèrement plus conservateur
        },
        "ollama": {
            "temperature": +0.1,  # Légèrement plus créatif
            "num_ctx": 2048,
            "repeat_penalty": 1.1,
        }
    }
    
    def __init__(self, role: str, model_type: str = "mistral"):
        """
        Initialise la configuration avec les paramètres appropriés pour le rôle et le modèle.
        
        Args:
            role: Rôle de l'agent ('json_analyser', 'image_sorter', etc.)
            model_type: Type de modèle ('mistral', 'pixtral', 'ollama', etc.)
        """
        self.role = role
        self.model_type = model_type.lower()
        
        # Récupérer la configuration de base pour ce rôle
        if role in self.DEFAULT_CONFIGS:
            self.config = self.DEFAULT_CONFIGS[role].copy()
        else:
            # Configuration par défaut
            self.config = {
                "temperature": 0.5,
                "top_p": 0.9,
                "max_tokens": 1000,
                "system_prompt": ""
            }
        
        # Appliquer les ajustements spécifiques au modèle
        if model_type.lower() in self.MODEL_ADJUSTMENTS:
            adjustments = self.MODEL_ADJUSTMENTS[model_type.lower()]
            for key, adjustment in adjustments.items():
                if key in self.config:
                    if isinstance(adjustment, (int, float)) and isinstance(self.config[key], (int, float)):
                        self.config[key] += adjustment  # Ajustement relatif
                    else:
                        self.config[key] = adjustment  # Remplacement direct
                else:
                    self.config[key] = adjustment  # Ajout d'un nouveau paramètre
    
    def get_params(self) -> Dict[str, Any]:
        """
        Retourne tous les paramètres sauf le prompt système.
        """
        params = self.config.copy()
        if "system_prompt" in params:
            params.pop("system_prompt")
        return params
    
    def get_system_prompt(self) -> str:
        """
        Retourne le prompt système.
        """
        return self.config.get("system_prompt", "")
    
    def to_dict(self) -> Dict[str, Any]:
        """
        Retourne toute la configuration sous forme de dictionnaire.
        """
        return {
            "role": self.role,
            "model_type": self.model_type,
            "config": self.config
        }
    
    def update(self, **kwargs) -> None:
        """
        Met à jour la configuration avec les nouveaux paramètres.
        """
        self.config.update(kwargs)