"""
Module de gestion des modèles Ollama pour systèmes d'agents IA.
Ce module fournit des classes pour interagir avec différents modèles Ollama
et paramétrer leurs comportements.
"""

import json
import requests
from typing import List, Dict, Any, Union, Optional


class OllamaModel:
    """Classe de base pour interagir avec les modèles Ollama."""
    
    def __init__(self, 
                 model_name: str,
                 base_url: str = "http://localhost:11434",
                 temperature: float = 0.7,
                 top_p: float = 0.9,
                 top_k: int = 40,
                 num_ctx: int = 4096,
                 num_predict: int = 128,
                 repeat_penalty: float = 1.1,
                 stream: bool = True):
        """
        Initialise un modèle Ollama.
        
        Args:
            model_name: Nom du modèle Ollama à utiliser
            base_url: URL de base de l'API Ollama
            temperature: Température pour la génération (0.0 à 1.0)
            top_p: Paramètre top_p pour la génération (0.0 à 1.0)
            top_k: Paramètre top_k pour la génération (1 à 100)
            num_ctx: Taille du contexte en tokens
            num_predict: Nombre maximum de tokens à générer
            repeat_penalty: Pénalité pour les répétitions
            stream: Si True, utilise le streaming pour les réponses
        """
        self.model_name = model_name
        self.base_url = base_url
        self.api_url = f"{base_url}/api"
        
        # Paramètres du modèle
        self.options = {
            "temperature": temperature,
            "top_p": top_p,
            "top_k": top_k,
            "num_ctx": num_ctx,
            "num_predict": num_predict,
            "repeat_penalty": repeat_penalty
        }
        
        self.stream = stream
        
    def generate(self, prompt: str) -> Union[str, Dict[str, Any]]:
        """
        Génère une réponse à partir d'un prompt.
        
        Args:
            prompt: Le texte à soumettre au modèle
            
        Returns:
            La réponse du modèle (texte ou dictionnaire complet)
        """
        url = f"{self.api_url}/generate"
        
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": self.stream,
            "options": self.options
        }
        
        try:
            if self.stream:
                return self._stream_response(url, payload)
            else:
                response = requests.post(url, json=payload)
                response.raise_for_status()
                return response.json()
        except Exception as e:
            return f"Erreur lors de la génération: {str(e)}"
    
    def _stream_response(self, url: str, payload: Dict[str, Any]) -> str:
        """
        Gère le streaming de la réponse et renvoie le texte complet.
        
        Args:
            url: URL de l'API
            payload: Données pour la requête
            
        Returns:
            La réponse complète sous forme de texte
        """
        full_response = ""
        try:
            with requests.post(url, json=payload, stream=True) as response:
                response.raise_for_status()
                for line in response.iter_lines():
                    if line:
                        chunk = json.loads(line)
                        if "response" in chunk:
                            print(chunk["response"], end="", flush=True)
                            full_response += chunk["response"]
                        if chunk.get("done", False):
                            break
            print()  # Nouvelle ligne après la réponse
            return full_response
        except Exception as e:
            return f"Erreur lors du streaming: {str(e)}"
    
    def chat(self, messages: List[Dict[str, str]]) -> Union[str, Dict[str, Any]]:
        """
        Utilise l'API chat pour communiquer avec le modèle.
        
        Args:
            messages: Liste de messages au format [{"role": "user", "content": "..."}]
            
        Returns:
            La réponse du modèle
        """
        url = f"{self.api_url}/chat"
        
        payload = {
            "model": self.model_name,
            "messages": messages,
            "stream": self.stream,
            "options": self.options
        }
        
        try:
            if self.stream:
                return self._stream_chat_response(url, payload)
            else:
                response = requests.post(url, json=payload)
                response.raise_for_status()
                return response.json()
        except Exception as e:
            return f"Erreur lors du chat: {str(e)}"
    
    def _stream_chat_response(self, url: str, payload: Dict[str, Any]) -> str:
        """
        Gère le streaming de la réponse de chat et renvoie le texte complet.
        
        Args:
            url: URL de l'API
            payload: Données pour la requête
            
        Returns:
            La réponse complète sous forme de texte
        """
        full_response = ""
        try:
            with requests.post(url, json=payload, stream=True) as response:
                response.raise_for_status()
                for line in response.iter_lines():
                    if line:
                        chunk = json.loads(line)
                        if "message" in chunk and "content" in chunk["message"]:
                            content = chunk["message"]["content"]
                            print(content, end="", flush=True)
                            full_response += content
                        if chunk.get("done", False):
                            break
            print()  # Nouvelle ligne après la réponse
            return full_response
        except Exception as e:
            return f"Erreur lors du streaming du chat: {str(e)}"
    
    def embed(self, input_text: Union[str, List[str]]) -> Dict[str, Any]:
        """
        Génère des embeddings pour un texte.
        
        Args:
            input_text: Texte ou liste de textes pour lesquels générer des embeddings
            
        Returns:
            Réponse contenant les embeddings
        """
        url = f"{self.api_url}/embed"
        
        payload = {
            "model": self.model_name,
            "prompt": input_text
        }
        
        try:
            response = requests.post(url, json=payload)
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}
    
    def update_options(self, **kwargs):
        """
        Met à jour les options du modèle.
        
        Args:
            **kwargs: Paires clé-valeur de paramètres à mettre à jour
        """
        for key, value in kwargs.items():
            if key in self.options:
                self.options[key] = value
    
    @staticmethod
    def list_models() -> Dict[str, Any]:
        """
        Liste tous les modèles disponibles.
        
        Returns:
            Dictionnaire contenant les informations sur les modèles disponibles
        """
        try:
            response = requests.get("http://localhost:11434/api/tags")
            response.raise_for_status()
            return response.json()
        except Exception as e:
            return {"error": str(e)}


class TextGenerationModel(OllamaModel):
    """Classe spécialisée pour les modèles de génération de texte."""
    
    def __init__(self, model_name: str, **kwargs):
        super().__init__(model_name, **kwargs)
    
    def create_completion(self, prompt: str, max_tokens: Optional[int] = None) -> str:
        """
        Crée une complétion de texte simple.
        
        Args:
            prompt: Le texte à compléter
            max_tokens: Nombre maximum de tokens à générer (remplace num_predict si spécifié)
            
        Returns:
            Le texte généré
        """
        # Sauvegarde des options existantes
        original_options = self.options.copy()
        
        if max_tokens is not None:
            self.options["num_predict"] = max_tokens
            
        result = self.generate(prompt)
        
        # Restauration des options d'origine
        self.options = original_options
        
        if isinstance(result, str):
            return result
        elif isinstance(result, dict) and "response" in result:
            return result["response"]
        else:
            return str(result)


class MultiModalModel(OllamaModel):
    """Classe pour les modèles multimodaux prenant en charge les images."""
    
    def __init__(self, model_name: str, **kwargs):
        super().__init__(model_name, **kwargs)
    
    def analyze_image(self, prompt: str, image_path: str) -> str:
        """
        Analyse une image et répond à une question à son sujet.
        
        Args:
            prompt: La question ou instruction concernant l'image
            image_path: Chemin vers le fichier image à analyser
            
        Returns:
            La réponse du modèle à propos de l'image
        """
        import base64
        
        # Lire et encoder l'image en base64
        try:
            with open(image_path, "rb") as image_file:
                image_data = base64.b64encode(image_file.read()).decode("utf-8")
        except Exception as e:
            return f"Erreur lors de la lecture de l'image: {str(e)}"
        
        url = f"{self.api_url}/generate"
        
        payload = {
            "model": self.model_name,
            "prompt": prompt,
            "stream": self.stream,
            "options": self.options,
            "images": [image_data]
        }
        
        try:
            if self.stream:
                return self._stream_response(url, payload)
            else:
                response = requests.post(url, json=payload)
                response.raise_for_status()
                result = response.json()
                return result.get("response", str(result))
        except Exception as e:
            return f"Erreur lors de l'analyse de l'image: {str(e)}"


class EmbeddingModel(OllamaModel):
    """Classe spécialisée pour les modèles d'embedding."""
    
    def __init__(self, model_name: str, **kwargs):
        super().__init__(model_name, **kwargs)
    
    def get_embeddings(self, texts: Union[str, List[str]]) -> List[List[float]]:
        """
        Obtient des embeddings pour un ou plusieurs textes.
        
        Args:
            texts: Texte unique ou liste de textes
            
        Returns:
            Liste des vecteurs d'embedding
        """
        result = self.embed(texts)
        
        if "error" in result:
            print(f"Erreur: {result['error']}")
            return []
        
        if "embedding" in result:
            # Cas d'un seul texte
            return [result["embedding"]]
        elif "embeddings" in result:
            # Cas d'une liste de textes
            return result["embeddings"]
        else:
            print(f"Format de réponse inattendu: {result}")
            return []


class AgentSystem:
    """Système de gestion d'agents basés sur les modèles Ollama."""
    
    def __init__(self):
        self.agents = {}
    
    def add_agent(self, agent_id: str, model: OllamaModel):
        """
        Ajoute un agent au système.
        
        Args:
            agent_id: Identifiant unique pour l'agent
            model: Modèle Ollama à utiliser pour cet agent
        """
        self.agents[agent_id] = model
    
    def get_agent(self, agent_id: str) -> Optional[OllamaModel]:
        """
        Récupère un agent par son ID.
        
        Args:
            agent_id: Identifiant de l'agent
            
        Returns:
            L'instance du modèle associée à l'agent, ou None si non trouvé
        """
        return self.agents.get(agent_id)
    
    def remove_agent(self, agent_id: str) -> bool:
        """
        Supprime un agent du système.
        
        Args:
            agent_id: Identifiant de l'agent à supprimer
            
        Returns:
            True si l'agent a été supprimé, False sinon
        """
        if agent_id in self.agents:
            del self.agents[agent_id]
            return True
        return False
    
    def list_agents(self) -> Dict[str, str]:
        """
        Liste tous les agents disponibles.
        
        Returns:
            Dictionnaire avec ID d'agent comme clé et nom du modèle comme valeur
        """
        return {agent_id: agent.model_name for agent_id, agent in self.agents.items()}


# Fonctions auxiliaires pour créer facilement des agents

def create_text_agent(agent_id: str, model_name: str, system: AgentSystem, **kwargs) -> OllamaModel:
    """
    Crée un agent de génération de texte et l'ajoute au système.
    
    Args:
        agent_id: Identifiant de l'agent
        model_name: Nom du modèle Ollama à utiliser
        system: Système d'agents auquel ajouter l'agent
        **kwargs: Paramètres supplémentaires pour configurer le modèle
        
    Returns:
        L'instance du modèle créée
    """
    agent = TextGenerationModel(model_name, **kwargs)
    system.add_agent(agent_id, agent)
    return agent

def create_multimodal_agent(agent_id: str, model_name: str, system: AgentSystem, **kwargs) -> OllamaModel:
    """
    Crée un agent multimodal et l'ajoute au système.
    
    Args:
        agent_id: Identifiant de l'agent
        model_name: Nom du modèle Ollama à utiliser
        system: Système d'agents auquel ajouter l'agent
        **kwargs: Paramètres supplémentaires pour configurer le modèle
        
    Returns:
        L'instance du modèle créée
    """
    agent = MultiModalModel(model_name, **kwargs)
    system.add_agent(agent_id, agent)
    return agent

def create_embedding_agent(agent_id: str, model_name: str, system: AgentSystem, **kwargs) -> OllamaModel:
    """
    Crée un agent d'embedding et l'ajoute au système.
    
    Args:
        agent_id: Identifiant de l'agent
        model_name: Nom du modèle Ollama à utiliser
        system: Système d'agents auquel ajouter l'agent
        **kwargs: Paramètres supplémentaires pour configurer le modèle
        
    Returns:
        L'instance du modèle créée
    """
    agent = EmbeddingModel(model_name, **kwargs)
    system.add_agent(agent_id, agent)
    return agent 