# Copyright 2022 The Debusine Developers
# See the AUTHORS file at the top-level directory of this distribution
#
# This file is part of Debusine. It is subject to the license terms
# in the LICENSE file found in the top-level directory of this
# distribution. No part of Debusine, including this file, may be copied,
# modified, propagated, or distributed except according to the terms
# contained in the LICENSE file.

"""Functions to get system information."""
import subprocess
from typing import Any

import psutil

from debusine.tasks.models import WorkerType


def total_physical_memory() -> int:
    """Return bytes of RAM memory in system."""
    return psutil.virtual_memory().total


def cpu_count() -> int:
    """Return number of CPUs in the system."""
    return psutil.cpu_count()


def architecture() -> str | None:
    """Return the architecture of the system or None if dpkg not available."""
    try:
        return subprocess.check_output(
            ["dpkg", "--print-architecture"], text=True
        ).strip()
    except (FileNotFoundError, subprocess.CalledProcessError):
        return None


def system_metadata(worker_type: WorkerType) -> dict[str, Any]:
    """Return system metadata formatted as a dictionary."""
    host_architecture = architecture()
    metadata = {
        "system:cpu_count": cpu_count(),
        "system:total_physical_memory": total_physical_memory(),
        "system:worker_type": worker_type,
    }
    if host_architecture := architecture():
        metadata.update(
            {
                "system:host_architecture": host_architecture,
                "system:architectures": [host_architecture],
            }
        )

    return metadata
