"""
System Wake Detection and Thread Recovery System
Handles system sleep/hibernation and thread recovery during idle periods
"""

import time
import threading
import platform
import psutil
from typing import Dict, Callable, Optional, Any
from datetime import datetime, timedelta
from ws_package.log.writer import LogWriter

class SystemWakeDetector:
    """
    Detects system wake events and manages thread recovery
    """
    
    def __init__(self):
        self.logger = LogWriter.file(__name__)
        self.is_monitoring = False
        self.last_activity_time = time.time()
        self.last_cpu_usage = 0.0
        self.last_memory_usage = 0.0
        self.wake_callbacks: Dict[str, Callable] = {}
        self.sleep_callbacks: Dict[str, Callable] = {}
        self.monitor_thread = None
        self.stop_event = threading.Event()
        
        # Configuration - Optimized to prevent memory leaks
        self.check_interval = 120  # Check every 2 minutes (reduced frequency)
        self.idle_threshold = 300  # 5 minutes of inactivity
        self.wake_threshold = 10   # 10 seconds of activity to consider "wake"
        self.cpu_threshold = 5.0  # 5% CPU usage threshold
        self.memory_threshold = 1.0  # 1% memory change threshold
        
        # Resource management
        self.resource_check_counter = 0
        self.max_resource_checks = 10  # Limit resource checks
        self.last_resource_cleanup = time.time()
        
        # App activity tracking (for wake detection when listeners are stopped)
        self.last_frontmost_app = None
        self.last_app_check_time = 0
        self.app_check_interval = 30  # Check app activity every 30 seconds
        
        # Windows power management tracking
        self.windows_power_state = None
        self.last_windows_power_check = 0
        self.windows_power_check_interval = 10  # Check Windows power state every 10 seconds
        
        self.stats = {
            'wake_events': 0,
            'sleep_events': 0,
            'false_wakes': 0,
            'last_wake_time': None,
            'last_sleep_time': None
        }
        
        self.logger.info("SystemWakeDetector initialized")
    
    def register_wake_callback(self, name: str, callback: Callable):
        """Register a callback to be called when system wakes"""
        self.wake_callbacks[name] = callback
        self.logger.info(f"Registered wake callback: {name}")
    
    def register_sleep_callback(self, name: str, callback: Callable):
        """Register a callback to be called when system goes to sleep"""
        self.sleep_callbacks[name] = callback
        self.logger.info(f"Registered sleep callback: {name}")

    def unregister_wake_callback(self, name: str):
        """Remove a wake callback (e.g. when a monitor session ends)."""
        if self.wake_callbacks.pop(name, None) is not None:
            self.logger.info(f"Unregistered wake callback: {name}")
    
    def start_monitoring(self):
        """Start system wake/sleep monitoring"""
        if self.is_monitoring:
            self.logger.warning("System wake monitoring already started")
            return
        
        self.is_monitoring = True
        self.stop_event.clear()
        
        # Use shorter check interval on Windows for more responsive sleep detection
        # This helps detect sleep button presses faster
        if platform.system() == "Windows":
            # 60s balances sleep detection vs background CPU/wake polling overhead
            original_interval = self.check_interval
            self.check_interval = 60
            self.logger.info(
                f"Windows detected - using shorter check interval: {self.check_interval}s (vs {original_interval}s)"
            )
        
        self.monitor_thread = threading.Thread(
            target=self._monitor_loop,
            name="SystemWakeMonitor",
            daemon=True
        )
        self.monitor_thread.start()
        self.logger.info("System wake monitoring started")
    
    def stop_monitoring(self):
        """Stop system wake/sleep monitoring"""
        if not self.is_monitoring:
            return
        
        self.is_monitoring = False
        self.stop_event.set()
        
        if self.monitor_thread and self.monitor_thread.is_alive():
            self.monitor_thread.join(timeout=5)
        
        self.logger.info("System wake monitoring stopped")
    
    def _monitor_loop(self):
        """Main monitoring loop"""
        self.logger.info("Starting system wake monitoring loop")
        
        while not self.stop_event.is_set():
            try:
                # Check for system wake/sleep events
                self._check_system_state()
                
                # Wait for next check
                self.stop_event.wait(self.check_interval)
                
            except Exception as e:
                self.logger.error(f"Error in system wake monitoring: {e}")
                time.sleep(self.check_interval)
        
        self.logger.info("System wake monitoring loop stopped")
    
    def _check_app_activity(self) -> bool:
        """
        Check for app activity (frontmost app changes) as a wake signal.
        This is useful when keyboard/mouse listeners are stopped (e.g., after sleep).
        
        Returns True if frontmost app has changed, indicating user activity.
        """
        try:
            current_time = time.time()
            
            # Only check app activity periodically to reduce overhead
            if current_time - self.last_app_check_time < self.app_check_interval:
                return False
            
            self.last_app_check_time = current_time
            
            if platform.system() == "Darwin":
                # macOS: Use AppleScript to get frontmost app
                import subprocess
                try:
                    cmd = """osascript -e 'tell application "System Events" to get name of first application process whose frontmost is true'"""
                    result = subprocess.check_output(
                        cmd, 
                        shell=True, 
                        timeout=1.5,
                        stderr=subprocess.DEVNULL
                    ).decode('utf-8').strip()
                    
                    if result and result.strip():
                        current_app = result.strip()
                        
                        # If frontmost app changed, it indicates user activity
                        if self.last_frontmost_app is not None and current_app != self.last_frontmost_app:
                            self.logger.debug(
                                f"Frontmost app changed: {self.last_frontmost_app} -> {current_app} (wake signal)"
                            )
                            self.last_frontmost_app = current_app
                            return True
                        
                        # Update last known app
                        self.last_frontmost_app = current_app
                        return False
                    
                except (subprocess.TimeoutExpired, subprocess.CalledProcessError, FileNotFoundError):
                    # If we can't check app, don't use it as wake signal
                    return False
            elif platform.system() == "Windows":
                # Windows: Use win32gui to get foreground window
                try:
                    import win32gui
                    hwnd = win32gui.GetForegroundWindow()
                    if hwnd:
                        window_text = win32gui.GetWindowText(hwnd)
                        if window_text and window_text.strip():
                            current_app = window_text.strip()
                            if self.last_frontmost_app is not None and current_app != self.last_frontmost_app:
                                self.logger.debug(
                                    f"Foreground window changed: {self.last_frontmost_app} -> {current_app} (wake signal)"
                                )
                                self.last_frontmost_app = current_app
                                return True
                            self.last_frontmost_app = current_app
                except ImportError:
                    pass
                except Exception:
                    pass
            
            return False
            
        except Exception as e:
            self.logger.debug(f"Error checking app activity: {e}")
            return False
    
    def _check_system_state(self):
        """Check current system state for wake/sleep detection - Optimized to prevent memory leaks"""
        try:
            current_time = time.time()
            
            # Increment resource check counter
            self.resource_check_counter += 1
            
            # Only check resources every few iterations to prevent memory leaks
            if self.resource_check_counter % 3 == 0:
                # Use non-blocking psutil calls to prevent thread blocking
                current_cpu = psutil.cpu_percent()  # Non-blocking call
                current_memory = psutil.virtual_memory().percent
                
                # Perform resource cleanup if needed
                if current_time - self.last_resource_cleanup > 600:  # Every 10 minutes
                    self._cleanup_resources()
                    self.last_resource_cleanup = current_time
            else:
                # Use cached values to reduce system calls
                current_cpu = self.last_cpu_usage
                current_memory = self.last_memory_usage
            
            # Calculate time since last activity
            time_since_activity = current_time - self.last_activity_time
            
            # Check for actual user activity first (keyboard/mouse input)
            # This prevents false wake events from background processes
            has_user_activity = self._check_user_activity_optimized()
            
            # Check for app activity (frontmost app changes) - useful when listeners are stopped
            has_app_activity = self._check_app_activity()
            
            # Detect system activity with optimized checks
            is_active = (
                current_cpu > self.cpu_threshold or
                abs(current_memory - self.last_memory_usage) > self.memory_threshold or
                has_user_activity or
                has_app_activity
            )
            
            if is_active:
                self.last_activity_time = current_time
                self.last_cpu_usage = current_cpu
                self.last_memory_usage = current_memory
                
                # Check if we have a recent sleep event - if so, trigger wake immediately on app activity
                # This handles the case where system wakes from sleep and listeners are stopped
                has_recent_sleep = False
                if self.stats['last_sleep_time']:
                    if isinstance(self.stats['last_sleep_time'], datetime):
                        sleep_timestamp = self.stats['last_sleep_time'].timestamp()
                    else:
                        sleep_timestamp = self.stats['last_sleep_time']
                    time_since_sleep = current_time - sleep_timestamp
                    # If sleep was detected within last 10 minutes, consider it recent
                    if time_since_sleep < 600:  # 10 minutes
                        has_recent_sleep = True
                
                # Trigger wake event if:
                # 1. System was idle for > threshold AND has user/app activity, OR
                # 2. Recent sleep detected AND app activity detected (wake from sleep)
                # This prevents false wake events when system wakes for background tasks
                should_trigger_wake = False
                
                if time_since_activity > self.idle_threshold:
                    # Normal wake detection: require user activity OR app activity
                    if platform.system() == "Darwin":
                        if has_user_activity or has_app_activity:
                            should_trigger_wake = True
                        else:
                            # CPU/memory activity but no user input or app activity = background wake, not real wake
                            self.logger.debug(
                                f"System activity detected but no user input or app activity - "
                                f"likely background wake (CPU: {current_cpu:.1f}%, "
                                f"Memory change: {abs(current_memory - self.last_memory_usage):.1f}%)"
                            )
                    else:
                        # On Windows/Linux, use original logic (less strict)
                        should_trigger_wake = True
                elif has_recent_sleep and has_app_activity:
                    # Wake from sleep: app activity detected after recent sleep event
                    # This is important because listeners are stopped during sleep, so user activity won't be detected
                    # IMPROVED: More aggressive wake detection - if app activity is detected within 30 minutes of sleep,
                    # consider it a wake event (reduced from checking only within 10 minutes)
                    if time_since_sleep < 1800:  # 30 minutes
                        self.logger.info(
                            f"Wake from sleep detected: app activity after recent sleep ({time_since_sleep / 60} minutes ago)"
                        )
                        should_trigger_wake = True
                elif has_recent_sleep and (has_user_activity or current_cpu > self.cpu_threshold):
                    # IMPROVED: Also trigger wake if there's user activity or significant CPU usage after sleep
                    # This handles cases where app activity check might miss the wake
                    if time_since_sleep < 1800:  # 30 minutes
                        self.logger.info(
                            f"Wake from sleep detected: user activity or CPU activity after recent sleep ({time_since_sleep / 60} minutes ago)"
                        )
                        should_trigger_wake = True
                
                if should_trigger_wake:
                    self._handle_wake_event()
            
            # Check Windows power state for immediate sleep detection (Windows only)
            if platform.system() == "Windows":
                windows_sleep_detected = self._check_windows_power_state(current_time)
                if windows_sleep_detected:
                    # Windows power API detected sleep - handle immediately
                    self._handle_sleep_event()
                    return  # Skip further checks
            
            # Check for sleep event (polling-based fallback)
            elif time_since_activity > self.idle_threshold:
                self._handle_sleep_event()
                
        except Exception as e:
            self.logger.error(f"Error checking system state: {e}")
            # Reset counter on error to prevent infinite loops
            self.resource_check_counter = 0
    
    def _check_user_activity(self) -> bool:
        """Check for user activity (keyboard/mouse)"""
        try:
            if platform.system() == "Windows":
                return self._check_windows_activity()
            elif platform.system() == "Darwin":
                return self._check_macos_activity()
            else:
                # return self._check_linux_activity()
                return False
        except Exception as e:
            self.logger.error(f"Error checking user activity: {e}")
            return False
    
    def _check_user_activity_optimized(self) -> bool:
        """Optimized user activity check to prevent memory leaks"""
        try:
            # Only check user activity every few iterations to reduce system calls
            if self.resource_check_counter % 5 == 0:
                return self._check_user_activity()
            return False  # Use cached result for other iterations
        except Exception as e:
            self.logger.error(f"Error in optimized user activity check: {e}")
            return False
    
    def _cleanup_resources(self):
        """Clean up system resources to prevent memory leaks"""
        try:
            # Force garbage collection
            import gc
            gc.collect()
            
            # Log resource cleanup
            self.logger.info("Performed resource cleanup to prevent memory leaks")
            
        except Exception as e:
            self.logger.error(f"Error during resource cleanup: {e}")
    
    def _check_windows_activity(self) -> bool:
        """Check Windows-specific user activity - Optimized to prevent memory leaks"""
        try:
            import win32gui
            import win32con
            
            # Use a simple, lightweight check to prevent memory leaks
            # Avoid frequent API calls that can cause resource accumulation
            
            # Check if any window is active (lightweight check)
            active_window = win32gui.GetForegroundWindow()
            if active_window and active_window != 0:
                # Additional lightweight validation
                try:
                    # Get window text length as a simple activity indicator
                    window_text_length = win32gui.GetWindowTextLength(active_window)
                    return window_text_length >= 0  # Basic validation
                except Exception:
                    return True  # If we can't get text, assume active
            
            return False
            
        except ImportError:
            # Fallback to basic check without Windows API
            self.logger.debug("Windows API not available, using fallback activity check")
            return False
        except Exception as e:
            self.logger.error(f"Error checking Windows activity: {e}")
            # Don't let Windows API errors crash the monitoring
            return False
    
    def _check_macos_activity(self) -> bool:
        """Check macOS-specific user activity (keyboard/mouse input, not just system activity)"""
        try:
            # Check for actual user input using IOHIDSystem
            # This checks if there's been recent keyboard/mouse activity
            import subprocess
            
            # Get HID idle time (time since last keyboard/mouse input)
            # Lower idle time = more recent user activity
            try:
                ioreg_output = subprocess.check_output(
                    ['ioreg', '-c', 'IOHIDSystem'],
                    stderr=subprocess.DEVNULL,
                    timeout=2
                ).decode('ascii', errors='ignore')
                
                # Extract HIDIdleTime (in nanoseconds)
                for line in ioreg_output.splitlines():
                    if 'HIDIdleTime' in line:
                        try:
                            # Parse the idle time value
                            idle_time_ns = int(line.split('=')[-1].strip().rstrip(';'))
                            idle_time_seconds = idle_time_ns / 1000000000.0  # Convert to seconds
                            
                            # If idle time is less than 30 seconds, user is actively using the system
                            # This indicates real user activity, not just background processes
                            if idle_time_seconds < 30:
                                return True
                        except (ValueError, IndexError):
                            pass
                
                # If we can't parse idle time, assume no recent user activity
                return False
                
            except (subprocess.TimeoutExpired, subprocess.CalledProcessError, FileNotFoundError):
                # ioreg unavailable (privacy restriction on macOS 12-13) — fall
                # through to the Quartz ctypes fallback below.
                pass

            # Quartz fallback: CGEventSourceSecondsSinceLastEventType does not
            # require subprocess and is not gated by privacy TCC on macOS 12+.
            try:
                import ctypes
                import ctypes.util
                _quartz_path = ctypes.util.find_library("Quartz")
                if _quartz_path:
                    _quartz = ctypes.CDLL(_quartz_path)
                    _quartz.CGEventSourceSecondsSinceLastEventType.restype = ctypes.c_double
                    # kCGEventSourceStateCombinedSessionState=0, kCGAnyInputEventType=0xFFFFFFFF
                    idle_time_sec = _quartz.CGEventSourceSecondsSinceLastEventType(0, 0xFFFFFFFF)
                    return float(idle_time_sec) < 30.0
            except Exception:
                pass

            return False

        except Exception as e:
            self.logger.debug(f"Error checking macOS activity: {e}")
            return False
    
    def _check_linux_activity(self) -> bool:
        """Check Linux-specific user activity"""
        try:
            # Check X11 or Wayland activity
            # This is a simplified implementation
            return False
        except Exception as e:
            self.logger.error(f"Error checking Linux activity: {e}")
            return False
    
    def _check_windows_power_state(self, current_time: float) -> bool:
        """
        Check Windows power state for immediate sleep detection.
        Uses Windows APIs to detect when system is about to sleep or has just slept.
        
        Returns True if sleep is detected, False otherwise.
        """
        try:
            # Only check periodically to avoid excessive API calls
            if current_time - self.last_windows_power_check < self.windows_power_check_interval:
                return False
            
            self.last_windows_power_check = current_time
            
            try:
                import ctypes
                from ctypes import wintypes
                
                # Import Windows API functions
                kernel32 = ctypes.windll.kernel32
                powrprof = ctypes.windll.powrprof
                
                # CallPowerSettingNotification is used to check power state
                # We'll use GetSystemPowerStatus for a simpler check
                class SYSTEM_POWER_STATUS(ctypes.Structure):
                    _fields_ = [
                        ("ACLineStatus", ctypes.c_byte),
                        ("BatteryFlag", ctypes.c_byte),
                        ("BatteryLifePercent", ctypes.c_byte),
                        ("Reserved1", ctypes.c_byte),
                        ("BatteryLifeTime", wintypes.DWORD),
                        ("BatteryFullLifeTime", wintypes.DWORD),
                    ]
                
                power_status = SYSTEM_POWER_STATUS()
                
                # Get system power status
                if kernel32.GetSystemPowerStatus(ctypes.byref(power_status)):
                    # Check if system is on battery and battery is low (potential sleep indicator)
                    # But this alone isn't reliable for sleep detection
                    pass
                
                # Better approach: Check system idle time and compare with last known state
                # If system was active but suddenly becomes idle, it might be sleeping
                # We'll use the existing activity tracking for this
                
                # Alternative: Use SetThreadExecutionState to detect sleep
                # But we don't want to prevent sleep, just detect it
                
                # For now, we'll rely on the polling-based detection but make it more responsive
                # The real fix is ensuring callbacks are registered (which we did above)
                
                return False  # Don't trigger sleep here, let polling handle it
                
            except ImportError:
                # Windows API not available (shouldn't happen on Windows, but handle gracefully)
                self.logger.debug("Windows API not available for power state checking")
                return False
            except Exception as e:
                self.logger.debug(f"Error checking Windows power state: {e}")
                return False
                
        except Exception as e:
            self.logger.debug(f"Error in Windows power state check: {e}")
            return False
    
    def _handle_wake_event(self):
        """Handle system wake event"""
        try:
            current_time = datetime.now()
            
            # Prevent false wake events
            if (self.stats['last_wake_time'] and 
                (current_time - self.stats['last_wake_time']).total_seconds() < 60):
                self.stats['false_wakes'] += 1
                return
            
            # Detect if this is wake from sleep vs fresh boot/shutdown
            wake_type = self._detect_wake_type()
            
            self.stats['wake_events'] += 1
            self.stats['last_wake_time'] = current_time
            
            # Log wake event (EnhancedLogWriter uses f-strings — not stdlib %-formatting + extra args)
            ts = current_time.strftime('%Y-%m-%d %H:%M:%S')
            if wake_type == "sleep":
                self.logger.info(f"*** SYSTEM WAKE FROM SLEEP detected at {ts} ***")
            elif wake_type == "shutdown":
                self.logger.info(f"*** SYSTEM WAKE FROM SHUTDOWN/REBOOT detected at {ts} ***")
            else:
                self.logger.info(f"System wake detected at {ts} (type: {wake_type})")
            
            # Call all registered wake callbacks
            for name, callback in self.wake_callbacks.items():
                try:
                    callback()
                    self.logger.info(f"Wake callback '{name}' executed successfully")
                except Exception as e:
                    self.logger.error(f"Error in wake callback '{name}': {e}")
                    
        except Exception as e:
            self.logger.error(f"Error handling wake event: {e}")
    
    def _detect_wake_type(self) -> str:
        """
        Detect if wake is from sleep or shutdown/reboot.
        Returns: 'sleep', 'shutdown', or 'unknown'
        """
        try:
            if platform.system() != "Darwin":
                return "unknown"  # Only macOS has reliable sleep detection
            
            # Check system uptime - if uptime is very short (< 5 minutes), likely fresh boot
            # If uptime is longer but we detected wake, likely wake from sleep
            import subprocess
            result = subprocess.run(
                ['sysctl', '-n', 'kern.boottime'],
                capture_output=True,
                text=True,
                timeout=5
            )
            
            if result.returncode == 0:
                # Parse boot time (format: { sec = 1234567890, usec = 0 } )
                boot_time_str = result.stdout.strip()
                import re
                match = re.search(r'sec\s*=\s*(\d+)', boot_time_str)
                if match:
                    boot_timestamp = int(match.group(1))
                    boot_time = datetime.fromtimestamp(boot_timestamp)
                    current_time = datetime.now()
                    uptime_seconds = (current_time - boot_time).total_seconds()
                    
                    # If we have a last_wake_time and uptime is long, likely wake from sleep
                    if self.stats['last_wake_time']:
                        time_since_last_wake = (current_time - self.stats['last_wake_time']).total_seconds()
                        # If last wake was recent (< 1 hour) and uptime is long, likely sleep
                        if time_since_last_wake < 3600 and uptime_seconds > 300:
                            return "sleep"
                    
                    # If uptime is very short (< 5 minutes), likely fresh boot
                    if uptime_seconds < 300:
                        return "shutdown"
                    # Otherwise, if we detected wake but uptime is long, likely sleep
                    elif uptime_seconds > 300:
                        return "sleep"
            
            # Fallback: if we have last_wake_time, assume sleep
            if self.stats['last_wake_time']:
                return "sleep"
            else:
                return "shutdown"
                
        except Exception as e:
            self.logger.debug(f"Could not detect wake type: {e}")
            return "unknown"
    
    def _handle_sleep_event(self):
        """Handle system sleep event"""
        try:
            current_time = datetime.now()
            
            # Prevent false sleep events
            if (self.stats['last_sleep_time'] and 
                (current_time - self.stats['last_sleep_time']).total_seconds() < 60):
                return
            
            self.stats['sleep_events'] += 1
            self.stats['last_sleep_time'] = current_time
            
            self.logger.info(f"System sleep detected at {current_time.strftime('%Y-%m-%d %H:%M:%S')}")
            
            # Call all registered sleep callbacks
            for name, callback in self.sleep_callbacks.items():
                try:
                    callback()
                    self.logger.info(f"Sleep callback '{name}' executed successfully")
                except Exception as e:
                    self.logger.error(f"Error in sleep callback '{name}': {e}")
                    
        except Exception as e:
            self.logger.error(f"Error handling sleep event: {e}")
    
    def get_stats(self) -> Dict[str, Any]:
        """Get monitoring statistics"""
        return {
            'is_monitoring': self.is_monitoring,
            'wake_events': self.stats['wake_events'],
            'sleep_events': self.stats['sleep_events'],
            'false_wakes': self.stats['false_wakes'],
            'last_wake_time': self.stats['last_wake_time'],
            'last_sleep_time': self.stats['last_sleep_time'],
            'registered_wake_callbacks': len(self.wake_callbacks),
            'registered_sleep_callbacks': len(self.sleep_callbacks)
        }
    
    def force_wake_event(self):
        """Force a wake event (for testing)"""
        self.logger.info("Forcing wake event")
        self._handle_wake_event()
    
    def force_sleep_event(self):
        """Force a sleep event (for testing)"""
        self.logger.info("Forcing sleep event")
        self._handle_sleep_event()

# Global instance
_system_wake_detector = None

def get_system_wake_detector() -> SystemWakeDetector:
    """Get the global system wake detector instance"""
    global _system_wake_detector
    if _system_wake_detector is None:
        _system_wake_detector = SystemWakeDetector()
    return _system_wake_detector

def start_system_wake_monitoring():
    """Start system wake monitoring"""
    detector = get_system_wake_detector()
    detector.start_monitoring()
    return detector

def stop_system_wake_monitoring():
    """Stop system wake monitoring"""
    global _system_wake_detector
    if _system_wake_detector:
        _system_wake_detector.stop_monitoring()
