#!/usr/bin/env python3
"""Collect Linux host metrics and store them in an RRD database."""

from __future__ import annotations

import argparse
import configparser
import os
from pathlib import Path
import subprocess
import time


DEFAULT_CONFIG = "/opt/cheap_system_monitor/monitor.conf"
DS_NAMES = (
    "load1", "load5", "load15", "cpu_user", "cpu_system", "cpu_nice",
    "net_in", "net_out", "disk_root", "disk_1", "disk_2", "memory", "uptime",
)


def read_config(filename: str) -> configparser.SectionProxy:
    parser = configparser.ConfigParser()
    parser.read_dict({"monitor": {
        "rrd_file": "/var/lib/cheap-system-monitor/system.rrd",
        "mountpoints": "/",
        "interfaces": "auto",
        "cpu_sample_seconds": "1.0",
    }})
    if not parser.read(filename):
        raise SystemExit(f"Configuration file not found: {filename}")
    return parser["monitor"]


def run_rrdtool(*arguments: str) -> None:
    try:
        subprocess.run(["rrdtool", *arguments], check=True)
    except FileNotFoundError:
        raise SystemExit("rrdtool is not installed (try: sudo apt install rrdtool)")
    except subprocess.CalledProcessError as error:
        raise SystemExit(f"rrdtool failed with exit status {error.returncode}")


def create_rrd(path: Path) -> None:
    path.parent.mkdir(parents=True, exist_ok=True)
    ds = [
        "DS:load1:GAUGE:180:0:U", "DS:load5:GAUGE:180:0:U",
        "DS:load15:GAUGE:180:0:U", "DS:cpu_user:GAUGE:180:0:100",
        "DS:cpu_system:GAUGE:180:0:100", "DS:cpu_nice:GAUGE:180:0:100",
        "DS:net_in:DERIVE:180:0:U", "DS:net_out:DERIVE:180:0:U",
        "DS:disk_root:GAUGE:180:0:100", "DS:disk_1:GAUGE:180:0:100",
        "DS:disk_2:GAUGE:180:0:100", "DS:memory:GAUGE:180:0:100",
        "DS:uptime:GAUGE:180:0:U",
    ]
    # 1-minute data for a day, 5-minute data for two weeks, hourly for a year.
    archives = [
        "RRA:AVERAGE:0.5:1:1440", "RRA:MAX:0.5:1:1440",
        "RRA:AVERAGE:0.5:5:4032", "RRA:MAX:0.5:5:4032",
        "RRA:AVERAGE:0.5:60:8784", "RRA:MAX:0.5:60:8784",
    ]
    run_rrdtool("create", str(path), "--step", "60", *ds, *archives)


def read_cpu() -> tuple[int, int, int, int]:
    fields = Path("/proc/stat").read_text().splitlines()[0].split()
    if fields[0] != "cpu" or len(fields) < 8:
        raise RuntimeError("unexpected /proc/stat format")
    values = [int(value) for value in fields[1:]]
    user, nice, system, idle = values[:4]
    idle += values[4] if len(values) > 4 else 0  # include iowait
    # Linux reports guest time as already included in user/nice, so only the
    # first eight fields belong in the total (through steal time).
    return user, system, nice, sum(values[:8])


def cpu_percentages(sample_seconds: float) -> tuple[float, float, float]:
    before = read_cpu()
    time.sleep(sample_seconds)
    after = read_cpu()
    total = after[3] - before[3]
    if total <= 0:
        return 0.0, 0.0, 0.0
    return tuple(100.0 * (after[i] - before[i]) / total for i in range(3))


def memory_percent() -> float:
    entries = {}
    for line in Path("/proc/meminfo").read_text().splitlines():
        key, value = line.split(":", 1)
        entries[key] = int(value.split()[0])
    total = entries["MemTotal"]
    available = entries.get("MemAvailable", entries.get("MemFree", 0))
    return 100.0 * (total - available) / total


def uptime_days() -> float:
    seconds = float(Path("/proc/uptime").read_text().split()[0])
    return seconds / 86400.0


def disk_percent(mountpoint: str) -> float:
    stats = os.statvfs(mountpoint)
    used = stats.f_blocks - stats.f_bfree
    available = stats.f_bavail
    return 100.0 * used / (used + available) if used + available else 0.0


def network_counters(interface_setting: str) -> tuple[int, int]:
    wanted = {item.strip() for item in interface_setting.split(",") if item.strip()}
    automatic = not wanted or wanted == {"auto"}
    received = transmitted = 0
    for line in Path("/proc/net/dev").read_text().splitlines()[2:]:
        interface, data = line.split(":", 1)
        interface = interface.strip()
        if (automatic and interface == "lo") or (not automatic and interface not in wanted):
            continue
        fields = data.split()
        received += int(fields[0])
        transmitted += int(fields[8])
    return received, transmitted


def collect(config: configparser.SectionProxy) -> list[float | int | str]:
    load = os.getloadavg()
    cpu = cpu_percentages(config.getfloat("cpu_sample_seconds"))
    net = network_counters(config["interfaces"])
    mounts = [item.strip() for item in config["mountpoints"].split(",") if item.strip()]
    if not mounts or mounts[0] != "/":
        mounts.insert(0, "/")
    if len(mounts) > 3:
        raise SystemExit("At most three mountpoints are supported (including /)")
    disks: list[float | str] = [disk_percent(item) for item in mounts]
    disks.extend(["U"] * (3 - len(disks)))
    return [*load, *cpu, *net, *disks, memory_percent(), uptime_days()]


def main() -> None:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--config", default=DEFAULT_CONFIG)
    parser.add_argument("--create-only", action="store_true")
    args = parser.parse_args()
    config = read_config(args.config)
    rrd_file = Path(config["rrd_file"])
    if not rrd_file.exists():
        create_rrd(rrd_file)
    if not args.create_only:
        values = collect(config)
        update = "N:" + ":".join(str(round(v, 3)) if isinstance(v, float) else str(v) for v in values)
        run_rrdtool("update", str(rrd_file), "--template", ":".join(DS_NAMES), update)


if __name__ == "__main__":
    main()