"""
ScreenAnalyzer - Construction complète d'un ScreenState depuis un screenshot

Orchestre les 4 niveaux du ScreenState :
  Niveau 1 (Raw)       : métadonnées de l'image
  Niveau 2 (Perception): OCR + embedding global
  Niveau 3 (UI)        : détection d'éléments UI
  Niveau 4 (Contexte)  : fenêtre active, workflow en cours

Ce module comble le chaînon manquant entre la capture brute (Couche 0)
et la construction d'embeddings (Couche 3).

=============================================================================
Thread-safety & partage multi-loops (Lot C — avril 2026)
=============================================================================
Cet analyseur peut être partagé entre plusieurs `ExecutionLoop` (singleton
`get_screen_analyzer()`). Pour éviter la contamination croisée :

  • `analyze()` NE MUTE JAMAIS `self._ocr`, `self._ui_detector`,
    `self._ocr_initialized`, `self._ui_detector_initialized` pour gérer les
    flags runtime (enable_ocr / enable_ui_detection). Ces flags sont par
    appel, résolus en variables locales.
  • `session_id` circule en paramètre d'appel et renseigne la metadata du
    ScreenState ; l'attribut `self.session_id` n'est qu'un défaut historique
    (rétrocompat) et n'est plus la source de vérité.
  • L'init lazy des composants lourds (OCR, UIDetector) est protégée par un
    `_init_lock` par instance pour empêcher une double initialisation
    concurrente.
"""

import contextlib
import logging
import os
import threading
import time
from datetime import datetime
from pathlib import Path
from typing import Optional, Dict, Any, List, Tuple

from PIL import Image

from core.models.screen_state import (
    ScreenState,
    RawLevel,
    PerceptionLevel,
    ContextLevel,
    WindowContext,
    EmbeddingRef,
)
from core.models.ui_element import UIElement

logger = logging.getLogger(__name__)


# Lock d'inférence local au module : sert de fallback si le GPUResourceManager
# n'est pas disponible (import error, tests). Partagé entre toutes les instances
# ScreenAnalyzer du process, cohérent avec le singleton get_screen_analyzer().
_ANALYZE_FALLBACK_LOCK = threading.Lock()


def _acquire_gpu_context(timeout: Optional[float] = None):
    """
    Retourne un context manager pour sérialiser les appels GPU.

    Préfère `GPUResourceManager.acquire_inference()` si disponible (coordination
    globale), sinon bascule sur un lock threading local au module.
    """
    try:
        from core.gpu import get_gpu_resource_manager

        manager = get_gpu_resource_manager()
        return manager.acquire_inference(timeout=timeout)
    except Exception as e:  # pragma: no cover - fallback defensif
        logger.debug(f"GPUResourceManager indisponible, fallback lock local: {e}")

        @contextlib.contextmanager
        def _fallback():
            if timeout is None:
                _ANALYZE_FALLBACK_LOCK.acquire()
                yield True
                _ANALYZE_FALLBACK_LOCK.release()
            else:
                got = _ANALYZE_FALLBACK_LOCK.acquire(timeout=timeout)
                try:
                    yield got
                finally:
                    if got:
                        _ANALYZE_FALLBACK_LOCK.release()

        return _fallback()


class ScreenAnalyzer:
    """
    Construit un ScreenState complet (4 niveaux) depuis un screenshot.

    Utilise le UIDetector pour la détection d'éléments et un OCR
    (docTR ou Tesseract) pour l'extraction de texte.

    Example:
        >>> analyzer = ScreenAnalyzer()
        >>> state = analyzer.analyze("/path/to/screenshot.png")
        >>> print(state.perception.detected_text)
        >>> print(len(state.ui_elements))

    Runtime overrides (kwargs-only) sur analyze() :
        >>> state = analyzer.analyze(
        ...     path,
        ...     enable_ocr=False,          # bypass OCR pour cet appel
        ...     enable_ui_detection=False, # bypass UIDetector
        ...     session_id="session_42",   # session par appel
        ... )
    """

    def __init__(
        self,
        ui_detector=None,
        ocr_engine: Optional[str] = None,
        session_id: str = "",
    ):
        """
        Args:
            ui_detector: Instance de UIDetector (créé si None)
            ocr_engine: Moteur OCR à utiliser ("doctr", "tesseract", None=auto)
            session_id: ID de session par défaut (rétrocompat ; préférer passer
                `session_id` en kwarg de `analyze()` pour chaque appel).
        """
        self._ui_detector = ui_detector
        self._ocr_engine_name = ocr_engine
        self._ocr = None
        # Session par défaut (rétrocompat). La source de vérité est désormais
        # le paramètre `session_id` de `analyze()`.
        self.session_id = session_id
        # Compteur d'états — protégé par _state_lock pour être safe en parallèle.
        self._state_counter = 0
        self._state_lock = threading.Lock()

        # Initialisation lazy pour éviter les imports lourds au démarrage.
        self._ui_detector_initialized = ui_detector is not None
        self._ocr_initialized = False

        # Lock dédié à l'init lazy : empêche deux threads d'initialiser
        # simultanément OCR ou UIDetector (double chargement GPU).
        self._init_lock = threading.Lock()

    # =========================================================================
    # API publique
    # =========================================================================

    def analyze(
        self,
        screenshot_path: str,
        window_info: Optional[Dict[str, Any]] = None,
        context: Optional[Dict[str, Any]] = None,
        *,
        enable_ocr: bool = True,
        enable_ui_detection: bool = True,
        session_id: str = "",
    ) -> ScreenState:
        """
        Analyser un screenshot et construire un ScreenState complet.

        Les flags `enable_ocr`, `enable_ui_detection` et `session_id` sont
        **par appel, kwargs-only**, pour ne pas polluer l'état partagé du
        singleton quand plusieurs `ExecutionLoop` se partagent l'analyseur.

        Args:
            screenshot_path: Chemin vers le fichier image
            window_info: Infos fenêtre active {"title": ..., "app_name": ...}
            context: Contexte métier optionnel
            enable_ocr: Active l'OCR pour cet appel (True par défaut).
                False → `detected_text=[]`, aucune init d'OCR déclenchée.
            enable_ui_detection: Active la détection UI pour cet appel
                (True par défaut). False → `ui_elements=[]`.
            session_id: ID de session pour cet appel. Si vide, on retombe sur
                `self.session_id` (rétrocompat). Cette valeur est propagée
                dans `ScreenState.session_id` et `metadata["session_id"]`.

        Returns:
            ScreenState avec les 4 niveaux remplis.
        """
        screenshot_path = str(screenshot_path)

        # Résolution de la session : priorité au kwarg, fallback sur l'état
        # interne (legacy). Variable locale uniquement — pas de mutation.
        effective_session_id = session_id or self.session_id

        # Compteur incrémenté sous lock pour identifiants uniques même en
        # parallèle. C'est la seule mutation tolérée : elle n'impacte pas le
        # comportement OCR/UI.
        with self._state_lock:
            self._state_counter += 1
            state_counter = self._state_counter

        state_id = (
            f"{effective_session_id}_state_{state_counter:04d}"
            if effective_session_id
            else f"state_{state_counter:04d}"
        )

        # Niveau 1 : Raw (léger, hors lock GPU)
        raw = self._build_raw_level(screenshot_path)

        # Résolution locale des instances OCR / UIDetector selon les flags.
        # Aucune mutation de self ici : on décide simplement ce qu'on utilise.
        ocr_instance = self._resolve_ocr_instance(enable_ocr=enable_ocr)
        ui_detector_instance = self._resolve_ui_detector_instance(
            enable_ui_detection=enable_ui_detection
        )

        # Niveaux 2 et 3 : OCR + détection UI sont les étapes lourdes en GPU.
        # On sérialise via GPUResourceManager.acquire_inference() pour éviter
        # que ExecutionLoop et stream_processor saturent simultanément la VRAM
        # sur RTX 5070 (12 Go). Timeout généreux : un appel peut prendre 15-20s.
        with _acquire_gpu_context(timeout=60.0) as acquired:
            if not acquired:
                logger.warning(
                    "Timeout en attendant le lock GPU pour ScreenAnalyzer.analyze() "
                    "→ exécution sans sérialisation (risque saturation VRAM)"
                )

            # Niveau 2 : Perception (OCR) — mesure du temps OCR
            ocr_t0 = time.time()
            detected_text = self._extract_text_with(ocr_instance, screenshot_path)
            ocr_ms = (time.time() - ocr_t0) * 1000

            # Niveau 3 : UI Elements — mesure du temps détection
            ui_t0 = time.time()
            ui_elements = self._detect_ui_elements_with(
                ui_detector_instance, screenshot_path, window_info
            )
            ui_ms = (time.time() - ui_t0) * 1000

        perception = PerceptionLevel(
            embedding=EmbeddingRef(
                provider="openclip_ViT-B-32",
                vector_id=f"data/embeddings/screens/{state_id}.npy",
                dimensions=512,
            ),
            detected_text=detected_text,
            text_detection_method=self._get_ocr_method_name(ocr_instance),
            confidence_avg=0.85 if detected_text else 0.0,
        )

        # Niveau 4 : Contexte
        window_ctx = self._build_window_context(window_info)
        context_level = self._build_context_level(context)

        state = ScreenState(
            screen_state_id=state_id,
            timestamp=datetime.now(),
            session_id=effective_session_id,
            window=window_ctx,
            raw=raw,
            perception=perception,
            context=context_level,
            metadata={
                "analyzer_version": "1.1",
                "session_id": effective_session_id,
                "ui_elements_count": len(ui_elements),
                "text_regions_count": len(detected_text),
                "ocr_ms": ocr_ms,
                "ui_ms": ui_ms,
                "ocr_enabled": enable_ocr,
                "ui_detection_enabled": enable_ui_detection,
            },
            ui_elements=ui_elements,
        )

        logger.info(
            f"ScreenState {state_id} construit: "
            f"{len(ui_elements)} éléments UI, {len(detected_text)} textes détectés "
            f"(ocr={enable_ocr}, ui={enable_ui_detection})"
        )
        return state

    def analyze_image(
        self,
        image: Image.Image,
        save_dir: str = "data/screens",
        window_info: Optional[Dict[str, Any]] = None,
        context: Optional[Dict[str, Any]] = None,
        *,
        enable_ocr: bool = True,
        enable_ui_detection: bool = True,
        session_id: str = "",
    ) -> ScreenState:
        """
        Analyser une PIL Image (utile quand on a déjà l'image en mémoire).

        Sauvegarde l'image sur disque puis appelle analyze(). Les flags
        runtime sont propagés à `analyze()` en kwargs-only.
        """
        save_path = Path(save_dir)
        save_path.mkdir(parents=True, exist_ok=True)

        timestamp = datetime.now().strftime("%Y%m%d_%H%M%S_%f")
        filename = f"screen_{timestamp}.png"
        filepath = save_path / filename

        image.save(str(filepath))
        return self.analyze(
            str(filepath),
            window_info=window_info,
            context=context,
            enable_ocr=enable_ocr,
            enable_ui_detection=enable_ui_detection,
            session_id=session_id,
        )

    # =========================================================================
    # Résolution des instances OCR / UI selon les flags d'appel
    # =========================================================================

    def _resolve_ocr_instance(self, *, enable_ocr: bool):
        """
        Retourner l'instance OCR à utiliser pour cet appel.

        - `enable_ocr=False` → None (pas d'init, pas d'appel OCR)
        - sinon → init lazy sous lock si nécessaire, puis retour de `self._ocr`

        Ne mute `self._ocr` / `self._ocr_initialized` QUE pendant l'init lazy
        réelle, jamais pour bypasser l'OCR d'un appel.
        """
        if not enable_ocr:
            return None
        if not self._ocr_initialized:
            with self._init_lock:
                # Double-check : un autre thread a pu initialiser entretemps.
                if not self._ocr_initialized:
                    self._ensure_ocr_locked()
        return self._ocr

    def _resolve_ui_detector_instance(self, *, enable_ui_detection: bool):
        """
        Retourner l'instance UIDetector pour cet appel (idem _resolve_ocr_instance).
        """
        if not enable_ui_detection:
            return None
        if not self._ui_detector_initialized:
            with self._init_lock:
                if not self._ui_detector_initialized:
                    self._ensure_ui_detector_locked()
        return self._ui_detector

    # =========================================================================
    # Niveau 1 : Raw
    # =========================================================================

    def _build_raw_level(self, screenshot_path: str) -> RawLevel:
        file_size = 0
        try:
            file_size = os.path.getsize(screenshot_path)
        except OSError:
            pass

        return RawLevel(
            screenshot_path=screenshot_path,
            capture_method="mss",
            file_size_bytes=file_size,
        )

    # =========================================================================
    # Niveau 2 : Perception — OCR
    # =========================================================================

    def _extract_text_with(self, ocr_callable, screenshot_path: str) -> List[str]:
        """Extraire le texte via un callable OCR donné (peut être None)."""
        if ocr_callable is None:
            return []
        try:
            return ocr_callable(screenshot_path)
        except Exception as e:
            logger.warning(f"OCR échoué: {e}")
            return []

    def _ensure_ocr_locked(self) -> None:
        """
        Initialiser le moteur OCR (appelé sous `self._init_lock`).

        Ne doit PAS être appelé hors de `_resolve_ocr_instance()`.
        """
        # Mutation intentionnelle : on installe l'instance OCR réelle.
        # Protégée par le lock d'init (pas le lock GPU).
        self._ocr_initialized = True

        engine = self._ocr_engine_name

        # Auto-détection : essayer docTR puis Tesseract
        if engine is None or engine == "doctr":
            try:
                self._ocr = self._create_doctr_ocr()
                logger.info("OCR initialisé avec docTR")
                return
            except Exception as e:
                if engine == "doctr":
                    logger.warning(f"docTR non disponible: {e}")
                    return

        if engine is None or engine == "tesseract":
            try:
                self._ocr = self._create_tesseract_ocr()
                logger.info("OCR initialisé avec Tesseract")
                return
            except Exception as e:
                logger.warning(f"Tesseract non disponible: {e}")

        logger.warning("Aucun moteur OCR disponible — detected_text sera vide")

    def _create_doctr_ocr(self):
        """Créer une fonction OCR basée sur docTR."""
        from doctr.io import DocumentFile
        from doctr.models import ocr_predictor

        predictor = ocr_predictor(det_arch="db_resnet50", reco_arch="crnn_vgg16_bn", pretrained=True)

        def ocr_func(image_path: str) -> List[str]:
            doc = DocumentFile.from_images(image_path)
            result = predictor(doc)
            texts = []
            for page in result.pages:
                for block in page.blocks:
                    for line in block.lines:
                        line_text = " ".join(word.value for word in line.words)
                        if line_text.strip():
                            texts.append(line_text.strip())
            return texts

        return ocr_func

    def _create_tesseract_ocr(self):
        """Créer une fonction OCR basée sur Tesseract."""
        import pytesseract

        def ocr_func(image_path: str) -> List[str]:
            img = Image.open(image_path)
            raw_text = pytesseract.image_to_string(img, lang="fra+eng")
            lines = [line.strip() for line in raw_text.split("\n") if line.strip()]
            return lines

        return ocr_func

    def _get_ocr_method_name(self, ocr_instance=None) -> str:
        """Nom du moteur OCR effectivement utilisé pour cet appel."""
        if ocr_instance is None:
            return "none"
        if self._ocr_engine_name:
            return self._ocr_engine_name
        return "doctr"

    # =========================================================================
    # Niveau 3 : UI Elements
    # =========================================================================

    def _detect_ui_elements_with(
        self,
        ui_detector,
        screenshot_path: str,
        window_info: Optional[Dict[str, Any]] = None,
    ) -> List[UIElement]:
        """Détecter les éléments UI via un détecteur donné (peut être None)."""
        if ui_detector is None:
            return []

        try:
            elements = ui_detector.detect(
                screenshot_path, window_context=window_info
            )
            return elements
        except Exception as e:
            logger.warning(f"Détection UI échouée: {e}")
            return []

    def _ensure_ui_detector_locked(self) -> None:
        """
        Initialiser le UIDetector (appelé sous `self._init_lock`).
        """
        self._ui_detector_initialized = True

        try:
            from core.detection.ui_detector import UIDetector, DetectionConfig

            config = DetectionConfig(
                use_owl_detection=False,  # Désactiver OWL par défaut (lourd)
                use_vlm_classification=True,
                confidence_threshold=0.6,
            )
            self._ui_detector = UIDetector(config)
            logger.info("UIDetector initialisé")
        except Exception as e:
            logger.warning(f"UIDetector non disponible: {e}")
            self._ui_detector = None

    # =========================================================================
    # Niveau 4 : Contexte
    # =========================================================================

    def _build_window_context(
        self, window_info: Optional[Dict[str, Any]] = None
    ) -> WindowContext:
        if window_info:
            return WindowContext(
                app_name=window_info.get("app_name", "unknown"),
                window_title=window_info.get("title", "Unknown"),
                screen_resolution=window_info.get("screen_resolution", [1920, 1080]),
                workspace=window_info.get("workspace", "main"),
                monitor_index=window_info.get("monitor_index", 0),
                dpi_scale=window_info.get("dpi_scale", 100),
                window_bounds=window_info.get("window_bounds"),
                monitors=window_info.get("monitors"),
                os_theme=window_info.get("os_theme", "unknown"),
                os_language=window_info.get("os_language", "unknown"),
            )
        return WindowContext(
            app_name="unknown",
            window_title="Unknown",
            screen_resolution=[1920, 1080],
            workspace="main",
        )

    def _build_context_level(
        self, context: Optional[Dict[str, Any]] = None
    ) -> ContextLevel:
        if context:
            return ContextLevel(
                current_workflow_candidate=context.get("workflow_candidate"),
                workflow_step=context.get("workflow_step"),
                user_id=context.get("user_id", ""),
                tags=context.get("tags", []),
                business_variables=context.get("business_variables", {}),
            )
        return ContextLevel()