#!/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",
    "io_wait",
    "net_in", "net_out", "disk_root", "disk_1", "disk_2", "memory", "uptime",
    "temperature", "throttled",
    "failed_services", "oom_kills",
    "free_root", "free_1", "free_2", "inode_root", "inode_1", "inode_2",
    "readonly_root", "readonly_1", "readonly_2",
    "storage_health",
)


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:io_wait: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",
        "DS:temperature:GAUGE:180:-50:200", "DS:throttled:GAUGE:180:0:1",
        "DS:failed_services:GAUGE:180:0:U", "DS:oom_kills:DERIVE:180:0:U",
        "DS:free_root:GAUGE:180:0:U", "DS:free_1:GAUGE:180:0:U",
        "DS:free_2:GAUGE:180:0:U", "DS:inode_root:GAUGE:180:0:100",
        "DS:inode_1:GAUGE:180:0:100", "DS:inode_2:GAUGE:180:0:100",
        "DS:readonly_root:GAUGE:180:0:1", "DS:readonly_1:GAUGE:180:0:1",
        "DS:readonly_2:GAUGE:180:0:1",
        "DS:storage_health:GAUGE:180:0:10",
    ]
    # 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:MIN:0.5:1:1440",
        "RRA:AVERAGE:0.5:5:4032", "RRA:MAX:0.5:5:4032",
        "RRA:MIN:0.5:5:4032",
        "RRA:AVERAGE:0.5:60:8784", "RRA:MAX:0.5:60:8784",
        "RRA:MIN:0.5:60:8784",
    ]
    run_rrdtool("create", str(path), "--step", "60", *ds, *archives)


def read_cpu() -> tuple[int, 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]
    iowait = values[4] if len(values) > 4 else 0
    idle += 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, iowait, sum(values[:8])


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


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 cpu_temperature() -> float | str:
    """Return the CPU temperature in Celsius, or an RRD unknown value."""
    thermal_root = Path("/sys/class/thermal")
    if not thermal_root.exists():
        return "U"

    zones = sorted(thermal_root.glob("thermal_zone*"))
    preferred = []
    others = []
    for zone in zones:
        try:
            zone_type = (zone / "type").read_text().strip().lower()
        except OSError:
            zone_type = ""
        target = preferred if any(name in zone_type for name in ("cpu", "soc", "package")) else others
        target.append(zone)

    for zone in preferred + others:
        try:
            value = float((zone / "temp").read_text().strip())
            # Linux thermal-zone temperatures are normally millidegrees.
            return value / 1000.0 if abs(value) >= 1000 else value
        except (OSError, ValueError):
            continue
    return "U"


def throttle_state() -> int | str:
    """Return 1 during active Raspberry Pi throttling, 0 otherwise."""
    try:
        result = subprocess.run(
            ["vcgencmd", "get_throttled"], check=True, capture_output=True,
            text=True, timeout=2,
        )
    except (FileNotFoundError, subprocess.CalledProcessError, subprocess.TimeoutExpired):
        return "U"

    try:
        flags = int(result.stdout.strip().split("=", 1)[1], 16)
    except (IndexError, ValueError):
        return "U"
    # Bits 1, 2, and 3 mean frequency capped, throttled, and soft temperature
    # limit active. Ignore the corresponding sticky "has occurred" bits.
    return int(bool(flags & 0x0E))


def failed_service_count() -> int | str:
    """Return the number of failed systemd service units."""
    try:
        result = subprocess.run(
            ["systemctl", "list-units", "--state=failed", "--type=service",
             "--no-legend", "--plain"],
            check=True, capture_output=True, text=True, timeout=5,
        )
    except (FileNotFoundError, subprocess.CalledProcessError, subprocess.TimeoutExpired):
        return "U"
    return sum(bool(line.strip()) for line in result.stdout.splitlines())


def oom_kill_count() -> int | str:
    """Return the kernel's cumulative OOM-kill counter since boot."""
    try:
        for line in Path("/proc/vmstat").read_text().splitlines():
            name, value = line.split()
            if name == "oom_kill":
                return int(value)
    except (OSError, ValueError):
        pass
    return "U"


def disk_metrics(mountpoint: str) -> tuple[float, float, float | str, int]:
    """Return used %, available GiB, inode used %, and read-only status."""
    stats = os.statvfs(mountpoint)
    used = stats.f_blocks - stats.f_bfree
    available = stats.f_bavail
    used_percent = 100.0 * used / (used + available) if used + available else 0.0
    block_size = stats.f_frsize or stats.f_bsize
    free_gib = available * block_size / (1024 ** 3)
    used_inodes = stats.f_files - stats.f_ffree
    available_inodes = stats.f_favail
    inode_total = used_inodes + available_inodes
    inode_percent: float | str = (
        100.0 * used_inodes / inode_total if inode_total else "U"
    )
    read_only = int(bool(stats.f_flag & os.ST_RDONLY))
    return used_percent, free_gib, inode_percent, read_only


def storage_health_score(
        storage: list[tuple[float, float, float | str, int]],
        io_wait: float) -> float:
    """Score current storage operability from 0 (unusable) to 10 (healthy)."""
    if any(read_only for _used, _free, _inodes, read_only in storage):
        return 0.0

    score = 10.0
    for used, _free, inodes, _read_only in storage:
        # No penalty through 75%; decline linearly to zero at 100%.
        space_score = max(0.0, min(10.0, (100.0 - used) / 25.0 * 10.0))
        score = min(score, space_score)
        if isinstance(inodes, float):
            inode_score = max(0.0, min(10.0, (100.0 - inodes) / 25.0 * 10.0))
            score = min(score, inode_score)

    # I/O wait is workload-sensitive, so it can reduce the score by at most
    # three points: no penalty through 10%, maximum penalty at 50%.
    io_penalty = max(0.0, min(3.0, (io_wait - 10.0) / 40.0 * 3.0))
    return max(0.0, score - io_penalty)


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 /)")
    storage = [disk_metrics(item) for item in mounts]
    disks: list[float | str] = [item[0] for item in storage]
    free: list[float | str] = [item[1] for item in storage]
    inodes: list[float | str] = [item[2] for item in storage]
    read_only: list[int | str] = [item[3] for item in storage]
    for values in (disks, free, inodes, read_only):
        values.extend(["U"] * (3 - len(values)))
    return [
        *load, *cpu, *net, *disks, memory_percent(), uptime_days(),
        cpu_temperature(), throttle_state(),
        failed_service_count(), oom_kill_count(),
        *free, *inodes, *read_only,
        storage_health_score(storage, cpu[3]),
    ]


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()