"""
Gestionnaire LLM pour le raisonnement visuel avec Ollama.
Interface vers les modèles vision-langage pour la prise de décision.
"""

import json
import base64
from io import BytesIO
from typing import Dict, List, Optional, Any, Tuple
import numpy as np
from PIL import Image

try:
    import ollama
except ImportError:
    ollama = None

from .logger import Logger


class LLMManager:
    """
    Gestionnaire LLM pour le raisonnement visuel utilisant Ollama.
    Supporte les modèles vision-langage comme Qwen 2.5-VL et CogVLM.
    """
    
    def __init__(
        self,
        model_name: str = "qwen2.5-vl:3b",
        ollama_host: str = "localhost:11434",
        logger: Optional[Logger] = None,
        fallback_to_vision: bool = True
    ):
        """
        Initialise le gestionnaire LLM.
        
        Args:
            model_name: Nom du modèle Ollama
            ollama_host: Hôte Ollama
            logger: Instance du logger
            fallback_to_vision: Utiliser la vision pure en cas d'échec LLM
        """
        self.model_name = model_name
        self.ollama_host = ollama_host
        self.logger = logger
        self.fallback_to_vision = fallback_to_vision
        
        # Initialiser le client Ollama
        self._init_client()
    
    def _init_client(self):
        """Initialise le client Ollama."""
        if ollama is None:
            raise ImportError(
                "Ollama n'est pas installé. "
                "Installez-le avec: pip install ollama"
            )
        
        try:
            self.client = ollama.Client(host=self.ollama_host)
            
            # Vérifier que le modèle est disponible
            models = self.client.list()
            model_names = [m.model for m in models.models] if hasattr(models, 'models') else []
            
            if self.model_name not in model_names:
                if self.logger:
                    self.logger.log_action({
                        "action": "model_not_found",
                        "model": self.model_name,
                        "available_models": model_names
                    })
                print(f"Avertissement: Le modèle {self.model_name} n'est pas trouvé.")
                print(f"Modèles disponibles: {model_names}")
            
            if self.logger:
                self.logger.log_action({
                    "action": "llm_client_initialized",
                    "model": self.model_name,
                    "host": self.ollama_host
                })
        
        except Exception as e:
            error_msg = f"Erreur lors de l'initialisation du client Ollama: {e}"
            if self.logger:
                self.logger.log_action({
                    "action": "llm_init_error",
                    "error": str(e)
                })
            if not self.fallback_to_vision:
                raise RuntimeError(error_msg)
            print(f"Avertissement: {error_msg}")
            self.client = None
    
    def _image_to_base64(self, image: np.ndarray) -> str:
        """
        Convertit une image numpy en base64.
        
        Args:
            image: Image numpy array (H, W, C)
            
        Returns:
            String base64 de l'image
        """
        # Convertir BGR vers RGB si nécessaire
        if len(image.shape) == 3 and image.shape[2] == 3:
            image_rgb = image[:, :, ::-1]
        else:
            image_rgb = image
        
        # Convertir en PIL Image
        pil_image = Image.fromarray(image_rgb.astype(np.uint8))
        
        # Convertir en base64
        buffered = BytesIO()
        pil_image.save(buffered, format="PNG")
        img_str = base64.b64encode(buffered.getvalue()).decode()
        
        return img_str
    
    def reason_about_detections(
        self,
        detections: List[Dict[str, Any]],
        context: Dict[str, Any],
        intent: str
    ) -> Dict[str, Any]:
        """
        Utilise le VLM pour sélectionner la meilleure action parmi les détections.
        
        Args:
            detections: Liste de détections avec labels, bbox, images ROI
            context: Contexte actuel (fenêtre, historique, etc.)
            intent: Intention utilisateur
            
        Returns:
            Dictionnaire avec l'élément sélectionné et le score de confiance
        """
        if not detections:
            return {
                "selected_element": None,
                "confidence": 0.0,
                "reasoning": "Aucune détection disponible"
            }
        
        # Fallback si pas de client Ollama
        if self.client is None and self.fallback_to_vision:
            return self._fallback_to_vision_only(detections)
        
        try:
            # Préparer le prompt
            elements_desc = [
                f"- Élément {i+1}: {d['label']} (confiance: {d['confidence']:.2f})"
                for i, d in enumerate(detections)
            ]
            
            prompt = f"""Tu es un assistant d'automatisation RPA. Analyse ces éléments UI détectés et détermine lequel correspond le mieux à l'intention de l'utilisateur.

Intention: {intent}
Contexte: Fenêtre '{context.get('window', 'Inconnue')}'

Éléments détectés:
{chr(10).join(elements_desc)}

Réponds UNIQUEMENT avec un JSON au format suivant:
{{
    "element_index": <index de l'élément (0-{len(detections)-1})>,
    "confidence": <score de confiance 0.0-1.0>,
    "reasoning": "<explication brève>"
}}"""
            
            # Préparer les images
            images = []
            for detection in detections:
                if 'roi_image' in detection and detection['roi_image'] is not None:
                    img_b64 = self._image_to_base64(detection['roi_image'])
                    images.append(img_b64)
            
            # Générer la réponse
            response = self.client.generate(
                model=self.model_name,
                prompt=prompt,
                images=images if images else None,
                stream=False
            )
            
            # Parser la réponse
            result = self._parse_llm_response(response['response'], detections)
            
            if self.logger:
                self.logger.log_action({
                    "action": "llm_reasoning",
                    "intent": intent,
                    "num_detections": len(detections),
                    "selected_index": result.get("element_index"),
                    "confidence": result.get("confidence")
                })
            
            return result
        
        except Exception as e:
            if self.logger:
                self.logger.log_action({
                    "action": "llm_reasoning_error",
                    "error": str(e)
                })
            
            if self.fallback_to_vision:
                return self._fallback_to_vision_only(detections)
            
            return {
                "selected_element": None,
                "confidence": 0.0,
                "reasoning": f"Erreur LLM: {str(e)}"
            }
    
    def _fallback_to_vision_only(
        self,
        detections: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Fallback vers la sélection basée uniquement sur la confiance vision.
        
        Args:
            detections: Liste de détections
            
        Returns:
            Dictionnaire avec l'élément le plus confiant
        """
        if not detections:
            return {
                "selected_element": None,
                "confidence": 0.0,
                "reasoning": "Aucune détection"
            }
        
        # Sélectionner la détection avec la confiance la plus élevée
        best_detection = max(detections, key=lambda d: d.get('confidence', 0.0))
        best_index = detections.index(best_detection)
        
        return {
            "element_index": best_index,
            "selected_element": best_detection,
            "confidence": best_detection.get('confidence', 0.0),
            "reasoning": "Sélection basée sur la confiance vision (fallback)",
            "llm_score": 0.0
        }
    
    def _parse_llm_response(
        self,
        response: str,
        detections: List[Dict[str, Any]]
    ) -> Dict[str, Any]:
        """
        Parse la réponse du LLM.
        
        Args:
            response: Réponse texte du LLM
            detections: Liste des détections originales
            
        Returns:
            Dictionnaire avec l'élément sélectionné et les métadonnées
        """
        try:
            # Extraire le JSON de la réponse
            response_clean = response.strip()
            
            # Chercher le JSON dans la réponse
            start_idx = response_clean.find('{')
            end_idx = response_clean.rfind('}') + 1
            
            if start_idx != -1 and end_idx > start_idx:
                json_str = response_clean[start_idx:end_idx]
                parsed = json.loads(json_str)
                
                element_index = parsed.get('element_index', 0)
                confidence = parsed.get('confidence', 0.5)
                reasoning = parsed.get('reasoning', '')
                
                # Valider l'index
                if 0 <= element_index < len(detections):
                    selected = detections[element_index]
                    
                    return {
                        "element_index": element_index,
                        "selected_element": selected,
                        "confidence": confidence,
                        "reasoning": reasoning,
                        "llm_score": confidence
                    }
            
            # Si le parsing échoue, fallback
            return self._fallback_to_vision_only(detections)
        
        except Exception as e:
            if self.logger:
                self.logger.log_action({
                    "action": "llm_parse_error",
                    "error": str(e),
                    "response": response
                })
            return self._fallback_to_vision_only(detections)
    
    def generate_with_vision(
        self,
        prompt: str,
        images: Optional[List[np.ndarray]] = None
    ) -> str:
        """
        Génération multi-modale avec images.
        
        Args:
            prompt: Prompt texte
            images: Liste d'images numpy arrays
            
        Returns:
            Réponse générée
        """
        if self.client is None:
            return "Erreur: Client Ollama non disponible"
        
        try:
            # Convertir les images en base64
            images_b64 = []
            if images:
                print(f"[LLM] Conversion de {len(images)} images en base64...")
                for i, img in enumerate(images):
                    try:
                        img_b64 = self._image_to_base64(img)
                        images_b64.append(img_b64)
                        print(f"[LLM] Image {i+1}/{len(images)} convertie ({len(img_b64)} bytes)")
                    except Exception as e:
                        print(f"[LLM] Erreur conversion image {i+1}: {e}")
                        raise
            
            # Générer
            print(f"[LLM] Appel Ollama avec modèle {self.model_name}...")
            print(f"[LLM] Prompt: {prompt[:100]}...")
            print(f"[LLM] Images: {len(images_b64) if images_b64 else 0}")
            
            response = self.client.generate(
                model=self.model_name,
                prompt=prompt,
                images=images_b64 if images_b64 else None,
                stream=False,
                options={
                    "temperature": 0.3,  # Basse température pour réponses plus déterministes
                    "num_predict": 20,   # Limiter à 20 tokens (environ 3-4 mots)
                    "top_p": 0.9,
                    "top_k": 40
                }
            )
            
            print(f"[LLM] Réponse brute: {response}")
            
            # Qwen3-VL peut mettre la réponse dans 'thinking' au lieu de 'response'
            result = response.get('response', '')
            
            # Si response est vide, essayer thinking
            if not result and 'thinking' in response:
                thinking = response['thinking']
                print(f"[LLM] Response vide, extraction depuis thinking: '{thinking}'")
                # Nettoyer les balises spéciales de Qwen
                result = thinking.replace('<|im_start|>', '').replace('<|im_end|>', '').replace('<think>', '').replace('</think>', '').strip()
            
            print(f"[LLM] Réponse extraite: '{result}' (longueur: {len(result)})")
            return result
        
        except Exception as e:
            print(f"[LLM] ❌ EXCEPTION: {e}")
            import traceback
            traceback.print_exc()
            if self.logger:
                self.logger.log_action({
                    "action": "generation_error",
                    "error": str(e)
                })
            return f"Erreur de génération: {str(e)}"
    
    def score_action_relevance(
        self,
        action: Dict[str, Any],
        intent: str
    ) -> float:
        """
        Calcule un score de pertinence pour une action donnée.
        
        Args:
            action: Dictionnaire décrivant l'action
            intent: Intention utilisateur
            
        Returns:
            Score de confiance (0.0-1.0)
        """
        if self.client is None:
            # Retourner un score neutre si pas de LLM
            return 0.5
        
        try:
            prompt = f"""Évalue la pertinence de cette action par rapport à l'intention utilisateur.

Intention: {intent}
Action: {action.get('action_type', 'unknown')} sur '{action.get('target_element', 'unknown')}'

Réponds UNIQUEMENT avec un score entre 0.0 et 1.0 (ex: 0.85)"""
            
            response = self.client.generate(
                model=self.model_name,
                prompt=prompt,
                stream=False
            )
            
            # Extraire le score
            response_text = response['response'].strip()
            
            # Chercher un nombre décimal
            import re
            match = re.search(r'0\.\d+|1\.0|0|1', response_text)
            if match:
                score = float(match.group())
                return max(0.0, min(1.0, score))
            
            return 0.5
        
        except Exception as e:
            if self.logger:
                self.logger.log_action({
                    "action": "scoring_error",
                    "error": str(e)
                })
            return 0.5
    
    def is_available(self) -> bool:
        """
        Vérifie si le service LLM est disponible.
        
        Returns:
            True si disponible, False sinon
        """
        if self.client is None:
            return False
        
        try:
            self.client.list()
            return True
        except:
            return False
    
    def get_model_info(self) -> Dict[str, Any]:
        """
        Retourne des informations sur le modèle.
        
        Returns:
            Dictionnaire d'informations
        """
        return {
            "model_name": self.model_name,
            "host": self.ollama_host,
            "available": self.is_available(),
            "fallback_enabled": self.fallback_to_vision
        }