"""Utilities for capturing environment information."""

import platform
import sys


def capture_environment() -> dict[str, object]:
    """Capture current environment information.

    Returns a dictionary with:
    - python_version: Python version string
    - platform: Operating system platform
    - cpu_count: Number of CPU cores
    - ram_gb: Total RAM in GB (if available)
    - gpu_count: Number of GPUs (if detectable)
    - gpu_names: List of GPU names (if detectable)

    Returns:
        Dictionary containing environment information
    """
    import os

    env_info: dict[str, object] = {
        "python_version": sys.version.split()[0],  # e.g., "3.12.1"
        "platform": platform.platform(),  # e.g., "Linux-5.15.0-91-generic-x86_64-with-glibc2.35"
        "cpu_count": os.cpu_count(),
    }

    # Try to get RAM size
    try:
        import psutil  # type: ignore[import-untyped]

        env_info["ram_gb"] = round(psutil.virtual_memory().total / (1024**3), 2)
    except ImportError:
        pass  # psutil not available

    # Try to detect GPUs
    try:
        import torch  # type: ignore[import-not-found]

        if torch.cuda.is_available():
            env_info["gpu_count"] = torch.cuda.device_count()
            env_info["gpu_names"] = [torch.cuda.get_device_name(i) for i in range(torch.cuda.device_count())]
    except ImportError:
        pass  # torch not available

    return env_info
