import threading
import time
from typing import Optional, Dict
import psutil

class SystemMetrics:
    """
    Tracks system metrics including CPU, GPU, memory and disk usage.
    Maintains rolling averages for CPU and GPU.
    
    Attributes:
        _stop_event: Threading event to signal monitor thread to stop
        _monitor_thread: Background thread that collects usage data
        _cpu_samples: List of CPU usage samples since last reading
        _gpu_samples: List of GPU usage samples since last reading
        _metrics_lock: Lock for thread-safe access to samples
    """
    
    def __init__(self, sample_interval: float = 0.1):
        """
        Initialize system metrics monitoring.
        
        Args:
            sample_interval: Seconds between usage samples
        """
        self._stop_event = threading.Event()
        self._cpu_samples = []
        self._gpu_samples = []
        self._metrics_lock = threading.Lock()
        self._sample_interval = sample_interval
        
        # Start monitoring thread
        self._monitor_thread = threading.Thread(
            target=self._monitor_loop,
            daemon=True
        )
        self._monitor_thread.start()
        
    def _get_gpu_usage(self) -> Optional[float]:
        """
        Get current GPU usage percentage.
        Returns None if GPU monitoring not available.
        """
        try:
            import pynvml
            pynvml.nvmlInit()
            handle = pynvml.nvmlDeviceGetHandleByIndex(0)
            info = pynvml.nvmlDeviceGetUtilizationRates(handle)
            pynvml.nvmlShutdown()
            return info.gpu
        except:
            return None

    def _monitor_loop(self):
        """
        Background loop that collects system usage samples.
        Runs until stop event is set.
        """
        while not self._stop_event.is_set():
            with self._metrics_lock:
                # Get CPU usage across all cores
                self._cpu_samples.append(psutil.cpu_percent())
                
                # Get GPU usage if available
                gpu_usage = self._get_gpu_usage()
                if gpu_usage is not None:
                    self._gpu_samples.append(gpu_usage)
                    
            time.sleep(self._sample_interval)

    def get_metrics(self) -> Dict:
        """
        Get current system metrics.
        Clears usage samples after calculating averages.
        
        Returns:
            Dict containing:
                - cpu: Average CPU usage percentage
                - gpu: Average GPU usage percentage
                - memory_used: Used memory in bytes
                - memory_total: Total memory in bytes
                - disk_used: Used disk space in bytes
                - disk_total: Total disk space in bytes
                - timestamp: Current UTC timestamp in ISO format
        """
        metrics = {}
        
        # Add timestamp
        metrics["timestamp"] = time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime())
        
        # Add usage metrics
        with self._metrics_lock:
            # CPU average (or 0 if no samples)
            metrics["cpu"] = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
            self._cpu_samples.clear()
            
            # GPU average (or 0 if no samples)
            metrics["gpu"] = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
            self._gpu_samples.clear()
        
        # Memory usage in bytes
        memory = psutil.virtual_memory()
        metrics["memory_used"] = memory.used
        metrics["memory_total"] = memory.total
        
        # Disk usage in bytes (root partition)
        disk = psutil.disk_usage('/')
        metrics["disk_used"] = disk.used
        metrics["disk_total"] = disk.total
        
        return metrics
        
    def stop(self):
        """Stop the monitoring thread and clean up."""
        self._stop_event.set()
        self._monitor_thread.join()