"""
SeeClick Adapter pour RPA Vision V3

Intègre le modèle SeeClick pour le visual grounding sur GUI.
SeeClick peut localiser un élément GUI à partir d'une description textuelle
et retourner des coordonnées (x, y) de clic.

Usage:
    adapter = SeeClickAdapter()
    x, y = adapter.ground(screenshot_pil, "bouton Valider")
    # Retourne les coordonnées de clic normalisées ou en pixels
"""

import os
import sys
from typing import Optional, Tuple, Dict, Any
from dataclasses import dataclass
from PIL import Image
import re

# Configuration
SEECLICK_MODEL_ID = "cckevinn/SeeClick"
SEECLICK_LOCAL_PATH = "/home/dom/ai/models/seeclick"  # Cache local si téléchargé


@dataclass
class GroundingResult:
    """Résultat de grounding SeeClick"""
    found: bool
    x: Optional[float]  # Coordonnée X normalisée (0-1) ou pixels
    y: Optional[float]  # Coordonnée Y normalisée (0-1) ou pixels
    x_pixel: Optional[int]  # Coordonnée X en pixels
    y_pixel: Optional[int]  # Coordonnée Y en pixels
    confidence: float
    raw_output: str  # Sortie brute du modèle


class SeeClickAdapter:
    """
    Adapter pour utiliser SeeClick comme grounding fallback.

    SeeClick est un modèle vision-langage basé sur Qwen-VL qui peut
    localiser des éléments GUI à partir de descriptions textuelles.

    Il retourne des coordonnées (x, y) normalisées entre 0 et 1.
    """

    _instance = None
    _initialized = False

    def __new__(cls):
        """Singleton pour éviter de charger le modèle plusieurs fois"""
        if cls._instance is None:
            cls._instance = super().__new__(cls)
        return cls._instance

    def __init__(self):
        """Initialise SeeClick (lazy loading)"""
        if SeeClickAdapter._initialized:
            return

        self.model = None
        self.tokenizer = None
        self.available = False
        self._device = None
        self._check_availability()

    def _check_availability(self):
        """Vérifie si SeeClick peut être chargé"""
        try:
            # Vérifier que transformers est disponible
            from transformers import AutoModelForCausalLM, AutoTokenizer

            # Vérifier si le modèle est en cache local ou accessible
            if os.path.exists(SEECLICK_LOCAL_PATH):
                print(f"✅ [SeeClick] Modèle local trouvé: {SEECLICK_LOCAL_PATH}")
                self.available = True
            else:
                # On suppose qu'on peut télécharger depuis HuggingFace
                print(f"ℹ️ [SeeClick] Modèle sera téléchargé depuis HuggingFace: {SEECLICK_MODEL_ID}")
                self.available = True

            print("✅ [SeeClick] Adapter disponible (chargement différé)")

        except ImportError as e:
            print(f"❌ [SeeClick] Dépendances manquantes: {e}")
            self.available = False
        except Exception as e:
            print(f"❌ [SeeClick] Erreur vérification: {e}")
            self.available = False

    def _load_model(self):
        """Charge le modèle SeeClick (lazy loading)"""
        if self.model is not None:
            return True

        if not self.available:
            return False

        try:
            import torch
            from transformers import AutoModelForCausalLM, AutoTokenizer

            # Déterminer le device
            self._device = 'cuda' if torch.cuda.is_available() else 'cpu'
            print(f"🔄 [SeeClick] Chargement du modèle sur {self._device}...")

            # Charger depuis le cache local ou HuggingFace
            model_path = SEECLICK_LOCAL_PATH if os.path.exists(SEECLICK_LOCAL_PATH) else SEECLICK_MODEL_ID

            self.tokenizer = AutoTokenizer.from_pretrained(
                model_path,
                trust_remote_code=True
            )

            self.model = AutoModelForCausalLM.from_pretrained(
                model_path,
                device_map=self._device,
                trust_remote_code=True,
                torch_dtype=torch.float16 if self._device == 'cuda' else torch.float32
            )

            SeeClickAdapter._initialized = True
            print(f"✅ [SeeClick] Modèle chargé avec succès sur {self._device}")
            return True

        except Exception as e:
            print(f"❌ [SeeClick] Erreur chargement modèle: {e}")
            import traceback
            traceback.print_exc()
            self.available = False
            return False

    def ground(
        self,
        image: Image.Image,
        description: str,
        return_pixels: bool = True
    ) -> GroundingResult:
        """
        Localise un élément GUI basé sur une description textuelle.

        Args:
            image: Image PIL du screenshot
            description: Description de l'élément à trouver
                Ex: "bouton Valider", "icône OnlyOffice", "champ de recherche"
            return_pixels: Si True, retourne les coordonnées en pixels

        Returns:
            GroundingResult avec les coordonnées trouvées
        """
        if not self._load_model():
            return GroundingResult(
                found=False,
                x=None, y=None,
                x_pixel=None, y_pixel=None,
                confidence=0.0,
                raw_output="Model not available"
            )

        try:
            W, H = image.size

            # Formatter le prompt pour SeeClick
            # Format: "In this UI screenshot, what is the position of the element..."
            prompt = f"In this UI screenshot, what is the position of the element corresponding to \"{description}\"?"

            # Préparer l'input (dépend de l'implémentation Qwen-VL)
            # SeeClick utilise le format chat de Qwen-VL
            messages = [
                {
                    "role": "user",
                    "content": [
                        {"type": "image", "image": image},
                        {"type": "text", "text": prompt}
                    ]
                }
            ]

            # Tokeniser et générer
            text = self.tokenizer.apply_chat_template(
                messages,
                tokenize=False,
                add_generation_prompt=True
            )

            inputs = self.tokenizer(text, return_tensors="pt")
            if self._device == 'cuda':
                inputs = {k: v.cuda() for k, v in inputs.items()}

            outputs = self.model.generate(
                **inputs,
                max_new_tokens=100,
                do_sample=False
            )

            # Décoder la sortie
            response = self.tokenizer.decode(outputs[0], skip_special_tokens=True)
            raw_output = response.split("assistant")[-1].strip() if "assistant" in response else response

            print(f"🔍 [SeeClick] Sortie brute: {raw_output}")

            # Parser les coordonnées (format: (0.xxx, 0.yyy) ou [0.xxx, 0.yyy])
            coords = self._parse_coordinates(raw_output)

            if coords:
                x_norm, y_norm = coords
                x_pixel = int(x_norm * W)
                y_pixel = int(y_norm * H)

                print(f"✅ [SeeClick] Trouvé: ({x_pixel}, {y_pixel}) = ({x_norm:.3f}, {y_norm:.3f}) normalisé")

                return GroundingResult(
                    found=True,
                    x=x_pixel if return_pixels else x_norm,
                    y=y_pixel if return_pixels else y_norm,
                    x_pixel=x_pixel,
                    y_pixel=y_pixel,
                    confidence=0.8,  # SeeClick ne donne pas de score explicite
                    raw_output=raw_output
                )
            else:
                print(f"⚠️ [SeeClick] Pas de coordonnées trouvées dans: {raw_output}")
                return GroundingResult(
                    found=False,
                    x=None, y=None,
                    x_pixel=None, y_pixel=None,
                    confidence=0.0,
                    raw_output=raw_output
                )

        except Exception as e:
            print(f"❌ [SeeClick] Erreur grounding: {e}")
            import traceback
            traceback.print_exc()
            return GroundingResult(
                found=False,
                x=None, y=None,
                x_pixel=None, y_pixel=None,
                confidence=0.0,
                raw_output=str(e)
            )

    def _parse_coordinates(self, text: str) -> Optional[Tuple[float, float]]:
        """
        Parse les coordonnées depuis la sortie du modèle.

        Formats supportés:
        - (0.123, 0.456)
        - [0.123, 0.456]
        - click(0.123, 0.456)
        - x: 0.123, y: 0.456
        """
        # Pattern pour (x, y) ou [x, y]
        pattern1 = r'[\(\[]?\s*(\d+\.?\d*)\s*[,;]\s*(\d+\.?\d*)\s*[\)\]]?'

        # Pattern pour click(x, y)
        pattern2 = r'click\s*[\(\[]?\s*(\d+\.?\d*)\s*[,;]\s*(\d+\.?\d*)'

        # Pattern pour x: 0.xxx, y: 0.yyy
        pattern3 = r'x\s*[:=]\s*(\d+\.?\d*).*?y\s*[:=]\s*(\d+\.?\d*)'

        for pattern in [pattern2, pattern1, pattern3]:
            matches = re.findall(pattern, text, re.IGNORECASE)
            if matches:
                x_str, y_str = matches[0]
                try:
                    x = float(x_str)
                    y = float(y_str)

                    # Si les valeurs sont > 1, elles sont probablement en pixels
                    # Dans ce cas, on ne peut pas les normaliser sans connaître la taille
                    if x > 1 or y > 1:
                        # Essayer de normaliser si > 1 (supposer pixels)
                        # Mais on ne sait pas la taille de l'image ici...
                        # On retourne tel quel et on laisse l'appelant gérer
                        pass

                    return (x, y)
                except ValueError:
                    continue

        return None

    def ground_with_anchor(
        self,
        image: Image.Image,
        anchor_image: Image.Image,
        anchor_description: Optional[str] = None
    ) -> GroundingResult:
        """
        Localise un élément similaire à une ancre.

        Si une description est fournie, l'utilise.
        Sinon, essaie de décrire l'ancre automatiquement.

        Args:
            image: Screenshot actuel
            anchor_image: Image de l'ancre à trouver
            anchor_description: Description optionnelle de l'ancre

        Returns:
            GroundingResult avec les coordonnées trouvées
        """
        if anchor_description:
            description = anchor_description
        else:
            # Générer une description générique
            # On pourrait utiliser un captioning model ici
            description = "the element that looks like the given target"

        return self.ground(image, description)


# Instance globale (singleton)
_seeclick_instance: Optional[SeeClickAdapter] = None


def get_seeclick() -> SeeClickAdapter:
    """Retourne l'instance singleton de SeeClick"""
    global _seeclick_instance
    if _seeclick_instance is None:
        _seeclick_instance = SeeClickAdapter()
    return _seeclick_instance


def ground_element(
    image: Image.Image,
    description: str,
    return_pixels: bool = True
) -> GroundingResult:
    """Fonction utilitaire pour localiser un élément"""
    return get_seeclick().ground(image, description, return_pixels)