import threading
import time
from typing import Optional
import xml.etree.ElementTree as ET
import psutil

class SystemMetrics:
    """
    Tracks system metrics including CPU, GPU, memory and disk usage.
    Maintains rolling averages for CPU and GPU between context generations.
    
    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 context
        _gpu_samples: List of GPU usage samples since last context
        _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 generate_context(self, context_usage: float, stdin_buffer_length: int) -> ET.Element:
        """
        Generate XML context element with current system metrics.
        Clears usage samples after generating averages.
        
        Args:
            context_usage: Current context size usage (0-1)
            
        Returns:
            ET.Element: Context element with system metrics
        """
        # Create context element
        context = ET.Element("context")
        
        # Add timestamp
        context.set("time", time.strftime("%Y-%m-%dT%H:%M:%SZ", time.gmtime()))
        
        # Add usage metrics
        with self._metrics_lock:
            # CPU average (or 0 if no samples)
            cpu_avg = round(sum(self._cpu_samples) / len(self._cpu_samples)) if self._cpu_samples else 0
            context.set("cpu", str(cpu_avg))
            self._cpu_samples.clear()
            
            # GPU average (or 0 if no samples)
            gpu_avg = round(sum(self._gpu_samples) / len(self._gpu_samples)) if self._gpu_samples else 0
            context.set("gpu", str(gpu_avg))
            self._gpu_samples.clear()
        
        # Memory usage in bytes
        memory = psutil.virtual_memory()
        context.set("memory_used", str(memory.used))
        context.set("memory_total", str(memory.total))
        
        # Disk usage in bytes (root partition)
        disk = psutil.disk_usage('/')
        context.set("disk_used", str(disk.used))
        context.set("disk_total", str(disk.total))
        
        # Context usage (0-100)
        context.set("context_usage", str(round(context_usage * 100)))
        
        # Standard input buffer size 
        context.set("stdin", stdin_buffer_length)
        
        return context
        
    def stop(self):
        """Stop the monitoring thread and clean up."""
        self._stop_event.set()
        self._monitor_thread.join()