from core.base_llm import BaseLLM
import requests
import json
import os
import base64
from typing import Dict, List, Any, Optional
from deep_translator import GoogleTranslator

class LlamaVision90B(BaseLLM):
    def __init__(self):
        model_name = "llama3.2-vision:90b"
        engine = "Ollama"
        
        self.api_url = "http://217.182.105.173:11434/api/chat"
        
        # Paramètres enrichis basés sur l'ancien projet
        default_params = {
            # Paramètres de créativité
            "temperature": 0.3,          # Créativité basse pour analyse technique
            "top_p": 1.0,                # Conserve toute la distribution
            "top_k": 40,                 # Limite vocabulaire
            
            # Paramètres de qualité
            "repeat_penalty": 1.1,       # Réduction des répétitions
            "min_p": 0.0,                # Seuil minimal pour la probabilité des tokens
            
            # Paramètres de contrôle avancé
            "mirostat": 0,               # 0=désactivé, 1=v1, 2=v2
            "mirostat_eta": 0.1,         # Taux d'apprentissage pour mirostat
            "mirostat_tau": 5.0,         # Cible pour mirostat
            
            # Paramètres de taille
            "num_predict": 512,          # Longueur max sortie
            "num_ctx": 4096,             # Contexte étendu
            
            # Paramètres de contrôle
            "seed": 0,                   # Graine pour reproductibilité (0=aléatoire)
            "stop": [],                  # Séquences d'arrêt
            
            # Paramètres de format
            "format": "json",            # Réponse structurée JSON (optionnel)
            "stream": False,             # Réponse d'un seul bloc
            "raw": False,                # Laisse le formatage systèmes
            "keep_alive": "5m"           # Durée de vie de la connexion
        }
        
        super().__init__(model_name=model_name, engine=engine, base_params=default_params)
        
        # Attributs spécifiques pour Llama Vision
        self.image_data = None           # Données de l'image en bytes
        self.image_path = ""             # Chemin vers l'image
        self.json_data = {}              # Données JSON associées
        self.dernier_resultat = {}       # Stockage du dernier résultat d'analyse
        self.historique_resultats = []   # Historique des résultats d'analyse
    
    def set_image(self, image_path: str) -> bool:
        """
        Définit l'image à analyser à partir d'un chemin de fichier
        
        Args:
            image_path: Chemin vers l'image à analyser
            
        Returns:
            True si l'image a été chargée avec succès, False sinon
        """
        if not os.path.exists(image_path):
            return False
        
        try:
            with open(image_path, "rb") as f:
                self.image_data = f.read()
            self.image_path = image_path
            return True
        except Exception as e:
            print(f"Erreur lors du chargement de l'image: {e}")
            return False
    
    def set_image_data(self, image_data: bytes) -> bool:
        """
        Définit directement les données de l'image à analyser
        
        Args:
            image_data: Données binaires de l'image
            
        Returns:
            True si les données ont été définies avec succès, False sinon
        """
        try:
            self.image_data = image_data
            self.image_path = ""
            return True
        except Exception as e:
            print(f"Erreur lors de la définition des données d'image: {e}")
            return False
    
    def set_json_data(self, json_data: Dict[str, Any]) -> bool:
        """
        Définit les données JSON à associer à l'image
        
        Args:
            json_data: Dictionnaire de données à associer
            
        Returns:
            True si les données ont été définies avec succès, False sinon
        """
        try:
            self.json_data = json_data
            return True
        except Exception as e:
            print(f"Erreur lors de la définition des données JSON: {e}")
            return False
        
    def generate(self, user_prompt: str, images: Optional[List[str]] = None, translate: bool = False):
        """
        Génère une réponse à partir d'un prompt utilisateur et d'images
        
        Args:
            user_prompt: Texte du prompt utilisateur
            images: Liste d'images encodées en base64 (optionnel)
            translate: Indique si la réponse doit être traduite
            
        Returns:
            Réponse du modèle, ou tuple (réponse, traduction) si translate=True
        """
        prompt = self._format_prompt(user_prompt)
        
        # Si des images sont fournies directement, utilisez-les
        images_to_use = images if images is not None else []
        
        # Si image_data est défini et aucune image n'est fournie explicitement
        if self.image_data is not None and (images is None or len(images) == 0):
            # Encodage en base64 si ce n'est pas déjà fait
            try:
                if isinstance(self.image_data, bytes):
                    encoded_image = base64.b64encode(self.image_data).decode('utf-8')
                    images_to_use = [encoded_image]
            except Exception as e:
                print(f"Erreur lors de l'encodage de l'image: {e}")
        
        # Ajout des données JSON dans le prompt si disponibles
        if self.json_data:
            json_str = json.dumps(self.json_data, ensure_ascii=False, indent=2)
            # On ajoute le JSON au prompt pour qu'il soit traité avec l'image
            prompt = f"{prompt}\n\nVoici des données JSON associées à cette image:\n```json\n{json_str}\n```"
        
        payload = {
            "model": self.model,
            "messages": [
                {
                    "role": "user",
                    "content": prompt,
                    "images": images_to_use
                }
            ],
            "options": self.params,
            "stream": self.params.get("stream", False)
        }
        
        response = requests.post(self.api_url, json=payload)
        
        if not response.ok:
            raise Exception(f"Erreur API Ollama : {response.status_code} - {response.text}")
        
        result_data = response.json()
        result_text = result_data.get("message", {}).get("content", "")
        
        # Stockage du résultat pour fusion ultérieure
        self.dernier_resultat = result_data
        self.historique_resultats.append(result_data)
        
        self._log_result(user_prompt, result_text)
        
        if translate:
            result_fr = GoogleTranslator(source="auto", target="fr").translate(result_text)
            return result_text, result_fr
        
        return result_text
    
    def fusionner_json_avec_resultats(self) -> Dict[str, Any]:
        """
        Fusionne les données JSON existantes avec les résultats de l'analyse d'image
        
        Returns:
            Dictionnaire combinant le JSON initial et les résultats d'analyse
        """
        if not hasattr(self, 'dernier_resultat') or not self.dernier_resultat:
            return self.json_data
        
        # Créer une copie du JSON original
        resultat_fusionne = self.json_data.copy() if self.json_data else {}
        
        # Ajouter le résultat de l'analyse d'image
        if "analyse_image" not in resultat_fusionne:
            resultat_fusionne["analyse_image"] = []
        
        # Ajouter le résultat à la liste des analyses
        nouvelle_analyse = {
            "modele": self.model,
            "reponse": self.dernier_resultat.get("message", {}).get("content", ""),
            "parametres": {
                "temperature": self.params.get("temperature"),
                "top_p": self.params.get("top_p"),
                "num_ctx": self.params.get("num_ctx")
            }
        }
        
        resultat_fusionne["analyse_image"].append(nouvelle_analyse)
        
        return resultat_fusionne
    
    def sauvegarder_resultats(self, chemin_fichier: str) -> bool:
        """
        Sauvegarde les résultats fusionnés dans un fichier JSON
        
        Args:
            chemin_fichier: Chemin où sauvegarder le fichier de résultats
            
        Returns:
            True si la sauvegarde a réussi, False sinon
        """
        try:
            resultats_fusionnes = self.fusionner_json_avec_resultats()
            
            # Création du dossier parent si nécessaire
            os.makedirs(os.path.dirname(os.path.abspath(chemin_fichier)), exist_ok=True)
            
            with open(chemin_fichier, 'w', encoding='utf-8') as f:
                json.dump(resultats_fusionnes, f, ensure_ascii=False, indent=2)
            
            return True
        except Exception as e:
            print(f"Erreur lors de la sauvegarde des résultats: {e}")
            return False