"""
Performance Monitoring and Instrumentation for AEGIS-AI.
Tracks timing and resource usage for various pipeline stages.
"""
import os
import time
import logging
from typing import Dict, Optional, Any
from contextlib import contextmanager

try:
    import psutil
    HAS_PSUTIL = True
except ImportError:
    HAS_PSUTIL = False

logger = logging.getLogger(__name__)

class PerformanceMonitor:
    """
    Monitors timing and memory usage across different stages of execution.
    """
    def __init__(self):
        """Initialize the performance monitor."""
        self.timings: Dict[str, float] = {}
        self._start_times: Dict[str, float] = {}
        
        # Valid stages defined in implementation plan
        self.VALID_STAGES = {
            'upload_validation',
            'face_detection',
            'model_inference',
            'heatmap_generation',
            'audio_processing',
            'metadata_analysis',
            'fusion_engine',
            'total_analysis',
            'model_load_time'
        }
        
    def start(self, stage: str) -> None:
        """
        Start timing a specific stage.
        
        Args:
            stage (str): Name of the stage to start timing.
        """
        if stage not in self.VALID_STAGES:
            logger.warning(f"Unknown performance stage started: {stage}")
            
        self._start_times[stage] = time.perf_counter()
        logger.debug(f"Started timing stage: {stage}")
        
    def stop(self, stage: str) -> None:
        """
        Stop timing a specific stage and record the duration.
        
        Args:
            stage (str): Name of the stage to stop timing.
        """
        if stage not in self._start_times:
            logger.warning(f"Attempted to stop unstarted stage: {stage}")
            return
            
        duration = time.perf_counter() - self._start_times.pop(stage)
        duration_ms = duration * 1000.0
        
        # Accumulate time if stage is run multiple times
        if stage in self.timings:
            self.timings[stage] += duration_ms
        else:
            self.timings[stage] = duration_ms
            
        logger.debug(f"Completed stage {stage} in {duration_ms:.2f}ms")
        
    def get_peak_ram_mb(self) -> float:
        """
        Get the current RAM usage of the process in MB.
        Note: This is current RAM, which serves as a proxy for peak if called at end.
        
        Returns:
            float: RAM usage in MB.
        """
        if not HAS_PSUTIL:
            return 0.0
        process = psutil.Process(os.getpid())
        return process.memory_info().rss / (1024 * 1024)
        
    def get_report(self) -> Dict[str, Any]:
        """
        Get a full report of all timings and peak RAM.
        
        Returns:
            Dict[str, Any]: Dictionary containing timings (ms) and memory stats.
        """
        report = {k: round(v, 2) for k, v in self.timings.items()}
        report['peak_ram_mb'] = round(self.get_peak_ram_mb(), 2)
        return report
        
    @contextmanager
    def track(self, stage: str):
        """
        Context manager for timing a block of code.
        
        Args:
            stage (str): Name of the stage to track.
            
        Yields:
            None
        """
        self.start(stage)
        try:
            yield
        finally:
            self.stop(stage)