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"
        
        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
        self.json_data = {}
        
    def set_image(self, image_path: str) -> bool:
        """
        Définit l'image à analyser à partir d'un chemin de fichier
        Retourne True si l'image a été chargée avec succès
        """
        if not os.path.exists(image_path):
            return False
        
        try:
            with open(image_path, "rb") as f:
                self.image_data = f.read()
            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
        """
        try:
            self.image_data = image_data
            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
        """
        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: list = None, translate: bool = False):
        prompt = self._format_prompt(user_prompt)
        
        # Si des images sont fournies directement, utilisez-les
        images_to_use = images if images else []
        
        # Si image_data est défini et aucune image n'est fournie explicitement
        if self.image_data is not None and not images:
            # Encodage en base64 si ce n'est pas déjà fait
            if isinstance(self.image_data, bytes):
                encoded_image = base64.b64encode(self.image_data).decode('utf-8')
                images_to_use = [encoded_image]
        
        # 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", "")
        
        self._log_result(user_prompt, result_text)
        
        # Stockage du résultat pour fusion ultérieure
        self.dernier_resultat = result_data
        
        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
        """
        if not hasattr(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
        """
        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