bigbrother/utils/resource.py

#!/usr/bin/env python3
"""Hardware tier detection and resource monitoring via /proc and /sys."""

import os
import re
from dataclasses import dataclass, field
from pathlib import Path
from typing import Dict, List, Optional, Tuple

TIER_OPI_ZERO3 = "opi_zero3"
TIER_PI_ZERO = "pi_zero"
TIER_GENERIC = "generic"


@dataclass
class HardwareInfo:
    tier: str
    model: str
    cpu_model: str
    cpu_cores: int
    total_ram_mb: int
    architecture: str


@dataclass
class ResourceUsage:
    ram_used_mb: float
    ram_total_mb: float
    ram_percent: float
    cpu_percent: float  # System-wide, from /proc/stat delta
    disk_used_pct: float
    disk_free_mb: float
    temperature_c: Optional[float]


@dataclass
class ProcessResources:
    pid: int
    name: str
    rss_mb: float
    vms_mb: float
    cpu_ticks: int
    state: str
    threads: int


def get_hardware_tier() -> str:
    info = detect_hardware()
    return info.tier


def detect_hardware() -> HardwareInfo:
    model = _read_device_model()
    cpu_model, cpu_cores, arch = _read_cpuinfo()

    if "Orange Pi Zero3" in model or "Orange Pi Zero 3" in model:
        tier = TIER_OPI_ZERO3
    elif "Raspberry Pi Zero 2" in model or "Raspberry Pi Zero2" in model:
        tier = TIER_PI_ZERO
    elif "Raspberry Pi" in model and ("Zero" in model or "zero" in model):
        tier = TIER_PI_ZERO
    elif "Allwinner" in cpu_model or "H618" in cpu_model:
        # Fallback detection via CPU for OPi Zero 3
        tier = TIER_OPI_ZERO3
    elif "BCM2835" in cpu_model or "BCM2710" in cpu_model:
        tier = TIER_PI_ZERO
    else:
        tier = TIER_GENERIC

    total_ram = _read_total_ram_mb()

    return HardwareInfo(
        tier=tier,
        model=model or "Unknown",
        cpu_model=cpu_model,
        cpu_cores=cpu_cores,
        total_ram_mb=total_ram,
        architecture=arch,
    )


def _read_device_model() -> str:
    dt_model = Path("/proc/device-tree/model")
    if dt_model.exists():
        try:
            return dt_model.read_text().strip().rstrip("\x00")
        except IOError:
            pass
    return ""


def _read_cpuinfo() -> Tuple[str, int, str]:
    cpu_model = ""
    cores = 0
    arch = "unknown"

    try:
        with open("/proc/cpuinfo", "r") as f:
            content = f.read()

        for line in content.splitlines():
            if line.startswith("model name") or line.startswith("Hardware"):
                cpu_model = line.split(":", 1)[1].strip()
            if line.startswith("processor"):
                cores += 1

        # Get architecture
        with open("/proc/cpuinfo", "r") as f:
            first_lines = f.read(4096)
        if "aarch64" in first_lines.lower() or "ARMv8" in first_lines:
            arch = "aarch64"
        elif "armv7" in first_lines.lower():
            arch = "armv7l"
        else:
            import platform
            arch = platform.machine()
    except IOError:
        pass

    return cpu_model, max(cores, 1), arch


def _read_total_ram_mb() -> int:
    try:
        with open("/proc/meminfo", "r") as f:
            for line in f:
                if line.startswith("MemTotal:"):
                    kb = int(line.split()[1])
                    return kb // 1024
    except IOError:
        pass
    return 0


def get_memory_info() -> Dict[str, int]:
    info = {}
    try:
        with open("/proc/meminfo", "r") as f:
            for line in f:
                parts = line.split()
                key = parts[0].rstrip(":")
                val = int(parts[1])  # in kB
                info[key] = val
    except IOError:
        pass
    return info


def get_resource_usage(disk_path: str = "/") -> ResourceUsage:
    mem = get_memory_info()
    total = mem.get("MemTotal", 0)
    available = mem.get("MemAvailable", mem.get("MemFree", 0))
    used = total - available
    pct = (used / total * 100) if total > 0 else 0.0

    stat = os.statvfs(disk_path)
    disk_total = stat.f_blocks * stat.f_frsize
    disk_free = stat.f_bavail * stat.f_frsize
    disk_used_pct = ((disk_total - disk_free) / disk_total * 100) if disk_total > 0 else 0.0

    temp = get_cpu_temperature()

    return ResourceUsage(
        ram_used_mb=used / 1024,
        ram_total_mb=total / 1024,
        ram_percent=pct,
        cpu_percent=_get_cpu_percent(),
        disk_used_pct=disk_used_pct,
        disk_free_mb=disk_free / (1024 * 1024),
        temperature_c=temp,
    )


def get_cpu_temperature() -> Optional[float]:
    thermal_zones = Path("/sys/class/thermal")
    if thermal_zones.exists():
        for zone in sorted(thermal_zones.iterdir()):
            temp_file = zone / "temp"
            type_file = zone / "type"
            if temp_file.exists():
                try:
                    temp_raw = int(temp_file.read_text().strip())
                    return temp_raw / 1000.0
                except (IOError, ValueError):
                    continue
    # Fallback for some SBCs
    hwmon = Path("/sys/class/hwmon")
    if hwmon.exists():
        for hw in sorted(hwmon.iterdir()):
            temp_file = hw / "temp1_input"
            if temp_file.exists():
                try:
                    return int(temp_file.read_text().strip()) / 1000.0
                except (IOError, ValueError):
                    continue
    return None


def _get_cpu_percent() -> float:
    try:
        with open("/proc/stat", "r") as f:
            line = f.readline()
        parts = line.split()
        # user, nice, system, idle, iowait, irq, softirq, steal
        total = sum(int(p) for p in parts[1:8])
        idle = int(parts[4])
        # Snapshot-based; for delta, caller should sample twice
        if total > 0:
            return ((total - idle) / total) * 100
    except (IOError, IndexError, ValueError):
        pass
    return 0.0


def get_process_resources(pid: int) -> Optional[ProcessResources]:
    try:
        stat_path = f"/proc/{pid}/stat"
        status_path = f"/proc/{pid}/status"

        with open(stat_path, "r") as f:
            stat_line = f.read()

        # Parse /proc/PID/stat — comm field can contain spaces/parens
        match = re.match(r"(\d+) \((.+)\) (.+)", stat_line)
        if not match:
            return None

        name = match.group(2)
        fields = match.group(3).split()
        state = fields[0]
        utime = int(fields[11])
        stime = int(fields[12])
        threads = int(fields[17])
        vsize = int(fields[20])
        rss_pages = int(fields[21])

        page_size = os.sysconf("SC_PAGE_SIZE")
        rss_mb = (rss_pages * page_size) / (1024 * 1024)
        vms_mb = vsize / (1024 * 1024)

        return ProcessResources(
            pid=pid,
            name=name,
            rss_mb=rss_mb,
            vms_mb=vms_mb,
            cpu_ticks=utime + stime,
            state=state,
            threads=threads,
        )
    except (IOError, OSError, IndexError, ValueError):
        return None


def get_disk_usage(path: str = "/") -> Tuple[float, float, float]:
    stat = os.statvfs(path)
    total_gb = (stat.f_blocks * stat.f_frsize) / (1024 ** 3)
    free_gb = (stat.f_bavail * stat.f_frsize) / (1024 ** 3)
    used_pct = ((stat.f_blocks - stat.f_bavail) / stat.f_blocks * 100) if stat.f_blocks > 0 else 0
    return total_gb, free_gb, used_pct