monitor_engine.py

import asyncio
import logging
from datetime import datetime, timedelta
from typing import Dict, List, Optional, Set, Callable, Any
from dataclasses import dataclass, field
from enum import Enum
import json
import time
from concurrent.futures import ThreadPoolExecutor
import threading
from collections import defaultdict, deque

from config import get_config, MonitoringConfig
from data_models import (
    SpaceInfo, SpaceStatusInfo, SpaceStatus, MonitorEvent, EventType,
    AlertLevel, DatabaseManager, BuildLog, WebhookEvent, Alert, AlertRule
)
from huggingface_client_v2 import HuggingFaceClient, RetryClient, WebhookHandler


class MonitorState(Enum):
    MONITOR_STATE = "monitor_state"
    STARTING = "starting"
    RUNNING = "running"
    PAUSED = "paused"
    STOPPING = "stopping"
    STOPPED = "stopped"
    ERROR = "error"


@dataclass
class SpaceMonitor:
    space_id: str
    config: Dict[str, Any]
    last_check: Optional[datetime] = None
    last_status: Optional[SpaceStatus] = None
    error_count: int = 0
    consecutive_errors: int = 0
    events: List[MonitorEvent] = field(default_factory=list)
    active: bool = True
    task: Optional[asyncio.Task] = None


class MonitorEngine:
    HuggingFace Spaces 监控引擎
    
    def __init__(self, config: Optional[MonitoringConfig] = None):
        self.config = config or get_config().monitoring
        self.logger = logging.getLogger(__name__)
        
        self.client = RetryClient(HuggingFaceClient())
        self.db_manager = DatabaseManager()
        self.webhook_handler = None
        
        self.monitored_spaces: Dict[str, SpaceMonitor] = {}
        self.state = MonitorState.STOPPED
        self.monitor_task: Optional[asyncio.Task] = None
        self.event_queue = asyncio.Queue()
        
        self.event_callbacks: Dict[EventType, List[Callable]] = defaultdict(list)
        self.alert_rules: Dict[str, AlertRule] = {}
        
        self._shutdown_event = asyncio.Event()
        self._executor = ThreadPoolExecutor(max_workers=4)
        
        self.stats = {
            'total_checks': 0,
            'successful_checks': 0,
            'failed_checks': 0,
            'events_generated': 0,
            'alerts_triggered': 0,
            'start_time': None,
            'last_check_time': None
        }
    
    async def start(self) -> None:
        if self.state != MonitorState.STOPPED:
            self.logger.warning("监控引擎已在运行中")
            return
        
        try:
            self.state = MonitorState.STARTING
            self.logger.info("启动监控引擎...")
            
            await self.client.client.validate_token()
            
            self._shutdown_event.clear()
            self.monitor_task = asyncio.create_task(self._monitor_loop())
            
            self.state = MonitorState.RUNNING
            self.stats['start_time'] = datetime.now()
            
            self.logger.info("监控引擎启动成功")
            
        except Exception as e:
            self.state = MonitorState.ERROR
            self.logger.error(f"启动监控引擎失败: {e}")
            raise
    
    async def stop(self) -> None:
        if self.state == MonitorState.STOPPED:
            return
        
        try:
            self.state = MonitorState.STOPPING
            self.logger.info("停止监控引擎...")
            
            self._shutdown_event.set()
            
            if self.monitor_task:
                self.monitor_task.cancel()
                try:
                    await self.monitor_task
                except asyncio.CancelledError:
                    pass
            
            for space_monitor in self.monitored_spaces.values():
                if space_monitor.task:
                    space_monitor.task.cancel()
                    try:
                        await space_monitor.task
                    except asyncio.CancelledError:
                        pass
            
            await self.client.client.close()
            
            self._executor.shutdown(wait=True)
            
            self.state = MonitorState.STOPPED
            self.logger.info("监控引擎已停止")
            
        except Exception as e:
            self.state = MonitorState.ERROR
            self.logger.error(f"停止监控引擎失败: {e}")
            raise
    
    async def add_space(self, space_id: str, config: Optional[Dict[str, Any]] = None) -> None:
        if space_id in self.monitored_spaces:
            self.logger.warning(f"Space {space_id} 已在监控中")
            return
        
        monitor_config = config or self.config.model_dump()
        
        try:
            space_info = await self.client.client.get_space_info(space_id)
            initial_status = await self.client.get_space_status(space_id)
            
            space_monitor = SpaceMonitor(
                space_id=space_id,
                config=monitor_config,
                last_check=datetime.now(),
                last_status=initial_status.status
            )
            
            self.monitored_spaces[space_id] = space_monitor
            await self.db_manager.save_space_info(space_info)
            
            event = MonitorEvent(
                space_id=space_id,
                event_type=EventType.BUILD_STARTED,
                timestamp=datetime.now(),
                message=f"开始监控 Space {space_id}，当前状态: {initial_status.status.value}",
                current_status=initial_status.status
            )
            
            await self._emit_event(event)
            
            self.logger.info(f"已添加 Space {space_id} 到监控列表")
            
        except Exception as e:
            self.logger.error(f"添加 Space {space_id} 失败: {e}")
            raise
    
    async def remove_space(self, space_id: str) -> None:
        if space_id not in self.monitored_spaces:
            self.logger.warning(f"Space {space_id} 不在监控中")
            return
        
        space_monitor = self.monitored_spaces[space_id]
        
        if space_monitor.task:
            space_monitor.task.cancel()
            try:
                await space_monitor.task
            except asyncio.CancelledError:
                pass
        
        del self.monitored_spaces[space_id]
        
        event = MonitorEvent(
            space_id=space_id,
            event_type=EventType.SPACE_STOPPED,
            timestamp=datetime.now(),
            message=f"停止监控 Space {space_id}"
        )
        
        await self._emit_event(event)
        
        self.logger.info(f"已从监控列表移除 Space {space_id}")
    
    async def get_monitored_spaces(self) -> List[str]:
        return list(self.monitored_spaces.keys())
    
    async def get_space_status(self, space_id: str) -> Optional[SpaceStatusInfo]:
        if space_id not in self.monitored_spaces:
            return None
        
        return await self.client.get_space_status(space_id)
    
    async def get_space_events(self, space_id: str, limit: int = 100) -> List[MonitorEvent]:
        return await self.db_manager.get_recent_events(space_id, limit)
    
    async def add_alert_rule(self, rule: AlertRule) -> None:
        self.alert_rules[rule.rule_id] = rule
        
        await self.db_manager.save_alert(Alert(
            rule_id=rule.rule_id,
            space_id=rule.space_id or "*",
            severity=rule.severity,
            title=f"告警规则创建: {rule.name}",
            message=f"告警规则 '{rule.name}' 已创建",
            timestamp=datetime.now()
        ))
        
        self.logger.info(f"已添加告警规则: {rule.name}")
    
    async def remove_alert_rule(self, rule_id: str) -> None:
        if rule_id in self.alert_rules:
            del self.alert_rules[rule_id]
            self.logger.info(f"已移除告警规则: {rule_id}")
    
    def register_event_callback(self, event_type: EventType, callback: Callable) -> None:
        self.event_callbacks[event_type].append(callback)
    
    def unregister_event_callback(self, event_type: EventType, callback: Callable) -> None:
        if callback in self.event_callbacks[event_type]:
            self.event_callbacks[event_type].remove(callback)
    
    async def _monitor_loop(self) -> None:
        self.logger.info("监控循环开始")
        
        while not self._shutdown_event.is_set():
            try:
                if self.monitored_spaces:
                    await self._check_all_spaces()
                
                await asyncio.sleep(1)
                
                if self._shutdown_event.is_set():
                    break
                    
            except asyncio.CancelledError:
                break
            except Exception as e:
                self.logger.error(f"监控循环异常: {e}")
                await asyncio.sleep(5)
        
        self.logger.info("监控循环结束")
    
    async def _check_all_spaces(self) -> None:
        tasks = []
        
        for space_id, monitor in self.monitored_spaces.items():
            if not monitor.active:
                continue
            
            interval = monitor.config.get('check_interval_seconds', self.config.default_check_interval)
            
            if monitor.last_check and (datetime.now() - monitor.last_check).total_seconds() < interval:
                continue
            
            task = asyncio.create_task(self._check_space(space_id, monitor))
            tasks.append(task)
        
        if tasks:
            results = await asyncio.gather(*tasks, return_exceptions=True)
            
            for i, result in enumerate(results):
                if isinstance(result, Exception):
                    space_id = list(self.monitored_spaces.keys())[i]
                    self.logger.error(f"检查 Space {space_id} 失败: {result}")
    
    async def _check_space(self, space_id: str, monitor: SpaceMonitor) -> None:
        try:
            self.stats['total_checks'] += 1
            
            status_info = await self.client.get_space_status(space_id)
            monitor.last_check = datetime.now()
            self.stats['last_check_time'] = monitor.last_check
            
            status_changed = monitor.last_status != status_info.status
            
            if status_changed:
                await self._handle_status_change(space_id, monitor.last_status, status_info)
                monitor.last_status = status_info.status
            
            if status_info.status == SpaceStatus.ERROR:
                monitor.error_count += 1
                monitor.consecutive_errors += 1
                
                await self._handle_error_state(space_id, status_info, monitor)
                
                threshold = monitor.config.get('error_threshold', self.config.error_threshold)
                if monitor.consecutive_errors >= threshold:
                    await self._trigger_error_alert(space_id, monitor)
            else:
                monitor.consecutive_errors = 0
            
            await self.db_manager.save_status_history(status_info)
            
            self.stats['successful_checks'] += 1
            
        except Exception as e:
            self.stats['failed_checks'] += 1
            monitor.consecutive_errors += 1
            
            self.logger.error(f"检查 Space {space_id} 异常: {e}")
            
            event = MonitorEvent(
                space_id=space_id,
                event_type=EventType.ERROR_DETECTED,
                timestamp=datetime.now(),
                message=f"检查失败: {str(e)}",
                severity=AlertLevel.HIGH
            )
            
            await self._emit_event(event)
    
    async def _handle_status_change(self, space_id: str, old_status: Optional[SpaceStatus], new_status: SpaceStatusInfo) -> None:
        event_type = None
        message = f"状态变化: {old_status.value if old_status else 'UNKNOWN'} → {new_status.status.value}"
        
        if new_status.status == SpaceStatus.BUILDING:
            event_type = EventType.BUILD_STARTED
        elif new_status.status == SpaceStatus.RUNNING:
            event_type = EventType.SPACE_STARTED
        elif new_status.status == SpaceStatus.STOPPED:
            event_type = EventType.SPACE_STOPPED
        elif new_status.status == SpaceStatus.ERROR:
            event_type = EventType.BUILD_FAILED
        
        if event_type:
            event = MonitorEvent(
                space_id=space_id,
                event_type=event_type,
                timestamp=datetime.now(),
                previous_status=old_status,
                current_status=new_status.status,
                message=message
            )
            
            await self._emit_event(event)
    
    async def _handle_error_state(self, space_id: str, status_info: SpaceStatusInfo, monitor: SpaceMonitor) -> None:
        try:
            logs = await self.client.get_space_logs(space_id, lines=monitor.config.get('log_lines_count', 50))
            error_entries = [entry for entry in logs.entries if 'error' in entry.message.lower()]
            
            if error_entries:
                latest_error = error_entries[-1].message
                
                event = MonitorEvent(
                    space_id=space_id,
                    event_type=EventType.ERROR_DETECTED,
                    timestamp=datetime.now(),
                    current_status=SpaceStatus.ERROR,
                    message=f"检测到错误: {latest_error[:200]}...",
                    severity=AlertLevel.HIGH,
                    data={'error_logs': [entry.message for entry in error_entries[-5:]]}
                )
                
                await self._emit_event(event)
        
        except Exception as e:
            self.logger.error(f"处理错误状态失败: {e}")
    
    async def _trigger_error_alert(self, space_id: str, monitor: SpaceMonitor) -> None:
        alert = Alert(
            rule_id="auto_error_threshold",
            space_id=space_id,
            severity=AlertLevel.CRITICAL,
            title=f"Space {space_id} 连续错误",
            message=f"Space {space_id} 连续 {monitor.consecutive_errors} 次检查失败，超过阈值 {monitor.config.get('error_threshold', self.config.error_threshold)}",
            timestamp=datetime.now()
        )
        
        await self.db_manager.save_alert(alert)
        self.stats['alerts_triggered'] += 1
        
        self.logger.critical(f"触发严重告警: {alert.title}")
    
    async def _emit_event(self, event: MonitorEvent) -> None:
        self.stats['events_generated'] += 1
        
        await self.db_manager.save_monitor_event(event)
        await self.event_queue.put(event)
        
        for callback in self.event_callbacks[event.event_type]:
            try:
                if asyncio.iscoroutinefunction(callback):
                    await callback(event)
                else:
                    await self._executor.submit(callback, event)
            except Exception as e:
                self.logger.error(f"事件回调执行失败: {e}")
        
        for rule_id, rule in self.alert_rules.items():
            if await self._should_trigger_alert(rule, event):
                await self._create_alert(rule, event)
    
    async def _should_trigger_alert(self, rule: AlertRule, event: MonitorEvent) -> bool:
        if not rule.enabled:
            return False
        
        if rule.space_id and rule.space_id != event.space_id:
            return False
        
        if rule.last_triggered:
            cooldown = timedelta(minutes=rule.cooldown_minutes)
            if datetime.now() - rule.last_triggered < cooldown:
                return False
        
        condition = rule.condition
        
        if 'event_type' in condition and condition['event_type'] != event.event_type.value:
            return False
        
        if 'severity' in condition:
            required_severity = AlertLevel(condition['severity'])
            if event.severity.value < required_severity.value:
                return False
        
        return True
    
    async def _create_alert(self, rule: AlertRule, event: MonitorEvent) -> None:
        alert = Alert(
            rule_id=rule.rule_id,
            space_id=event.space_id,
            severity=rule.severity,
            title=f"告警: {rule.name}",
            message=f"{rule.description or ''} - {event.message}",
            timestamp=datetime.now(),
            metadata={'event_id': event.event_id, 'rule_name': rule.name}
        )
        
        await self.db_manager.save_alert(alert)
        self.stats['alerts_triggered'] += 1
        
        rule.last_triggered = datetime.now()
        
        self.logger.warning(f"触发告警: {alert.title}")
    
    async def get_stats(self) -> Dict[str, Any]:
        uptime = None
        if self.stats['start_time']:
            uptime = (datetime.now() - self.stats['start_time']).total_seconds()
        
        return {
            **self.stats,
            'state': self.state.value,
            'monitored_spaces_count': len(self.monitored_spaces),
            'active_spaces_count': sum(1 for m in self.monitored_spaces.values() if m.active),
            'uptime_seconds': uptime,
            'alerts_rules_count': len(self.alert_rules)
        }
    
    async def set_webhook_handler(self, handler: WebhookHandler) -> None:
        self.webhook_handler = handler
    
    async def handle_webhook(self, payload: Dict[str, Any], headers: Dict[str, str]) -> WebhookEvent:
        if not self.webhook_handler:
            raise ValueError("Webhook 处理器未设置")
        
        webhook_event = await self.webhook_handler.handle_webhook(payload, headers)
        await self.db_manager.save_webhook_event(webhook_event)
        
        space_id = webhook_event.space_id
        if space_id in self.monitored_spaces:
            await self._check_space(space_id, self.monitored_spaces[space_id])
        
        return webhook_event
    
    async def pause_monitoring(self, space_id: Optional[str] = None) -> None:
        if space_id:
            if space_id in self.monitored_spaces:
                self.monitored_spaces[space_id].active = False
                self.logger.info(f"已暂停监控 Space {space_id}")
        else:
            self.state = MonitorState.PAUSED
            for monitor in self.monitored_spaces.values():
                monitor.active = False
            self.logger.info("已暂停所有监控")
    
    async def resume_monitoring(self, space_id: Optional[str] = None) -> None:
        if space_id:
            if space_id in self.monitored_spaces:
                self.monitored_spaces[space_id].active = True
                self.logger.info(f"已恢复监控 Space {space_id}")
        else:
            self.state = MonitorState.RUNNING
            for monitor in self.monitored_spaces.values():
                monitor.active = True
            self.logger.info("已恢复所有监控")


class HealthChecker:
    健康检查器
    
    def __init__(self, engine: MonitorEngine):
        self.engine = engine
        self.logger = logging.getLogger(__name__)
    
    async def check_health(self) -> Dict[str, Any]:
        health_status = {
            'status': 'healthy',
            'timestamp': datetime.now().isoformat(),
            'checks': {}
        }
        
        try:
            stats = await self.engine.get_stats()
            health_status['checks']['engine'] = {
                'status': 'healthy' if stats['state'] == 'running' else 'unhealthy',
                'details': stats
            }
            
            token_valid = await self.engine.client.client.validate_token()
            health_status['checks']['token'] = {
                'status': 'healthy' if token_valid else 'unhealthy',
                'details': {'valid': token_valid}
            }
            
            try:
                await self.engine.db_manager._init_database()
                health_status['checks']['database'] = {
                    'status': 'healthy',
                    'details': {'connection': 'ok'}
                }
            except Exception as e:
                health_status['checks']['database'] = {
                    'status': 'unhealthy',
                    'details': {'error': str(e)}
                }
            
            overall_status = 'healthy'
            for check in health_status['checks'].values():
                if check['status'] != 'healthy':
                    overall_status = 'unhealthy'
                    break
            
            health_status['status'] = overall_status
            
        except Exception as e:
            health_status['status'] = 'unhealthy'
            health_status['error'] = str(e)
        
        return health_status