alert_system.py

#!/usr/bin/env python3
"""
Alert and Notification System

Provides comprehensive alerting for GPU monitoring with multiple notification
channels, threshold management, and alert history tracking.
"""

import time
import json
import logging
import smtplib
import subprocess
import threading
from typing import Dict, List, Optional, Callable, Any
from dataclasses import dataclass, asdict
from datetime import datetime, timedelta
from email.mime.text import MIMEText
from email.mime.multipart import MIMEMultipart
from pathlib import Path

from gpu_monitoring import GPUStatus, GPUDataManager

logger = logging.getLogger(__name__)


@dataclass
class AlertThreshold:
    """Alert threshold configuration."""
    metric: str
    threshold: float
    operator: str  # '>', '<', '>=', '<=', '=='
    duration: int  # seconds to maintain threshold before alerting
    enabled: bool = True
    cooldown: int = 300  # seconds between same alerts


@dataclass
class Alert:
    """Alert instance."""
    id: str
    timestamp: float
    gpu_name: str
    metric: str
    value: float
    threshold: float
    message: str
    severity: str  # 'info', 'warning', 'critical', 'emergency'
    acknowledged: bool = False
    resolved: bool = False
    resolved_at: Optional[float] = None


class NotificationChannel:
    """Base class for notification channels."""
    
    def __init__(self, name: str, enabled: bool = True):
        self.name = name
        self.enabled = enabled
    
    def send(self, alert: Alert) -> bool:
        """Send notification for alert."""
        raise NotImplementedError


class LogNotification(NotificationChannel):
    """Log-based notification channel."""
    
    def send(self, alert: Alert) -> bool:
        """Send alert to log."""
        severity_colors = {
            'info': '\033[94m',      # Blue
            'warning': '\033[93m',   # Yellow
            'critical': '\033[91m',  # Red
            'emergency': '\033[41m'  # Red background
        }
        reset_color = '\033[0m'
        
        color = severity_colors.get(alert.severity, '')
        
        message = (
            f"{color}[{alert.severity.upper()}] GPU Alert: {alert.gpu_name} - "
            f"{alert.metric}: {alert.value} (threshold: {alert.threshold}){reset_color}"
        )
        
        if alert.severity in ['critical', 'emergency']:
            logger.critical(message)
        elif alert.severity == 'warning':
            logger.warning(message)
        else:
            logger.info(message)
        
        return True


class DesktopNotification(NotificationChannel):
    """Desktop notification channel using notify-send."""
    
    def send(self, alert: Alert) -> bool:
        """Send desktop notification."""
        try:
            urgency = 'normal'
            if alert.severity == 'critical':
                urgency = 'critical'
            elif alert.severity == 'warning':
                urgency = 'normal'
            else:
                urgency = 'low'
            
            # Create notification command
            cmd = [
                'notify-send',
                '--app-name=GPU Monitor',
                f'--urgency={urgency}',
                f'--icon=video-display',
                f'GPU Alert - {alert.gpu_name}',
                f'{alert.message}'
            ]
            
            subprocess.run(cmd, check=True)
            return True
            
        except subprocess.CalledProcessError as e:
            logger.error(f"Failed to send desktop notification: {e}")
            return False
        except FileNotFoundError:
            logger.warning("notify-send not found, desktop notifications disabled")
            return False


class EmailNotification(NotificationChannel):
    """Email notification channel."""
    
    def __init__(self, name: str, smtp_config: Dict[str, Any], recipients: List[str]):
        super().__init__(name)
        self.smtp_config = smtp_config
        self.recipients = recipients
    
    def send(self, alert: Alert) -> bool:
        """Send email notification."""
        try:
            # Create message
            msg = MIMEMultipart()
            msg['From'] = self.smtp_config['sender']
            msg['To'] = ', '.join(self.recipients)
            msg['Subject'] = f'GPU Alert: {alert.gpu_name} - {alert.severity.upper()}'
            
            # Create email body
            body = f"""
GPU Monitoring Alert

GPU: {alert.gpu_name}
Metric: {alert.metric}
Current Value: {alert.value}
Threshold: {alert.threshold}
Severity: {alert.severity.upper()}
Time: {datetime.fromtimestamp(alert.timestamp).strftime('%Y-%m-%d %H:%M:%S')}

Message: {alert.message}

This is an automated message from the GPU Monitoring System.
"""
            
            msg.attach(MIMEText(body, 'plain'))
            
            # Send email
            with smtplib.SMTP(self.smtp_config['server'], self.smtp_config['port']) as server:
                if self.smtp_config.get('use_tls', True):
                    server.starttls()
                
                if 'username' in self.smtp_config and 'password' in self.smtp_config:
                    server.login(self.smtp_config['username'], self.smtp_config['password'])
                
                server.send_message(msg)
            
            return True
            
        except Exception as e:
            logger.error(f"Failed to send email notification: {e}")
            return False


class WebhookNotification(NotificationChannel):
    """Webhook notification channel."""
    
    def __init__(self, name: str, webhook_url: str, headers: Optional[Dict[str, str]] = None):
        super().__init__(name)
        self.webhook_url = webhook_url
        self.headers = headers or {}
    
    def send(self, alert: Alert) -> bool:
        """Send webhook notification."""
        try:
            import requests
            
            payload = {
                'alert_id': alert.id,
                'timestamp': alert.timestamp,
                'gpu_name': alert.gpu_name,
                'metric': alert.metric,
                'value': alert.value,
                'threshold': alert.threshold,
                'message': alert.message,
                'severity': alert.severity,
                'acknowledged': alert.acknowledged
            }
            
            response = requests.post(
                self.webhook_url,
                json=payload,
                headers=self.headers,
                timeout=10
            )
            
            return response.status_code == 200
            
        except Exception as e:
            logger.error(f"Failed to send webhook notification: {e}")
            return False


class AlertManager:
    """Main alert management system."""
    
    def __init__(self, config_file: str = "config/alerts.json", db_path: str = "data/gpu_monitoring.db"):
        self.config_file = config_file
        self.db_manager = GPUDataManager(db_path)
        
        # Alert state tracking
        self.active_alerts = {}  # alert_id -> Alert
        self.threshold_states = {}  # (gpu, metric) -> {'value': float, 'start_time': float}
        self.last_alert_times = {}  # (gpu, metric) -> timestamp
        
        # Notification channels
        self.channels = []
        
        # Configuration
        self.thresholds = []
        self.alert_settings = {}
        
        # Threading
        self.running = False
        self.thread = None
        
        # Load configuration
        self.load_config()
        self.setup_channels()
    
    def load_config(self):
        """Load alert configuration from file."""
        try:
            if Path(self.config_file).exists():
                with open(self.config_file, 'r') as f:
                    config = json.load(f)
                
                # Load thresholds
                self.thresholds = []
                for threshold_data in config.get('thresholds', []):
                    threshold = AlertThreshold(**threshold_data)
                    self.thresholds.append(threshold)
                
                # Load settings
                self.alert_settings = config.get('settings', {
                    'check_interval': 5.0,
                    'cleanup_interval': 3600.0,
                    'max_alert_age': 86400.0  # 24 hours
                })
                
                logger.info(f"Loaded {len(self.thresholds)} alert thresholds")
            else:
                self.create_default_config()
                
        except Exception as e:
            logger.error(f"Error loading alert config: {e}")
            self.create_default_config()
    
    def create_default_config(self):
        """Create default alert configuration."""
        default_config = {
            'thresholds': [
                {
                    'metric': 'temperature',
                    'threshold': 75.0,
                    'operator': '>=',
                    'duration': 10,
                    'enabled': True,
                    'cooldown': 300
                },
                {
                    'metric': 'temperature',
                    'threshold': 85.0,
                    'operator': '>=',
                    'duration': 5,
                    'enabled': True,
                    'cooldown': 600
                },
                {
                    'metric': 'load',
                    'threshold': 90.0,
                    'operator': '>=',
                    'duration': 30,
                    'enabled': True,
                    'cooldown': 600
                },
                {
                    'metric': 'power_draw',
                    'threshold': 200.0,
                    'operator': '>=',
                    'duration': 10,
                    'enabled': True,
                    'cooldown': 300
                },
                {
                    'metric': 'fan_speed',
                    'threshold': 95.0,
                    'operator': '>=',
                    'duration': 60,
                    'enabled': True,
                    'cooldown': 1800
                }
            ],
            'settings': {
                'check_interval': 5.0,
                'cleanup_interval': 3600.0,
                'max_alert_age': 86400.0
            },
            'notifications': {
                'log': {'enabled': True},
                'desktop': {'enabled': True},
                'email': {
                    'enabled': False,
                    'smtp': {
                        'server': 'smtp.gmail.com',
                        'port': 587,
                        'use_tls': True,
                        'sender': 'your-email@gmail.com',
                        'username': 'your-username',
                        'password': 'your-app-password'
                    },
                    'recipients': ['admin@example.com']
                },
                'webhook': {
                    'enabled': False,
                    'url': 'https://your-webhook-url.com/alerts',
                    'headers': {'Authorization': 'Bearer your-token'}
                }
            }
        }
        
        # Save default config
        Path(self.config_file).parent.mkdir(parents=True, exist_ok=True)
        with open(self.config_file, 'w') as f:
            json.dump(default_config, f, indent=2)
        
        logger.info("Created default alert configuration")
    
    def setup_channels(self):
        """Setup notification channels."""
        try:
            if Path(self.config_file).exists():
                with open(self.config_file, 'r') as f:
                    config = json.load(f)
                
                notifications = config.get('notifications', {})
                
                # Log channel
                if notifications.get('log', {}).get('enabled', True):
                    self.channels.append(LogNotification('log'))
                
                # Desktop channel
                if notifications.get('desktop', {}).get('enabled', True):
                    self.channels.append(DesktopNotification('desktop'))
                
                # Email channel
                email_config = notifications.get('email', {})
                if email_config.get('enabled', False):
                    smtp_config = email_config.get('smtp', {})
                    recipients = email_config.get('recipients', [])
                    if smtp_config and recipients:
                        self.channels.append(EmailNotification(
                            'email', smtp_config, recipients
                        ))
                
                # Webhook channel
                webhook_config = notifications.get('webhook', {})
                if webhook_config.get('enabled', False):
                    url = webhook_config.get('url')
                    headers = webhook_config.get('headers', {})
                    if url:
                        self.channels.append(WebhookNotification('webhook', url, headers))
                
                logger.info(f"Setup {len(self.channels)} notification channels")
                
        except Exception as e:
            logger.error(f"Error setting up notification channels: {e}")
    
    def add_threshold(self, threshold: AlertThreshold):
        """Add a new alert threshold."""
        self.thresholds.append(threshold)
        self.save_config()
    
    def remove_threshold(self, metric: str, threshold_value: float):
        """Remove an alert threshold."""
        self.thresholds = [
            t for t in self.thresholds 
            if not (t.metric == metric and t.threshold == threshold_value)
        ]
        self.save_config()
    
    def save_config(self):
        """Save current configuration to file."""
        try:
            config = {
                'thresholds': [asdict(t) for t in self.thresholds],
                'settings': self.alert_settings,
                'notifications': {
                    'log': {'enabled': any(isinstance(c, LogNotification) for c in self.channels)},
                    'desktop': {'enabled': any(isinstance(c, DesktopNotification) for c in self.channels)},
                }
            }
            
            # Add email and webhook configs if channels exist
            email_channel = next((c for c in self.channels if isinstance(c, EmailNotification)), None)
            if email_channel:
                config['notifications']['email'] = {
                    'enabled': True,
                    'smtp': email_channel.smtp_config,
                    'recipients': email_channel.recipients
                }
            
            webhook_channel = next((c for c in self.channels if isinstance(c, WebhookNotification)), None)
            if webhook_channel:
                config['notifications']['webhook'] = {
                    'enabled': True,
                    'url': webhook_channel.webhook_url,
                    'headers': webhook_channel.headers
                }
            
            with open(self.config_file, 'w') as f:
                json.dump(config, f, indent=2)
            
        except Exception as e:
            logger.error(f"Error saving alert config: {e}")
    
    def check_thresholds(self, gpu_name: str, status: GPUStatus):
        """Check all thresholds against current status."""
        current_time = time.time()
        
        # Check each threshold
        for threshold in self.thresholds:
            if not threshold.enabled:
                continue
            
            # Get metric value
            metric_value = self.get_metric_value(status, threshold.metric)
            if metric_value is None:
                continue
            
            # Check if threshold is exceeded
            if self.check_operator(metric_value, threshold.operator, threshold.threshold):
                # Start or continue threshold state
                state_key = (gpu_name, threshold.metric, threshold.threshold)
                
                if state_key not in self.threshold_states:
                    self.threshold_states[state_key] = {
                        'value': metric_value,
                        'start_time': current_time
                    }
                else:
                    self.threshold_states[state_key]['value'] = metric_value
                
                # Check if duration threshold is met
                state = self.threshold_states[state_key]
                duration = current_time - state['start_time']
                
                if duration >= threshold.duration:
                    # Check cooldown
                    last_alert_key = state_key
                    if last_alert_key in self.last_alert_times:
                        time_since_last = current_time - self.last_alert_times[last_alert_key]
                        if time_since_last < threshold.cooldown:
                            continue
                    
                    # Create alert
                    alert = self.create_alert(gpu_name, threshold, metric_value)
                    self.trigger_alert(alert)
                    
                    # Update cooldown
                    self.last_alert_times[last_alert_key] = current_time
            
            else:
                # Reset threshold state if condition is no longer met
                state_key = (gpu_name, threshold.metric, threshold.threshold)
                if state_key in self.threshold_states:
                    del self.threshold_states[state_key]
    
    def get_metric_value(self, status: GPUStatus, metric: str) -> Optional[float]:
        """Get metric value from GPU status."""
        if metric == 'temperature':
            return status.temperature
        elif metric == 'load':
            return status.load
        elif metric == 'power_draw':
            return status.power_draw
        elif metric == 'fan_speed':
            return (status.fan_pwm / 255) * 100  # Convert PWM to percentage
        elif metric == 'memory_usage':
            if status.memory_total > 0:
                return (status.memory_used / status.memory_total) * 100
            return 0
        elif metric == 'efficiency':
            return status.efficiency
        else:
            return None
    
    def check_operator(self, value: float, operator: str, threshold: float) -> bool:
        """Check if value meets threshold condition."""
        if operator == '>':
            return value > threshold
        elif operator == '<':
            return value < threshold
        elif operator == '>=':
            return value >= threshold
        elif operator == '<=':
            return value <= threshold
        elif operator == '==':
            return value == threshold
        else:
            return False
    
    def create_alert(self, gpu_name: str, threshold: AlertThreshold, value: float) -> Alert:
        """Create an alert instance."""
        alert_id = f"{gpu_name}_{threshold.metric}_{threshold.threshold}_{int(time.time())}"
        
        # Determine severity
        severity = 'info'
        if threshold.metric == 'temperature':
            if threshold.threshold >= 85:
                severity = 'emergency'
            elif threshold.threshold >= 75:
                severity = 'critical'
            else:
                severity = 'warning'
        elif threshold.metric == 'load':
            severity = 'warning'
        elif threshold.metric == 'power_draw':
            severity = 'critical'
        elif threshold.metric == 'fan_speed':
            severity = 'warning'
        
        message = f"{threshold.metric} ({value}) exceeded threshold ({threshold.threshold}) for {threshold.duration}s"
        
        alert = Alert(
            id=alert_id,
            timestamp=time.time(),
            gpu_name=gpu_name,
            metric=threshold.metric,
            value=value,
            threshold=threshold.threshold,
            message=message,
            severity=severity
        )
        
        return alert
    
    def trigger_alert(self, alert: Alert):
        """Trigger an alert and send notifications."""
        # Store alert in database
        self.db_manager.save_alert(alert)
        
        # Add to active alerts
        self.active_alerts[alert.id] = alert
        
        # Send notifications
        for channel in self.channels:
            if channel.enabled:
                try:
                    success = channel.send(alert)
                    if not success:
                        logger.warning(f"Failed to send alert via {channel.name}")
                except Exception as e:
                    logger.error(f"Error sending alert via {channel.name}: {e}")
        
        logger.info(f"Alert triggered: {alert.message}")
    
    def resolve_alert(self, alert_id: str):
        """Resolve an active alert."""
        if alert_id in self.active_alerts:
            alert = self.active_alerts[alert_id]
            alert.resolved = True
            alert.resolved_at = time.time()
            
            # Remove from active alerts
            del self.active_alerts[alert_id]
            
            logger.info(f"Alert resolved: {alert.message}")
    
    def get_active_alerts(self) -> List[Alert]:
        """Get list of active alerts."""
        return list(self.active_alerts.values())
    
    def get_alert_history(self, hours: int = 24) -> List[Alert]:
        """Get alert history from database."""
        try:
            cutoff_time = time.time() - (hours * 3600)
            
            with sqlite3.connect(self.db_manager.db_path) as conn:
                cursor = conn.cursor()
                
                cursor.execute('''
                    SELECT id, timestamp, gpu_name, metric, value, threshold, message, severity,
                           acknowledged, resolved, resolved_at
                    FROM alerts 
                    WHERE timestamp >= ?
                    ORDER BY timestamp DESC
                ''', (cutoff_time,))
                
                rows = cursor.fetchall()
                
                if not rows:
                    return []
                
                alerts = []
                for row in rows:
                    alert = Alert(
                        id=row[0],
                        timestamp=row[1],
                        gpu_name=row[2],
                        metric=row[3],
                        value=row[4],
                        threshold=row[5],
                        message=row[6],
                        severity=row[7],
                        acknowledged=bool(row[8]),
                        resolved=bool(row[9]),
                        resolved_at=row[10]
                    )
                    alerts.append(alert)
                
                return alerts
                
        except Exception as e:
            logger.error(f"Error getting alert history: {e}")
            return []
    
    def cleanup_old_alerts(self):
        """Clean up old alerts and threshold states."""
        current_time = time.time()
        max_age = self.alert_settings.get('max_alert_age', 86400.0)
        
        # Clean up old active alerts
        old_alerts = []
        for alert_id, alert in self.active_alerts.items():
            if current_time - alert.timestamp > max_age:
                old_alerts.append(alert_id)
        
        for alert_id in old_alerts:
            del self.active_alerts[alert_id]
        
        # Clean up old threshold states
        old_states = []
        for state_key, state in self.threshold_states.items():
            if current_time - state['start_time'] > max_age:
                old_states.append(state_key)
        
        for state_key in old_states:
            del self.threshold_states[state_key]
        
        # Clean up old last alert times
        old_cooldowns = []
        for key, timestamp in self.last_alert_times.items():
            if current_time - timestamp > max_age:
                old_cooldowns.append(key)
        
        for key in old_cooldowns:
            del self.last_alert_times[key]
    
    def start(self):
        """Start the alert manager."""
        if self.running:
            return
        
        self.running = True
        self.thread = threading.Thread(target=self.run, daemon=True)
        self.thread.start()
        logger.info("Alert manager started")
    
    def stop(self):
        """Stop the alert manager."""
        self.running = False
        if self.thread:
            self.thread.join()
        logger.info("Alert manager stopped")
    
    def run(self):
        """Main alert manager loop."""
        check_interval = self.alert_settings.get('check_interval', 5.0)
        cleanup_interval = self.alert_settings.get('cleanup_interval', 3600.0)
        last_cleanup = time.time()
        
        while self.running:
            try:
                # Get current GPU status
                gpu_manager = GPUManager()
                if gpu_manager.initialize():
                    status_dict = gpu_manager.get_status()
                    
                    # Check thresholds for each GPU
                    for gpu_name, status in status_dict.items():
                        if status:
                            self.check_thresholds(gpu_name, status)
                
                # Cleanup old data periodically
                current_time = time.time()
                if current_time - last_cleanup >= cleanup_interval:
                    self.cleanup_old_alerts()
                    last_cleanup = current_time
                
                time.sleep(check_interval)
                
            except Exception as e:
                logger.error(f"Error in alert manager loop: {e}")
                time.sleep(5)  # Wait before retrying


class AlertAPI:
    """API interface for alert management."""
    
    def __init__(self, alert_manager: AlertManager):
        self.alert_manager = alert_manager
    
    def get_active_alerts(self) -> List[Dict[str, Any]]:
        """Get active alerts."""
        alerts = self.alert_manager.get_active_alerts()
        return [asdict(alert) for alert in alerts]
    
    def get_alert_history(self, hours: int = 24) -> List[Dict[str, Any]]:
        """Get alert history."""
        alerts = self.alert_manager.get_alert_history(hours)
        return [asdict(alert) for alert in alerts]
    
    def acknowledge_alert(self, alert_id: str) -> bool:
        """Acknowledge an alert."""
        if alert_id in self.alert_manager.active_alerts:
            self.alert_manager.active_alerts[alert_id].acknowledged = True
            return True
        return False
    
    def resolve_alert(self, alert_id: str) -> bool:
        """Resolve an alert."""
        self.alert_manager.resolve_alert(alert_id)
        return True
    
    def add_threshold(self, threshold_data: Dict[str, Any]) -> bool:
        """Add a new threshold."""
        try:
            threshold = AlertThreshold(**threshold_data)
            self.alert_manager.add_threshold(threshold)
            return True
        except Exception as e:
            logger.error(f"Error adding threshold: {e}")
            return False
    
    def remove_threshold(self, metric: str, threshold_value: float) -> bool:
        """Remove a threshold."""
        try:
            self.alert_manager.remove_threshold(metric, threshold_value)
            return True
        except Exception as e:
            logger.error(f"Error removing threshold: {e}")
            return False


if __name__ == "__main__":
    # Test alert system
    logging.basicConfig(level=logging.INFO)
    
    alert_manager = AlertManager()
    alert_manager.start()
    
    try:
        while True:
            time.sleep(1)
    except KeyboardInterrupt:
        alert_manager.stop()