import asyncio
import psutil
import pynvml
import time
from typing import Dict, Any

async def get_system_metrics(args: Dict[str, Any]) -> Dict[str, Any]:
    """
    Retrieves system metrics including CPU usage, RAM usage, GPU usage (for NVIDIA GPUs),
    system uptime, and network traffic of the primary network interface.

    Args:
        args: Dictionary of arguments (not used in this implementation).

    Returns:
        Dictionary with CPU usage (%), RAM used (GB), RAM total (GB), GPU usage (%),
        system uptime (seconds), network traffic (Mbit/s), and input arguments.

    Raises:
        OSError: If retrieving metrics fails.
    """
    try:
        # Get CPU usage (averaged over 0.1 seconds for responsiveness)
        cpu_usage = await asyncio.to_thread(psutil.cpu_percent, interval=1)

        # Get RAM usage
        memory = psutil.virtual_memory()
        ram_used = memory.used / (1024 ** 3)  # Convert bytes to GB
        ram_total = memory.total / (1024 ** 3)  # Convert bytes to GB

        # Get system uptime
        uptime = time.time() - psutil.boot_time()  # Time since system boot in seconds

        # Get network traffic for primary interface
        network_traffic = 0.0
        net_io = psutil.net_io_counters(pernic=True)
        primary_interface = None
        max_bytes = 0

        # Find primary network interface (with most bytes sent/received)
        for interface, stats in net_io.items():
            total_bytes = stats.bytes_sent + stats.bytes_recv
            if total_bytes > max_bytes:
                max_bytes = total_bytes
                primary_interface = interface

        if primary_interface:
            # Measure traffic over 1 second
            initial_stats = net_io[primary_interface]
            await asyncio.sleep(1)
            final_stats = psutil.net_io_counters(pernic=True)[primary_interface]
            
            # Calculate traffic in Mbit/s
            bytes_diff = (final_stats.bytes_sent - initial_stats.bytes_sent +
                         final_stats.bytes_recv - initial_stats.bytes_recv)
            network_traffic = (bytes_diff * 8) / (1024 ** 2)  # Convert bytes to Mbit/s

        # Try to get GPU usage (NVIDIA only)
        gpu_usage = 0.0
        try:
            pynvml.nvmlInit()
            device_count = pynvml.nvmlDeviceGetCount()
            if device_count > 0:
                handle = pynvml.nvmlDeviceGetHandleByIndex(0)  # Use first GPU
                utilization = pynvml.nvmlDeviceGetUtilizationRates(handle)
                gpu_usage = utilization.gpu
            pynvml.nvmlShutdown()
        except pynvml.NVMLError:
            gpu_usage = -1  # Indicate GPU metrics unavailable

        return {
            "message": "System metrics retrieved successfully",
            "metrics": {
                "cpu_usage_percent": round(cpu_usage, 2),
                "ram_used_gb": round(ram_used, 2),
                "ram_total_gb": round(ram_total, 2),
                "gpu_usage_percent": gpu_usage,
                "uptime_seconds": round(uptime, 2),
                "network_traffic_mbps": round(network_traffic, 2)
            },
            "args": args
        }

    except Exception as e:
        raise OSError(f"Error while retrieving system metrics: {str(e)}")