#!/usr/bin/env python3
# -*- coding: utf-8 -*-

"""
Factory pour les agents LLM.
Centralise l'initialisation des agents pour faciliter l'adaptation à différents modèles LLM.
"""

import logging
import time
from typing import Dict, Any, Optional, List, Tuple, Union, Type

# Import des agents disponibles
from agents.mistral_large.agent_ticket_analyser import AgentTicketAnalyser
from agents.pixtral_large.agent_image_sorter import AgentImageSorter
from agents.pixtral_large.agent_image_analyser import AgentImageAnalyser
from agents.mistral_large.agent_report_generator import AgentReportGenerator

# Import des implémentations LLM
from llm_classes.base_llm import BaseLLM
from llm_classes.mistral_large import MistralLarge
from llm_classes.mistral_medium import MistralMedium
from llm_classes.pixtral_large import PixtralLarge
from llm_classes.pixtral_12b import Pixtral12b
from llm_classes.deepseek import DeepSeek
from llm_classes.llama_vision import LlamaVision
from llm_classes.qwen2_5 import Qwen2_5

# Configuration du logging
logger = logging.getLogger("AgentFactory")

# Mapping des noms de modèles aux classes
LLM_CLASSES = {
    "mistral_large": MistralLarge,
    "mistral_medium": MistralMedium,
    "pixtral_large": PixtralLarge,
    "pixtral12b": Pixtral12b,
    "deepseek": DeepSeek,
    "llama_vision": LlamaVision,
    "qwen2_5": Qwen2_5
}

# Catégories de modèles par capacité
TEXT_MODELS = ["mistral_large", "mistral_medium", "deepseek", "qwen2_5"]
VISION_MODELS = ["pixtral_large", "pixtral12b", "llama_vision"]

class AgentFactory:
    """
    Factory pour créer et configurer les agents avec différents modèles LLM.
    """
    
    @staticmethod
    def create_llm(model_name: str, **kwargs) -> BaseLLM:
        """
        Crée une instance d'un modèle LLM spécifique.
        
        Args:
            model_name: Nom du modèle à créer
            **kwargs: Arguments supplémentaires à passer au constructeur
            
        Returns:
            Instance du modèle LLM demandé
            
        Raises:
            ValueError: Si le modèle demandé n'existe pas
        """
        if model_name not in LLM_CLASSES:
            raise ValueError(f"Modèle LLM inconnu: {model_name}. Options disponibles: {', '.join(LLM_CLASSES.keys())}")
        
        logger.info(f"Initialisation du modèle {model_name}")
        start_time = time.time()
        
        model = LLM_CLASSES[model_name](**kwargs)
        
        init_time = time.time() - start_time
        logger.info(f"Modèle {model_name} initialisé en {init_time:.2f} secondes")
        
        return model
    
    @staticmethod
    def get_best_text_model() -> BaseLLM:
        """
        Retourne le meilleur modèle texte disponible.
        
        Returns:
            Instance du meilleur modèle texte
        """
        for model_name in TEXT_MODELS:
            try:
                return AgentFactory.create_llm(model_name)
            except Exception as e:
                logger.warning(f"Impossible d'initialiser {model_name}: {e}")
        
        # Si aucun modèle n'a pu être initialisé
        raise ValueError("Aucun modèle texte n'a pu être initialisé")
    
    @staticmethod
    def get_best_vision_model() -> BaseLLM:
        """
        Retourne le meilleur modèle vision disponible.
        
        Returns:
            Instance du meilleur modèle vision
        """
        for model_name in VISION_MODELS:
            try:
                return AgentFactory.create_llm(model_name)
            except Exception as e:
                logger.warning(f"Impossible d'initialiser {model_name}: {e}")
        
        # Si aucun modèle n'a pu être initialisé
        raise ValueError("Aucun modèle vision n'a pu être initialisé")
    
    @staticmethod
    def create_ticket_analyser(llm: Optional[BaseLLM] = None) -> AgentTicketAnalyser:
        """
        Crée un agent d'analyse de tickets.
        
        Args:
            llm: Instance du modèle LLM à utiliser (si None, utilise le meilleur modèle texte disponible)
            
        Returns:
            Agent d'analyse de tickets
        """
        if llm is None:
            llm = AgentFactory.get_best_text_model()
            
        logger.info("Création de l'agent d'analyse de tickets")
        start_time = time.time()
        
        agent = AgentTicketAnalyser(llm)
        
        init_time = time.time() - start_time
        logger.info(f"Agent d'analyse de tickets créé en {init_time:.2f} secondes")
        
        return agent
    
    @staticmethod
    def create_image_sorter(llm: Optional[BaseLLM] = None) -> AgentImageSorter:
        """
        Crée un agent de tri d'images.
        
        Args:
            llm: Instance du modèle LLM à utiliser (si None, utilise le meilleur modèle vision disponible)
            
        Returns:
            Agent de tri d'images
        """
        if llm is None:
            llm = AgentFactory.get_best_vision_model()
            
        logger.info("Création de l'agent de tri d'images")
        start_time = time.time()
        
        agent = AgentImageSorter(llm)
        
        init_time = time.time() - start_time
        logger.info(f"Agent de tri d'images créé en {init_time:.2f} secondes")
        
        return agent
    
    @staticmethod
    def create_image_analyser(llm: Optional[BaseLLM] = None) -> AgentImageAnalyser:
        """
        Crée un agent d'analyse d'images.
        
        Args:
            llm: Instance du modèle LLM à utiliser (si None, utilise le meilleur modèle vision disponible)
            
        Returns:
            Agent d'analyse d'images
        """
        if llm is None:
            llm = AgentFactory.get_best_vision_model()
            
        logger.info("Création de l'agent d'analyse d'images")
        start_time = time.time()
        
        agent = AgentImageAnalyser(llm)
        
        init_time = time.time() - start_time
        logger.info(f"Agent d'analyse d'images créé en {init_time:.2f} secondes")
        
        return agent
    
    @staticmethod
    def create_report_generator(llm: Optional[BaseLLM] = None) -> AgentReportGenerator:
        """
        Crée un agent de génération de rapport.
        
        Args:
            llm: Instance du modèle LLM à utiliser (si None, utilise le meilleur modèle texte disponible)
            
        Returns:
            Agent de génération de rapport
        """
        if llm is None:
            llm = AgentFactory.get_best_text_model()
            
        logger.info("Création de l'agent de génération de rapport")
        start_time = time.time()
        
        agent = AgentReportGenerator(llm)
        
        init_time = time.time() - start_time
        logger.info(f"Agent de génération de rapport créé en {init_time:.2f} secondes")
        
        return agent
    
    @staticmethod
    def create_all_agents(text_model_name: Optional[str] = None, vision_model_name: Optional[str] = None) -> Dict[str, Any]:
        """
        Crée tous les agents nécessaires avec les modèles spécifiés.
        
        Args:
            text_model_name: Nom du modèle texte à utiliser
            vision_model_name: Nom du modèle vision à utiliser
            
        Returns:
            Dictionnaire contenant tous les agents
        """
        # Initialiser les modèles LLM
        text_llm = None
        if text_model_name:
            text_llm = AgentFactory.create_llm(text_model_name)
        else:
            text_llm = AgentFactory.get_best_text_model()
        
        vision_llm = None
        if vision_model_name:
            vision_llm = AgentFactory.create_llm(vision_model_name)
        else:
            vision_llm = AgentFactory.get_best_vision_model()
        
        # Créer les agents
        agents = {
            "ticket_analyser": AgentFactory.create_ticket_analyser(text_llm),
            "image_sorter": AgentFactory.create_image_sorter(vision_llm),
            "image_analyser": AgentFactory.create_image_analyser(vision_llm),
            "report_generator": AgentFactory.create_report_generator(text_llm)
        }
        
        return agents