backend/performance/resource_monitor.py

"""
Resource Monitor for AegisLM SaaS Backend.

Monitors system resources including CPU, memory, disk usage,
and queue sizes to prevent system overload.
"""

import asyncio
import psutil
import time
from typing import Dict, Any, Optional, List
from dataclasses import dataclass
from datetime import datetime, timezone
import logging
from collections import deque

logger = logging.getLogger(__name__)


@dataclass
class SystemResources:
    """System resource snapshot."""
    cpu_percent: float
    memory_percent: float
    memory_available_gb: float
    disk_percent: float
    disk_free_gb: float
    load_average: List[float]
    active_processes: int
    timestamp: datetime


@dataclass
class ResourceThresholds:
    """Resource monitoring thresholds."""
    cpu_warning: float = 70.0
    cpu_critical: float = 90.0
    memory_warning: float = 75.0
    memory_critical: float = 90.0
    disk_warning: float = 80.0
    disk_critical: float = 95.0
    queue_warning: int = 50
    queue_critical: int = 100


class ResourceMonitor:
    """
    Monitors system resources and queue status.
    
    Features:
    - Real-time resource monitoring
    - Threshold-based alerts
    - Historical data tracking
    - Automatic resource-based scaling decisions
    """
    
    def __init__(self, check_interval_seconds: int = 30, history_size: int = 1000):
        """
        Initialize resource monitor.
        
        Args:
            check_interval_seconds: Monitoring interval
            history_size: Number of historical data points to keep
        """
        self.check_interval = check_interval_seconds
        self.history_size = history_size
        
        # Historical data
        self.cpu_history: deque = deque(maxlen=history_size)
        self.memory_history: deque = deque(maxlen=history_size)
        self.disk_history: deque = deque(maxlen=history_size)
        
        # Current state
        self.current_resources: Optional[SystemResources] = None
        self.thresholds = ResourceThresholds()
        
        # Monitoring task
        self._monitor_task: Optional[asyncio.Task] = None
        self._running = False
        
        # Alert callbacks
        self._alert_callbacks: List[callable] = []
        
        logger.info(f"ResourceMonitor initialized with {check_interval_seconds}s interval")
    
    async def start(self):
        """Start resource monitoring."""
        if self._running:
            return
        
        self._running = True
        self._monitor_task = asyncio.create_task(self._monitor_loop())
        logger.info("ResourceMonitor started")
    
    async def stop(self):
        """Stop resource monitoring."""
        self._running = False
        if self._monitor_task:
            self._monitor_task.cancel()
            try:
                await self._monitor_task
            except asyncio.CancelledError:
                pass
        logger.info("ResourceMonitor stopped")
    
    async def _monitor_loop(self):
        """Main monitoring loop."""
        while self._running:
            try:
                await self._collect_metrics()
                await self._check_thresholds()
                await asyncio.sleep(self.check_interval)
            except asyncio.CancelledError:
                break
            except Exception as e:
                logger.error(f"Resource monitoring error: {e}")
                await asyncio.sleep(self.check_interval)
    
    async def _collect_metrics(self):
        """Collect current system metrics."""
        try:
            # CPU metrics
            cpu_percent = psutil.cpu_percent(interval=1)
            
            # Memory metrics
            memory = psutil.virtual_memory()
            memory_percent = memory.percent
            memory_available_gb = memory.available / (1024**3)
            
            # Disk metrics
            disk = psutil.disk_usage('/')
            disk_percent = (disk.used / disk.total) * 100
            disk_free_gb = disk.free / (1024**3)
            
            # Load average (Unix-like systems)
            try:
                load_average = list(psutil.getloadavg())
            except (AttributeError, OSError):
                # Windows or system without load average
                load_average = [0.0, 0.0, 0.0]
            
            # Process count
            active_processes = len(psutil.pids())
            
            # Create resource snapshot
            self.current_resources = SystemResources(
                cpu_percent=cpu_percent,
                memory_percent=memory_percent,
                memory_available_gb=memory_available_gb,
                disk_percent=disk_percent,
                disk_free_gb=disk_free_gb,
                load_average=load_average,
                active_processes=active_processes,
                timestamp=datetime.now(timezone.utc)
            )
            
            # Store in history
            timestamp = time.time()
            self.cpu_history.append((timestamp, cpu_percent))
            self.memory_history.append((timestamp, memory_percent))
            self.disk_history.append((timestamp, disk_percent))
            
            logger.debug(f"Collected metrics: CPU={cpu_percent}%, MEM={memory_percent}%")
            
        except Exception as e:
            logger.error(f"Error collecting metrics: {e}")
    
    async def _check_thresholds(self):
        """Check resource thresholds and trigger alerts."""
        if not self.current_resources:
            return
        
        alerts = []
        
        # CPU checks
        if self.current_resources.cpu_percent >= self.thresholds.cpu_critical:
            alerts.append({
                "type": "cpu_critical",
                "value": self.current_resources.cpu_percent,
                "threshold": self.thresholds.cpu_critical,
                "message": f"CPU usage critical: {self.current_resources.cpu_percent:.1f}%"
            })
        elif self.current_resources.cpu_percent >= self.thresholds.cpu_warning:
            alerts.append({
                "type": "cpu_warning",
                "value": self.current_resources.cpu_percent,
                "threshold": self.thresholds.cpu_warning,
                "message": f"CPU usage high: {self.current_resources.cpu_percent:.1f}%"
            })
        
        # Memory checks
        if self.current_resources.memory_percent >= self.thresholds.memory_critical:
            alerts.append({
                "type": "memory_critical",
                "value": self.current_resources.memory_percent,
                "threshold": self.thresholds.memory_critical,
                "message": f"Memory usage critical: {self.current_resources.memory_percent:.1f}%"
            })
        elif self.current_resources.memory_percent >= self.thresholds.memory_warning:
            alerts.append({
                "type": "memory_warning",
                "value": self.current_resources.memory_percent,
                "threshold": self.thresholds.memory_warning,
                "message": f"Memory usage high: {self.current_resources.memory_percent:.1f}%"
            })
        
        # Disk checks
        if self.current_resources.disk_percent >= self.thresholds.disk_critical:
            alerts.append({
                "type": "disk_critical",
                "value": self.current_resources.disk_percent,
                "threshold": self.thresholds.disk_critical,
                "message": f"Disk usage critical: {self.current_resources.disk_percent:.1f}%"
            })
        elif self.current_resources.disk_percent >= self.thresholds.disk_warning:
            alerts.append({
                "type": "disk_warning",
                "value": self.current_resources.disk_percent,
                "threshold": self.thresholds.disk_warning,
                "message": f"Disk usage high: {self.current_resources.disk_percent:.1f}%"
            })
        
        # Trigger alerts
        for alert in alerts:
            await self._trigger_alert(alert)
    
    async def _trigger_alert(self, alert: Dict[str, Any]):
        """Trigger resource alert."""
        logger.warning(alert["message"])
        
        # Call registered callbacks
        for callback in self._alert_callbacks:
            try:
                await callback(alert)
            except Exception as e:
                logger.error(f"Alert callback error: {e}")
    
    def add_alert_callback(self, callback: callable):
        """
        Add alert callback function.
        
        Args:
            callback: Async function to call on alerts
        """
        self._alert_callbacks.append(callback)
    
    def remove_alert_callback(self, callback: callable):
        """
        Remove alert callback function.
        
        Args:
            callback: Callback function to remove
        """
        if callback in self._alert_callbacks:
            self._alert_callbacks.remove(callback)
    
    async def get_current_resources(self) -> Optional[SystemResources]:
        """
        Get current system resources.
        
        Returns:
            SystemResources: Current resource snapshot
        """
        return self.current_resources
    
    async def get_resource_history(self, minutes: int = 60) -> Dict[str, List[Dict[str, Any]]]:
        """
        Get historical resource data.
        
        Args:
            minutes: Number of minutes of history to return
            
        Returns:
            Dict: Historical resource data
        """
        cutoff_time = time.time() - (minutes * 60)
        
        def filter_history(history: deque) -> List[Dict[str, Any]]:
            return [
                {
                    "timestamp": timestamp,
                    "value": value,
                    "datetime": datetime.fromtimestamp(timestamp, timezone.utc).isoformat()
                }
                for timestamp, value in history if timestamp >= cutoff_time
            ]
        
        return {
            "cpu": filter_history(self.cpu_history),
            "memory": filter_history(self.memory_history),
            "disk": filter_history(self.disk_history)
        }
    
    async def get_resource_summary(self) -> Dict[str, Any]:
        """
        Get resource usage summary.
        
        Returns:
            Dict: Resource summary
        """
        if not self.current_resources:
            return {"status": "no_data"}
        
        # Calculate averages over last hour
        history = await self.get_resource_history(60)
        
        avg_cpu = sum(point["value"] for point in history["cpu"]) / len(history["cpu"]) if history["cpu"] else 0
        avg_memory = sum(point["value"] for point in history["memory"]) / len(history["memory"]) if history["memory"] else 0
        avg_disk = sum(point["value"] for point in history["disk"]) / len(history["disk"]) if history["disk"] else 0
        
        return {
            "current": {
                "cpu_percent": self.current_resources.cpu_percent,
                "memory_percent": self.current_resources.memory_percent,
                "memory_available_gb": self.current_resources.memory_available_gb,
                "disk_percent": self.current_resources.disk_percent,
                "disk_free_gb": self.current_resources.disk_free_gb,
                "load_average": self.current_resources.load_average,
                "active_processes": self.current_resources.active_processes,
                "timestamp": self.current_resources.timestamp.isoformat()
            },
            "averages_last_hour": {
                "cpu_percent": avg_cpu,
                "memory_percent": avg_memory,
                "disk_percent": avg_disk
            },
            "thresholds": {
                "cpu_warning": self.thresholds.cpu_warning,
                "cpu_critical": self.thresholds.cpu_critical,
                "memory_warning": self.thresholds.memory_warning,
                "memory_critical": self.thresholds.memory_critical,
                "disk_warning": self.thresholds.disk_warning,
                "disk_critical": self.thresholds.disk_critical
            },
            "status": self._get_system_status()
        }
    
    def _get_system_status(self) -> str:
        """
        Get overall system status based on current resources.
        
        Returns:
            str: System status (healthy, warning, critical)
        """
        if not self.current_resources:
            return "unknown"
        
        # Check for critical conditions
        if (self.current_resources.cpu_percent >= self.thresholds.cpu_critical or
            self.current_resources.memory_percent >= self.thresholds.memory_critical or
            self.current_resources.disk_percent >= self.thresholds.disk_critical):
            return "critical"
        
        # Check for warning conditions
        if (self.current_resources.cpu_percent >= self.thresholds.cpu_warning or
            self.current_resources.memory_percent >= self.thresholds.memory_warning or
            self.current_resources.disk_percent >= self.thresholds.disk_warning):
            return "warning"
        
        return "healthy"
    
    async def check_can_handle_load(self, additional_tasks: int = 1) -> Dict[str, Any]:
        """
        Check if system can handle additional load.
        
        Args:
            additional_tasks: Number of additional tasks to evaluate
            
        Returns:
            Dict: Load assessment result
        """
        if not self.current_resources:
            return {"can_handle": False, "reason": "no_resource_data"}
        
        reasons = []
        can_handle = True
        
        # CPU check
        if self.current_resources.cpu_percent >= self.thresholds.cpu_warning:
            can_handle = False
            reasons.append(f"CPU usage too high: {self.current_resources.cpu_percent:.1f}%")
        
        # Memory check
        if self.current_resources.memory_percent >= self.thresholds.memory_warning:
            can_handle = False
            reasons.append(f"Memory usage too high: {self.current_resources.memory_percent:.1f}%")
        
        # Disk check
        if self.current_resources.disk_percent >= self.thresholds.disk_warning:
            can_handle = False
            reasons.append(f"Disk usage too high: {self.current_resources.disk_percent:.1f}%")
        
        return {
            "can_handle": can_handle,
            "reason": "; ".join(reasons) if reasons else "system_healthy",
            "current_resources": {
                "cpu_percent": self.current_resources.cpu_percent,
                "memory_percent": self.current_resources.memory_percent,
                "disk_percent": self.current_resources.disk_percent
            }
        }
    
    def set_thresholds(self, thresholds: ResourceThresholds):
        """
        Update resource monitoring thresholds.
        
        Args:
            thresholds: New threshold values
        """
        self.thresholds = thresholds
        logger.info("Resource monitoring thresholds updated")


# Global resource monitor instance
resource_monitor = ResourceMonitor()


def get_resource_monitor() -> ResourceMonitor:
    """
    Get the global resource monitor instance.
    
    Returns:
        ResourceMonitor: Global instance
    """
    return resource_monitor


# Alert callback for automatic load management
async def auto_load_management_alert(alert: Dict[str, Any]):
    """
    Automatic load management alert handler.
    
    Args:
        alert: Alert information
    """
    from performance.concurrency_manager import get_concurrency_manager
    from performance.metrics_tracker import get_metrics_tracker
    
    if alert["type"] in ["cpu_critical", "memory_critical"]:
        # Reduce concurrent tasks in critical conditions
        manager = get_concurrency_manager()
        if manager.max_concurrent_tasks > 1:
            manager.max_concurrent_tasks = max(1, manager.max_concurrent_tasks - 1)
            logger.warning(f"Reduced max concurrent tasks to {manager.max_concurrent_tasks} due to {alert['type']}")
    
    # Record system metrics
    monitor = get_resource_monitor()
    resources = await monitor.get_current_resources()
    if resources:
        tracker = get_metrics_tracker()
        await tracker.record_system_metrics(resources.cpu_percent, resources.memory_percent)


# Register auto load management
resource_monitor.add_alert_callback(auto_load_management_alert)