from .base_llm import BaseLLM
import requests
from datetime import datetime, timedelta
from typing import Dict, Any
import os
import json

class Qwen2_5(BaseLLM):
    """
    Classe pour interagir avec le modèle Qwen 2.5 via Ollama.
    """
    
    def __init__(self):
        """
        Initialise une instance du modèle Qwen 2.5.
        """
        # Initialiser avec le nom correct du modèle
        super().__init__("qwen2-7b")
        
        # Définir les attributs spécifiques
        self.modele = "qwen2-7b"
        self.version = "7B"
        
        # Paramètres optimisés pour Qwen 2.5
        self.params = {
            "temperature": 0.3,
            "top_p": 0.8,
            "top_k": 40,
            "num_ctx": 4096,
            "repeat_penalty": 1.2,
            "repeat_last_n": 128,
            "mirostat": 0,
            "mirostat_eta": 0.1,
            "mirostat_tau": 5,
            "num_predict": 4000,
            "min_p": 0.05,
            "seed": 0,
            "stop": ["</answer>", "###", "\n\n\n"],
            "stream": False
        }
        
        # Timeout de requête
        self.request_timeout = 180  # 3 minutes
        
        # État de la dernière requête
        self.heureDepart = None
        self.heureFin = None
        self.dureeTraitement = timedelta(0)
        self.reponseErreur = False
        
        # Prompt système par défaut
        self.prompt_system = "Tu es un assistant IA expert et précis. Fournis des réponses complètes mais concises."
    
    def urlBase(self) -> str:
        """
        Retourne l'URL de base de l'API Ollama.
        """
        return "http://217.182.105.173:11434/"

    def cleAPI(self) -> str:
        """
        Ollama ne nécessite pas de clé API.
        """
        return ""

    def urlFonction(self) -> str:
        """
        Retourne l'URL spécifique à Ollama pour générer une réponse.
        """
        return "api/generate"
    
    def _preparer_contenu(self, question: str) -> Dict[str, Any]:
        """
        Prépare le contenu de la requête pour Qwen 2.5.
        
        Args:
            question: La question ou instruction à envoyer au modèle
            
        Returns:
            Dictionnaire formaté pour l'API Ollama
        """
        # Optimiser le prompt avec le format spécifique pour Qwen
        prompt_optimise = self._optimiser_prompt(question)
        
        contenu = {
            "model": self.modele,
            "prompt": prompt_optimise,
            "options": {
                "temperature": self.params["temperature"],
                "top_p": self.params["top_p"],
                "top_k": self.params["top_k"],
                "num_ctx": self.params["num_ctx"],
                "repeat_penalty": self.params["repeat_penalty"],
                "repeat_last_n": self.params["repeat_last_n"],
                "mirostat": self.params["mirostat"],
                "mirostat_eta": self.params["mirostat_eta"],
                "mirostat_tau": self.params["mirostat_tau"],
                "num_predict": self.params["num_predict"],
                "min_p": self.params["min_p"],
                "seed": self.params["seed"],
                "stop": self.params["stop"],
            },
            "stream": self.params["stream"]
        }
        return contenu
    
    def _optimiser_prompt(self, question: str) -> str:
        """
        Optimise le format du prompt pour Qwen 2.5.
        
        Args:
            question: La question ou instruction originale
            
        Returns:
            Prompt optimisé pour de meilleures performances
        """
        # Formater avec le format spécifique à Qwen
        formatted_prompt = f"""<system>
{self.prompt_system}
</system>

<human>
{question}
</human>

<answer>
"""
        return formatted_prompt
    
    def _traiter_reponse(self, reponse: requests.Response) -> str:
        """
        Traite la réponse fournie par l'API.
        
        Args:
            reponse: Réponse HTTP de l'API
            
        Returns:
            Texte de la réponse
        """
        try:
            data = reponse.json()
            response_text = data.get("response", "")
            
            # Nettoyer la réponse des tags spécifiques à Qwen si présents
            response_text = response_text.replace("</answer>", "").strip()
            
            # Retirer les parties répétitives potentielles à la fin
            if "<human>" in response_text:
                response_text = response_text.split("<human>")[0].strip()
            
            return response_text
        except Exception as e:
            self.reponseErreur = True
            return f"Erreur de traitement de la réponse: {str(e)}"
    
    def interroger(self, question: str) -> str:
        """
        Interroge le modèle Qwen 2.5.
        
        Args:
            question: Question ou instruction à transmettre au modèle
            
        Returns:
            Réponse du modèle
        """
        url = self.urlBase() + self.urlFonction()
        headers = {"Content-Type": "application/json"}
        contenu = self._preparer_contenu(question)
        
        try:
            self.heureDepart = datetime.now()
            response = requests.post(url=url, headers=headers, json=contenu, timeout=self.request_timeout)
            self.heureFin = datetime.now()
            
            if self.heureDepart is not None:
                self.dureeTraitement = self.heureFin - self.heureDepart
                
            if response.status_code in [200, 201]:
                self.reponseErreur = False
                return self._traiter_reponse(response)
            else:
                self.reponseErreur = True
                return f"Erreur API ({response.status_code}): {response.text}"
                
        except Exception as e:
            self.heureFin = datetime.now()
            if self.heureDepart is not None:
                self.dureeTraitement = self.heureFin - self.heureDepart
            self.reponseErreur = True
            return f"Erreur lors de l'interrogation: {str(e)}"
    
    def interroger_avec_image(self, image_path: str, question: str) -> str:
        """
        Cette méthode est présente pour maintenir la compatibilité avec l'interface BaseLLM,
        mais le modèle Qwen2-7b ne prend pas en charge l'analyse d'images.
        
        Args:
            image_path: Chemin vers l'image (non utilisé)
            question: Question concernant l'image
            
        Returns:
            Message d'erreur indiquant que l'analyse d'image n'est pas prise en charge
        """
        self.reponseErreur = True
        return "Le modèle Qwen2-7b ne prend pas en charge l'analyse d'images."
    
    def configurer(self, **kwargs):
        """
        Configure les paramètres du modèle.
        
        Args:
            **kwargs: Paramètres à configurer (temperature, top_p, etc.)
        """
        # Appliquer les paramètres
        for key, value in kwargs.items():
            if key in self.params:
                self.params[key] = value
            elif key == "prompt_system" and isinstance(value, str):
                self.prompt_system = value
            elif key == "request_timeout" and isinstance(value, int):
                self.request_timeout = value
        
        return self