from .base_llm import BaseLLM
import requests
import base64
import os
from PIL import Image
import io
from datetime import datetime, timedelta
from typing import Dict, Any

class LlamaVision(BaseLLM):
    """
    Classe optimisée pour interagir avec l'API Llama Vision.
    """

    def __init__(self, modele: str = "llama3.2-vision:90b-instruct-q8_0"):
        super().__init__(modele)
        self.api_url = "http://217.182.105.173:11434/api/generate"
        
        # Paramètres optimisés pour réduire les timeouts
        self.params: Dict[str, Any] = {
            "temperature": 0.1,          # Réduit pour des réponses plus prévisibles
            "top_p": 0.8,                # Légèrement réduit pour accélérer
            "top_k": 30,                 # Réduit pour limiter les choix et accélérer
            "num_ctx": 1024,             # Contexte réduit pour des réponses plus rapides
            "repeat_penalty": 1.1,       
            "repeat_last_n": 64,
            "mirostat": 0,               # Désactivé pour accélérer
            "mirostat_eta": 0.1,
            "mirostat_tau": 5,
            "keep_alive": int(timedelta(minutes=2).total_seconds()),  # Réduit le temps de maintien
            "num_predict": 1024,         # Limite la longueur de sortie
            "min_p": 0,
            "seed": 0,
            "stop": ["\n\n", "###"],     # Ajout de tokens d'arrêt pour terminer plus tôt
            "stream": False
        }
        
        # Timeout de requête augmenté pour les images volumineuses
        self.request_timeout = 300  # 5 minutes

    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 par défaut.
        """
        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 spécifique pour Ollama avec le modèle Llama Vision.
        """
        # Ajout d'instructions pour réponses concises
        prompt_prefixe = "Réponds de manière concise et directe à la question suivante: "
        
        contenu = {
            "model": self.modele,
            "prompt": prompt_prefixe + question,
            "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"],
                "keep_alive": self.params["keep_alive"],
                "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 _traiter_reponse(self, reponse: requests.Response) -> str:
        """
        Traite et retourne la réponse fournie par Ollama.
        """
        data = reponse.json()
        return data.get("response", "")

    def _encoder_image_base64(self, image_path: str) -> str:
        """
        Encode une image en base64, avec optimisation de la taille si nécessaire.
        """
        try:
            # Vérifier la taille de l'image et la réduire si trop grande
            with Image.open(image_path) as img:
                # Si l'image est trop grande, la redimensionner
                max_dim = 800  # Dimension maximale
                width, height = img.size
                
                if width > max_dim or height > max_dim:
                    # Calculer le ratio pour conserver les proportions
                    ratio = min(max_dim / width, max_dim / height)
                    new_width = int(width * ratio)
                    new_height = int(height * ratio)
                    
                    # Redimensionner l'image
                    img = img.resize((new_width, new_height), Image.Resampling.LANCZOS)
                
                # Convertir en RGB si nécessaire (pour les formats comme PNG avec canal alpha)
                if img.mode in ("RGBA", "LA", "P"):
                    # Créer un fond blanc et composer l'image dessus pour gérer la transparence
                    background = Image.new("RGB", img.size, (255, 255, 255))
                    if img.mode == "P":
                        img = img.convert("RGBA")
                    background.paste(img, mask=img.split()[3] if img.mode == "RGBA" else None)
                    img = background
                elif img.mode != "RGB":
                    img = img.convert("RGB")
                
                # Sauvegarder temporairement l'image redimensionnée
                buffer = io.BytesIO()
                img.save(buffer, format="JPEG", quality=85)
                buffer.seek(0)
                
                # Encoder en base64
                encoded = base64.b64encode(buffer.read()).decode("utf-8")
                return encoded
            
        except Exception as e:
            print(f"Erreur lors de l'optimisation de l'image: {str(e)}")
            try:
                # Seconde tentative avec une approche plus simple
                with Image.open(image_path) as img:
                    # Convertir directement en RGB quelle que soit l'image
                    img = img.convert("RGB")
                    buffer = io.BytesIO()
                    img.save(buffer, format="JPEG", quality=75)
                    buffer.seek(0)
                    encoded = base64.b64encode(buffer.read()).decode("utf-8")
                    return encoded
            except Exception as e2:
                print(f"Deuxième erreur lors de l'optimisation de l'image: {str(e2)}")
                # Dernier recours: encoder l'image originale sans optimisation
                with open(image_path, "rb") as image_file:
                    encoded = base64.b64encode(image_file.read()).decode("utf-8")
                return encoded

    def _optimiser_prompt(self, question: str) -> str:
        """
        Optimise le prompt pour des réponses plus rapides.
        """
        # Ajouter des instructions pour limiter la longueur et être direct
        optimised_question = f"""Réponds à cette question de façon concise et directe. Limite ta réponse à 3-4 phrases maximum.

Question: {question}"""
        return optimised_question

    def interroger_avec_image(self, image_path: str, question: str) -> str:
        """
        Interroge le modèle Llama Vision avec une image et une question via Ollama.
        Optimisé pour éviter les timeouts.
        
        Args:
            image_path: Chemin vers l'image à analyser
            question: Question ou instructions pour l'analyse
            
        Returns:
            Réponse du modèle à la question
        """
        url = self.urlBase() + self.urlFonction()
        headers = {"Content-Type": "application/json"}

        try:
            # Encoder l'image en base64 avec optimisation de taille
            image_b64 = self._encoder_image_base64(image_path)
            
            # Optimiser la question pour des réponses plus courtes et plus rapides
            optimised_question = self._optimiser_prompt(question)
            
            # Préparer le prompt avec le format spécial pour Ollama multimodal
            prompt = f"""<image>
{image_b64}
</image>

{optimised_question}"""

            contenu = {
                "model": self.modele,
                "prompt": prompt,
                "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"],
                    "keep_alive": self.params["keep_alive"],
                    "num_predict": self.params["num_predict"],
                    "min_p": self.params["min_p"],
                    "seed": self.params["seed"],
                    "stop": self.params["stop"],
                },
                "stream": self.params["stream"]
            }

            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
            else:
                self.dureeTraitement = timedelta(0)

            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 requests.exceptions.Timeout:
            self.heureFin = datetime.now()
            if self.heureDepart is not None:
                self.dureeTraitement = self.heureFin - self.heureDepart
            self.reponseErreur = True
            return "Timeout lors de l'appel à l'API. L'analyse de l'image a pris trop de temps."
            
        except Exception as e:
            self.heureFin = datetime.now()
            if self.heureDepart is not None:
                self.dureeTraitement = self.heureFin - self.heureDepart
            else:
                self.dureeTraitement = timedelta(0)
            self.reponseErreur = True
            return f"Erreur lors de l'analyse de l'image: {str(e)}"
            
    def configurer(self, **kwargs):
        """
        Mise à jour facile des paramètres spécifiques à Llama Vision.
        """
        for key, value in kwargs.items():
            if key in self.params:
                self.params[key] = value
            elif key == "request_timeout" and isinstance(value, int):
                self.request_timeout = value