"""
Screen Capture Module - Capture d'écran continue pour RPA Vision V3

Fonctionnalités:
- Capture unique ou continue
- Buffer circulaire pour historique
- Détection de changement d'écran
- Support multi-moniteur
- Optimisation mémoire
"""
import numpy as np
from typing import Optional, Dict, List, Callable, Tuple
from dataclasses import dataclass, field
from datetime import datetime
from pathlib import Path
import threading
import time
import logging
import hashlib
from PIL import Image

logger = logging.getLogger(__name__)


@dataclass
class CaptureFrame:
    """Un frame capturé avec métadonnées"""
    image: np.ndarray
    timestamp: datetime
    frame_id: int
    hash: str
    window_info: Optional[Dict] = None
    changed_from_previous: bool = True


@dataclass
class CaptureStats:
    """Statistiques de capture"""
    total_captures: int = 0
    captures_per_second: float = 0.0
    unchanged_frames_skipped: int = 0
    average_capture_time_ms: float = 0.0
    buffer_size: int = 0
    memory_usage_mb: float = 0.0


class ScreenCapturer:
    """
    Capturer d'écran avancé avec mode continu.
    
    Modes:
    - Single: Capture unique à la demande
    - Continuous: Capture en boucle avec callback
    - Buffered: Maintient un historique des N derniers frames
    
    Example:
        >>> capturer = ScreenCapturer(buffer_size=10)
        >>> # Capture unique
        >>> frame = capturer.capture()
        >>> # Mode continu
        >>> capturer.start_continuous(callback=on_frame, interval_ms=500)
        >>> # ... plus tard ...
        >>> capturer.stop_continuous()
    """
    
    def __init__(
        self,
        buffer_size: int = 10,
        detect_changes: bool = True,
        change_threshold: float = 0.02,
        monitor_index: int = 1
    ):
        """
        Initialiser le capturer.
        
        Args:
            buffer_size: Nombre de frames à garder en mémoire
            detect_changes: Détecter si l'écran a changé
            change_threshold: Seuil de changement (0-1)
            monitor_index: Index du moniteur (1=principal)
        """
        self.buffer_size = buffer_size
        self.detect_changes = detect_changes
        self.change_threshold = change_threshold
        self.monitor_index = monitor_index
        
        # Buffer circulaire
        self._buffer: List[CaptureFrame] = []
        self._frame_counter = 0
        self._last_hash: Optional[str] = None
        
        # Mode continu
        self._continuous_running = False
        self._continuous_thread: Optional[threading.Thread] = None
        self._continuous_callback: Optional[Callable[[CaptureFrame], None]] = None
        self._continuous_interval_ms = 500
        self._lock = threading.Lock()
        
        # Stats
        self._stats = CaptureStats()
        self._capture_times: List[float] = []
        
        # Initialiser le backend de capture
        self._init_capture_backend()
        
        logger.info(f"ScreenCapturer initialized (buffer={buffer_size}, changes={detect_changes})")
    
    def _init_capture_backend(self) -> None:
        """Initialiser le backend de capture (mss ou pyautogui)."""
        # Ne plus garder self.sct - créer MSS à chaque capture (Option A - ultra stable)
        self.sct = None
        self.pyautogui = None
        self.method = None
        
        try:
            import mss
            # Test que mss fonctionne sans garder l'instance
            with mss.mss() as test_sct:
                pass
            self.method = "mss"
            logger.info("Using mss for screen capture (thread-safe mode)")
        except ImportError:
            try:
                import pyautogui
                self.pyautogui = pyautogui
                self.method = "pyautogui"
                logger.info("Using pyautogui for screen capture")
            except ImportError:
                raise ImportError("Neither mss nor pyautogui available for screen capture")
    
    # =========================================================================
    # Capture unique
    # =========================================================================
    
    def capture(self) -> Optional[np.ndarray]:
        """
        Capture unique de l'écran.
        
        Returns:
            Screenshot as numpy array (H, W, 3) RGB ou None si erreur
        """
        try:
            start_time = time.time()
            
            if self.method == "mss":
                img = self._capture_mss()
            else:
                img = self._capture_pyautogui()
            
            # Stats
            capture_time = (time.time() - start_time) * 1000
            self._capture_times.append(capture_time)
            if len(self._capture_times) > 100:
                self._capture_times.pop(0)
            self._stats.total_captures += 1
            self._stats.average_capture_time_ms = sum(self._capture_times) / len(self._capture_times)
            
            return img
            
        except Exception as e:
            logger.error(f"Capture failed: {e}")
            return None
    
    def capture_frame(self) -> Optional[CaptureFrame]:
        """
        Capture avec métadonnées complètes.
        
        Returns:
            CaptureFrame avec image, timestamp, hash, etc.
        """
        img = self.capture()
        return self._create_frame(img)
    
    def _capture_frame_threaded(self, thread_sct) -> Optional[CaptureFrame]:
        """
        Capture avec instance mss thread-local.
        
        Args:
            thread_sct: Instance mss créée dans le thread
        
        Returns:
            CaptureFrame ou None
        """
        try:
            start_time = time.time()
            
            if self.method == "mss" and thread_sct:
                monitor_idx = self.monitor_index if len(thread_sct.monitors) > self.monitor_index else 0
                monitor = thread_sct.monitors[monitor_idx]
                sct_img = thread_sct.grab(monitor)
                img = np.array(sct_img)
                img = img[:, :, :3][:, :, ::-1]  # BGRA to RGB
            else:
                img = self._capture_pyautogui()
            
            # Stats
            capture_time = (time.time() - start_time) * 1000
            self._capture_times.append(capture_time)
            if len(self._capture_times) > 100:
                self._capture_times.pop(0)
            self._stats.total_captures += 1
            self._stats.average_capture_time_ms = sum(self._capture_times) / len(self._capture_times)
            
            return self._create_frame(img)
            
        except Exception as e:
            logger.error(f"Threaded capture failed: {e}")
            return None
    
    def _create_frame(self, img: Optional[np.ndarray]) -> Optional[CaptureFrame]:
        """Créer un CaptureFrame à partir d'une image."""
        if img is None:
            return None
        
        # Calculer le hash pour détecter les changements
        img_hash = self._compute_hash(img)
        changed = True
        
        if self.detect_changes and self._last_hash:
            changed = img_hash != self._last_hash
            if not changed:
                self._stats.unchanged_frames_skipped += 1
        
        self._last_hash = img_hash
        self._frame_counter += 1
        
        frame = CaptureFrame(
            image=img,
            timestamp=datetime.now(),
            frame_id=self._frame_counter,
            hash=img_hash,
            window_info=self.get_active_window(),
            changed_from_previous=changed
        )
        
        # Ajouter au buffer
        self._add_to_buffer(frame)
        
        return frame
    
    def capture_screen(self) -> Optional[Image.Image]:
        """
        Capture et retourne une PIL Image (compatibilité avec ExecutionLoop).
        
        Returns:
            PIL Image ou None
        """
        img = self.capture()
        if img is None:
            return None
        return Image.fromarray(img)
    
    def _capture_mss(self) -> np.ndarray:
        """Capture using mss - Option A: créer MSS à chaque capture (ultra stable)."""
        import mss
        # Créer une nouvelle instance MSS à chaque capture - zéro surprise, marche dans n'importe quel thread
        with mss.mss() as sct:
            monitor_idx = 1 if len(sct.monitors) > 1 else 0
            monitor = sct.monitors[monitor_idx]
            sct_img = sct.grab(monitor)
            
            img = np.array(sct_img)
            # Convert BGRA to RGB
            img = img[:, :, :3][:, :, ::-1]
            
            if img.size == 0 or img.shape[0] == 0 or img.shape[1] == 0:
                raise ValueError("Captured image has invalid dimensions")
            
            return img
    
    def _capture_pyautogui(self) -> np.ndarray:
        """Capture using pyautogui."""
        screenshot = self.pyautogui.screenshot()
        img = np.array(screenshot)
        
        if img.size == 0 or img.shape[0] == 0 or img.shape[1] == 0:
            raise ValueError("Captured image has invalid dimensions")
        
        return img
    
    # =========================================================================
    # Mode continu
    # =========================================================================
    
    def start_continuous(
        self,
        callback: Callable[[CaptureFrame], None],
        interval_ms: int = 500,
        skip_unchanged: bool = True
    ) -> bool:
        """
        Démarrer la capture continue.
        
        Args:
            callback: Fonction appelée pour chaque frame
            interval_ms: Intervalle entre captures (ms)
            skip_unchanged: Ne pas appeler callback si écran inchangé
        
        Returns:
            True si démarré avec succès
        """
        with self._lock:
            if self._continuous_running:
                logger.warning("Continuous capture already running")
                return False
            
            self._continuous_callback = callback
            self._continuous_interval_ms = interval_ms
            self._skip_unchanged = skip_unchanged
            self._continuous_running = True
            
            self._continuous_thread = threading.Thread(
                target=self._continuous_loop,
                daemon=True
            )
            self._continuous_thread.start()
            
            logger.info(f"Started continuous capture (interval={interval_ms}ms)")
            return True
    
    def stop_continuous(self) -> None:
        """Arrêter la capture continue."""
        with self._lock:
            self._continuous_running = False
        
        if self._continuous_thread:
            self._continuous_thread.join(timeout=2.0)
            self._continuous_thread = None
        
        logger.info("Stopped continuous capture")
    
    def is_continuous_running(self) -> bool:
        """Vérifier si la capture continue est active."""
        return self._continuous_running
    
    def _continuous_loop(self) -> None:
        """Boucle de capture continue (thread)."""
        last_capture_time = 0
        captures_in_second = 0
        second_start = time.time()
        
        # Créer une nouvelle instance mss pour ce thread (requis pour X11)
        thread_sct = None
        if self.method == "mss":
            import mss
            thread_sct = mss.mss()
        
        while self._continuous_running:
            try:
                # Capturer avec l'instance thread-local
                frame = self._capture_frame_threaded(thread_sct)
                
                if frame:
                    # Calculer FPS
                    captures_in_second += 1
                    if time.time() - second_start >= 1.0:
                        self._stats.captures_per_second = captures_in_second
                        captures_in_second = 0
                        second_start = time.time()
                    
                    # Appeler callback si changement ou si on ne skip pas
                    if self._continuous_callback:
                        if frame.changed_from_previous or not self._skip_unchanged:
                            try:
                                self._continuous_callback(frame)
                            except Exception as e:
                                logger.error(f"Callback error: {e}")
                
                # Attendre l'intervalle
                elapsed = (time.time() - last_capture_time) * 1000
                sleep_time = max(0, self._continuous_interval_ms - elapsed) / 1000.0
                if sleep_time > 0:
                    time.sleep(sleep_time)
                last_capture_time = time.time()
                
            except Exception as e:
                logger.error(f"Continuous capture error: {e}")
                time.sleep(0.1)
        
        # Cleanup thread-local mss
        if thread_sct:
            try:
                thread_sct.close()
            except Exception:
                pass
    
    # =========================================================================
    # Buffer et historique
    # =========================================================================
    
    def _add_to_buffer(self, frame: CaptureFrame) -> None:
        """Ajouter un frame au buffer circulaire."""
        with self._lock:
            self._buffer.append(frame)
            if len(self._buffer) > self.buffer_size:
                self._buffer.pop(0)
            self._stats.buffer_size = len(self._buffer)
            
            # Calculer utilisation mémoire
            if self._buffer:
                frame_size = self._buffer[0].image.nbytes / (1024 * 1024)
                self._stats.memory_usage_mb = frame_size * len(self._buffer)
    
    def get_buffer(self) -> List[CaptureFrame]:
        """Obtenir une copie du buffer."""
        with self._lock:
            return list(self._buffer)
    
    def get_last_frame(self) -> Optional[CaptureFrame]:
        """Obtenir le dernier frame capturé."""
        with self._lock:
            return self._buffer[-1] if self._buffer else None
    
    def get_frame_by_id(self, frame_id: int) -> Optional[CaptureFrame]:
        """Obtenir un frame par son ID."""
        with self._lock:
            for frame in self._buffer:
                if frame.frame_id == frame_id:
                    return frame
            return None
    
    def clear_buffer(self) -> None:
        """Vider le buffer."""
        with self._lock:
            self._buffer.clear()
            self._stats.buffer_size = 0
    
    # =========================================================================
    # Utilitaires
    # =========================================================================
    
    def _compute_hash(self, img: np.ndarray) -> str:
        """Calculer un hash rapide de l'image pour détecter les changements."""
        # Sous-échantillonner pour un hash rapide
        small = img[::20, ::20, :].tobytes()
        return hashlib.md5(small).hexdigest()
    
    def get_active_window(self) -> Optional[Dict]:
        """Obtenir les infos de la fenêtre active."""
        try:
            import pygetwindow as gw
            active = gw.getActiveWindow()
            if active:
                return {
                    'title': active.title,
                    'x': active.left,
                    'y': active.top,
                    'width': active.width,
                    'height': active.height,
                    'app': getattr(active, '_app', 'unknown')
                }
        except Exception as e:
            logger.debug(f"Could not get active window: {e}")
        return None
    
    def get_screen_resolution(self) -> Tuple[int, int]:
        """Obtenir la résolution de l'écran."""
        if self.method == "mss":
            import mss
            # Créer une instance temporaire pour obtenir la résolution
            with mss.mss() as sct:
                monitor = sct.monitors[1] if len(sct.monitors) > 1 else sct.monitors[0]
                return (monitor['width'], monitor['height'])
        else:
            size = self.pyautogui.size()
            return (size.width, size.height)
    
    def get_stats(self) -> CaptureStats:
        """Obtenir les statistiques de capture."""
        return self._stats
    
    def save_frame(self, frame: CaptureFrame, path: str) -> bool:
        """Sauvegarder un frame sur disque."""
        try:
            img = Image.fromarray(frame.image)
            img.save(path)
            return True
        except Exception as e:
            logger.error(f"Failed to save frame: {e}")
            return False
    
    def __del__(self):
        """Cleanup."""
        self.stop_continuous()
        # Plus besoin de fermer self.sct car nous créons MSS à chaque capture