#!/usr/bin/env python3
"""Collect ZFS metrics for the Prometheus node_exporter textfile collector.

node_exporter's built-in zfs collector reads /proc/spl/kstat/zfs, which covers ARC
statistics, per-dataset objset counters and pool state -- but nothing about capacity,
fragmentation, vdev errors, scrub history, compression, encryption or snapshots. All of
that has to come from the zpool/zfs CLIs, which is what this script does.

Design notes:

* Each collection stage is independent. A stage that fails logs to stderr and simply
  contributes no samples, rather than taking the whole file down with it -- a broken
  snapshot listing should not cost you your pool capacity graphs.
* Output is written to a temporary file and renamed into place, so node_exporter never
  reads a half-written file.
* Samples are grouped into metric families by a registry rather than printed inline, so
  the exposition format is correct no matter what order stages emit things in.

Run it with --output PATH. Optionally pass --retention label=count (repeatable) to
publish the configured auto-snapshot retention targets alongside the observed counts.
"""

from __future__ import annotations

import argparse
import os
import re
import subprocess
import sys
import time
from collections import defaultdict

# zpool/zfs report state as a string; dashboards want something orderable.
STATE_CODES = {
    "ONLINE": 0,
    "DEGRADED": 1,
    "FAULTED": 2,
    "OFFLINE": 3,
    "UNAVAIL": 4,
    "REMOVED": 5,
    "SUSPENDED": 6,
    "AVAIL": 7,
    "INUSE": 7,
}
STATE_UNKNOWN = 8

# zfs-auto-snapshot (zfstools) names snapshots "zfs-auto-snap_<label>-<Y-m-d-HhM>".
AUTOSNAP_RE = re.compile(r"^zfs-auto-snap_([A-Za-z0-9]+)-")

# How stale a snapshot of each label is allowed to get before it counts as overdue.
# Mirrors the systemd timers NixOS creates for services.zfs.autoSnapshot.
LABEL_INTERVALS = {
    "frequent": 15 * 60,
    "hourly": 60 * 60,
    "daily": 24 * 60 * 60,
    "weekly": 7 * 24 * 60 * 60,
    "monthly": 31 * 24 * 60 * 60,
}


def warn(msg: str) -> None:
    print(f"zfs-metrics: {msg}", file=sys.stderr)


class Registry:
    """Collects samples and renders them in Prometheus text exposition format."""

    def __init__(self) -> None:
        self._families: dict[str, dict] = {}

    def add(self, name, value, labels=None, help="", type="gauge") -> None:
        if value is None:
            return
        fam = self._families.get(name)
        if fam is None:
            fam = {"help": help, "type": type, "samples": []}
            self._families[name] = fam
        fam["samples"].append((labels or {}, value))

    def render(self) -> str:
        out = []
        for name in sorted(self._families):
            fam = self._families[name]
            if not fam["samples"]:
                continue
            out.append(f"# HELP {name} {fam['help']}")
            out.append(f"# TYPE {name} {fam['type']}")
            for labels, value in fam["samples"]:
                out.append(f"{name}{render_labels(labels)} {render_value(value)}")
        return "\n".join(out) + "\n"


def render_labels(labels: dict) -> str:
    if not labels:
        return ""
    # Pool, dataset and vdev names cannot contain quotes or backslashes, so stripping
    # them is a no-op in practice while guaranteeing we never emit a malformed label.
    parts = [f'{k}="{str(v).translate(_LABEL_STRIP)}"' for k, v in sorted(labels.items())]
    return "{" + ",".join(parts) + "}"


_LABEL_STRIP = str.maketrans("", "", '"\\\n')


def render_value(value) -> str:
    if isinstance(value, float):
        if value != value or value in (float("inf"), float("-inf")):
            return "NaN"
        # Avoid scientific notation, which Prometheus accepts but is unpleasant to read,
        # while still keeping enough precision for byte counts.
        return f"{value:.6f}".rstrip("0").rstrip(".") or "0"
    return str(value)


def run(args: list[str]) -> str | None:
    """Run a command, returning stdout, or None if it failed."""
    try:
        proc = subprocess.run(
            args, capture_output=True, text=True, timeout=120, check=False
        )
    except (OSError, subprocess.TimeoutExpired) as exc:
        warn(f"{args[0]} {args[1] if len(args) > 1 else ''}: {exc}")
        return None
    if proc.returncode != 0:
        err = proc.stderr.strip().splitlines()
        warn(f"{' '.join(args)} exited {proc.returncode}: {err[0] if err else ''}")
        return None
    return proc.stdout


def num(value: str):
    """Parse a zfs/zpool numeric field, tolerating '-' and human-formatted suffixes."""
    if value in ("", "-", "none"):
        return None
    value = value.rstrip("%x")
    try:
        return float(value) if "." in value else int(value)
    except ValueError:
        return None


_SIZE_SUFFIXES = {"B": 1}
for _i, _s in enumerate("KMGTPEZ", start=1):
    _SIZE_SUFFIXES[_s] = 1024**_i


def nicebytes(value: str):
    """Parse zfs_nicebytes output, e.g. '0B', '512B', '1.05G'."""
    if not value or value == "-":
        return None
    match = re.fullmatch(r"([0-9]+(?:\.[0-9]+)?)\s*([BKMGTPEZ])?", value.strip())
    if not match:
        return None
    return float(match.group(1)) * _SIZE_SUFFIXES.get(match.group(2) or "B", 1)


def parse_ctime(value: str):
    """Parse a ctime(3) timestamp into a Unix epoch.

    zpool status prints scan timestamps with ctime(), whose format is fixed and
    locale-independent ("Sun Aug 10 05:36:35 2025"), so this is safe to parse. The day
    is space-padded, hence the whitespace normalisation.
    """
    value = " ".join(value.split())
    try:
        return int(time.mktime(time.strptime(value, "%a %b %d %H:%M:%S %Y")))
    except (ValueError, OverflowError):
        return None


def parse_dhms(value: str):
    """Parse secs_to_dhms output: 'HH:MM:SS' or 'N days HH:MM:SS'."""
    match = re.fullmatch(r"(?:(\d+) days )?(\d+):(\d{2}):(\d{2})", value.strip())
    if not match:
        return None
    days, hours, mins, secs = match.groups()
    return int(days or 0) * 86400 + int(hours) * 3600 + int(mins) * 60 + int(secs)


def split_dataset(name: str) -> tuple[str, str]:
    return name.split("/", 1)[0], name


# --------------------------------------------------------------------------------------
# Stage: pool capacity and health
# --------------------------------------------------------------------------------------

POOL_COLUMNS = [
    "name", "size", "allocated", "free", "checkpoint", "fragmentation",
    "capacity", "dedupratio", "health", "freeing", "leaked",
]


def collect_pools(reg: Registry) -> None:
    out = run(["zpool", "list", "-Hp", "-o", ",".join(POOL_COLUMNS)])
    if out is None:
        return
    for line in out.splitlines():
        fields = line.split("\t")
        if len(fields) != len(POOL_COLUMNS):
            continue
        row = dict(zip(POOL_COLUMNS, fields))
        pool = row["name"]
        lbl = {"pool": pool}

        reg.add("zfs_pool_size_bytes", num(row["size"]), lbl,
                "Total size of the pool.")
        reg.add("zfs_pool_allocated_bytes", num(row["allocated"]), lbl,
                "Space allocated in the pool.")
        reg.add("zfs_pool_free_bytes", num(row["free"]), lbl,
                "Space free in the pool.")
        reg.add("zfs_pool_checkpoint_bytes", num(row["checkpoint"]), lbl,
                "Space consumed by a pool checkpoint.")
        reg.add("zfs_pool_freeing_bytes", num(row["freeing"]), lbl,
                "Space still to be reclaimed in the background.")
        reg.add("zfs_pool_leaked_bytes", num(row["leaked"]), lbl,
                "Space leaked, which should normally be zero.")
        reg.add("zfs_pool_dedup_ratio", num(row["dedupratio"]), lbl,
                "Deduplication ratio.")

        frag = num(row["fragmentation"])
        if frag is not None:
            reg.add("zfs_pool_fragmentation_ratio", frag / 100, lbl,
                    "Free space fragmentation, 0-1.")
        cap = num(row["capacity"])
        if cap is not None:
            reg.add("zfs_pool_capacity_ratio", cap / 100, lbl,
                    "Allocated space as a fraction of pool size, 0-1.")

        health = row["health"]
        reg.add("zfs_pool_health", 1, {"pool": pool, "state": health},
                "Pool health, 1 for the currently reported state.")
        reg.add("zfs_pool_health_code", STATE_CODES.get(health, STATE_UNKNOWN), lbl,
                "Pool health as a code: 0 ONLINE, 1 DEGRADED, 2 FAULTED, 3 OFFLINE, "
                "4 UNAVAIL, 5 REMOVED, 6 SUSPENDED, 8 unknown.")


# --------------------------------------------------------------------------------------
# Stage: zpool status -- vdev topology, errors, scrub/resilver history
# --------------------------------------------------------------------------------------

CONFIG_HEADER_RE = re.compile(r"^\s*NAME\s+STATE\s+READ\s+WRITE\s+CKSUM")
POOL_RE = re.compile(r"^\s*pool:\s*(\S+)")
SCAN_RE = re.compile(r"^\s*scan:\s*(.*)$")
ERRORS_RE = re.compile(r"^\s*errors:\s*(.*)$")

SCAN_PATTERNS = [
    ("scrub", "finished",
     re.compile(r"^scrub repaired (\S+) in (.+?) with (\d+) errors on (.+)$")),
    ("resilver", "finished",
     re.compile(r"^resilvered (\S+) in (.+?) with (\d+) errors on (.+)$")),
    ("scrub", "canceled", re.compile(r"^scrub canceled on (.+)$")),
    ("resilver", "canceled", re.compile(r"^resilver canceled on (.+)$")),
    ("scrub", "scanning", re.compile(r"^scrub in progress since (.+)$")),
    ("scrub", "paused", re.compile(r"^scrub paused since (.+)$")),
    ("resilver", "scanning", re.compile(r"^resilver in progress since (.+)$")),
]


def collect_status(reg: Registry) -> None:
    out = run(["zpool", "status", "-p"])
    if out is None:
        return

    pool = None
    in_config = False
    seen_scan: set[str] = set()
    pools: list[str] = []

    for line in out.splitlines():
        match = POOL_RE.match(line)
        if match:
            pool = match.group(1)
            pools.append(pool)
            in_config = False
            continue

        if pool is None:
            continue

        match = SCAN_RE.match(line)
        if match and pool not in seen_scan:
            seen_scan.add(pool)
            parse_scan(reg, pool, match.group(1).strip())
            continue

        match = ERRORS_RE.match(line)
        if match:
            in_config = False
            reg.add("zfs_pool_data_errors", parse_data_errors(match.group(1)),
                    {"pool": pool},
                    "Data errors reported by zpool status; 0 means no known errors.")
            continue

        if CONFIG_HEADER_RE.match(line):
            in_config = True
            continue

        if in_config:
            if not line.strip():
                in_config = False
                continue
            parse_vdev(reg, pool, line)

    # A pool with no scan line at all has never been scrubbed; make that explicit
    # rather than leaving the series absent.
    for p in pools:
        if p not in seen_scan:
            set_scan_flags(reg, p, scrub=0, resilver=0)
            reg.add("zfs_pool_scan_state", 1, {"pool": p, "state": "none"},
                    "Current scan state, 1 for the active state.")


def parse_data_errors(text: str):
    text = text.strip()
    if not text:
        return None
    if "No known data errors" in text:
        return 0
    match = re.search(r"(\d+) data errors", text)
    if match:
        return int(match.group(1))
    # e.g. "Permanent errors have been detected in the following files:" -- we know
    # something is wrong but not how many, so report a non-zero count.
    return 1


def set_scan_flags(reg: Registry, pool: str, scrub: int, resilver: int) -> None:
    reg.add("zfs_pool_scrub_in_progress", scrub, {"pool": pool},
            "Whether a scrub is currently running.")
    reg.add("zfs_pool_resilver_in_progress", resilver, {"pool": pool},
            "Whether a resilver is currently running.")


def parse_scan(reg: Registry, pool: str, text: str) -> None:
    lbl = {"pool": pool}

    for kind, state, pattern in SCAN_PATTERNS:
        match = pattern.match(text)
        if not match:
            continue

        scanning = state in ("scanning", "paused")
        set_scan_flags(
            reg, pool,
            scrub=int(scanning and kind == "scrub"),
            resilver=int(scanning and kind == "resilver"),
        )
        reg.add("zfs_pool_scan_state", 1, {"pool": pool, "state": f"{kind}_{state}"},
                "Current scan state, 1 for the active state.")

        if state == "finished":
            repaired, duration, errors, when = match.groups()
            reg.add(f"zfs_pool_last_{kind}_end_timestamp_seconds", parse_ctime(when),
                    lbl, f"Unix time the last {kind} completed.")
            reg.add(f"zfs_pool_last_{kind}_duration_seconds", parse_dhms(duration),
                    lbl, f"How long the last completed {kind} took.")
            reg.add(f"zfs_pool_last_{kind}_errors", int(errors), lbl,
                    f"Unrepairable errors found by the last completed {kind}.")
            reg.add(f"zfs_pool_last_{kind}_repaired_bytes", nicebytes(repaired), lbl,
                    f"Bytes repaired by the last completed {kind}.")
        elif state == "scanning":
            reg.add(f"zfs_pool_{kind}_started_timestamp_seconds",
                    parse_ctime(match.group(1)), lbl,
                    f"Unix time the running {kind} started.")
        return

    # "none requested", or a format we do not recognise.
    set_scan_flags(reg, pool, scrub=0, resilver=0)
    state = "none" if "none requested" in text else "unknown"
    reg.add("zfs_pool_scan_state", 1, {"pool": pool, "state": state},
            "Current scan state, 1 for the active state.")


def parse_vdev(reg: Registry, pool: str, line: str) -> None:
    fields = line.split()
    # vdev rows are "NAME STATE READ WRITE CKSUM" plus an optional trailing note.
    # Section headers (logs/cache/spares) and spare rows have fewer fields.
    if len(fields) < 5:
        return
    name, state, read, write, cksum = fields[:5]
    if not (read.isdigit() and write.isdigit() and cksum.isdigit()):
        return

    reg.add("zfs_pool_vdev_state_code", STATE_CODES.get(state, STATE_UNKNOWN),
            {"pool": pool, "vdev": name},
            "vdev state as a code, same encoding as zfs_pool_health_code.")
    for kind, value in (("read", read), ("write", write), ("checksum", cksum)):
        reg.add("zfs_pool_vdev_errors_total", int(value),
                {"pool": pool, "vdev": name, "type": kind},
                "Cumulative vdev errors as reported by zpool status. "
                "Reset by zpool clear.", type="counter")


# --------------------------------------------------------------------------------------
# Stage: datasets -- space, compression, encryption
# --------------------------------------------------------------------------------------

DATASET_COLUMNS = [
    "name", "type", "used", "available", "referenced", "logicalused",
    "usedbysnapshots", "usedbydataset", "compressratio", "compression",
    "encryption", "keystatus", "quota", "refquota", "recordsize", "mounted",
    "readonly", "com.sun:auto-snapshot",
]

ONOFF = {"on": 1, "off": 0, "yes": 1, "no": 0, "true": 1, "false": 0}

# `zfs get keystatus` reports "none" for datasets that are not encrypted at all.
KEYSTATUS_CODES = {"none": 0, "-": 0, "available": 1, "unavailable": 2}


def collect_datasets(reg: Registry) -> None:
    out = run([
        "zfs", "list", "-Hp", "-t", "filesystem,volume",
        "-o", ",".join(DATASET_COLUMNS),
    ])
    if out is None:
        return

    for line in out.splitlines():
        fields = line.split("\t")
        if len(fields) != len(DATASET_COLUMNS):
            continue
        row = dict(zip(DATASET_COLUMNS, fields))
        pool, dataset = split_dataset(row["name"])
        lbl = {"pool": pool, "dataset": dataset}

        reg.add("zfs_dataset_used_bytes", num(row["used"]), lbl,
                "Space consumed by the dataset and its descendants.")
        reg.add("zfs_dataset_available_bytes", num(row["available"]), lbl,
                "Space available to the dataset.")
        reg.add("zfs_dataset_referenced_bytes", num(row["referenced"]), lbl,
                "Space referenced by the dataset itself.")
        reg.add("zfs_dataset_logical_used_bytes", num(row["logicalused"]), lbl,
                "Logical (pre-compression) space consumed.")
        reg.add("zfs_dataset_used_by_snapshots_bytes", num(row["usedbysnapshots"]), lbl,
                "Space consumed by this dataset's snapshots.")
        reg.add("zfs_dataset_used_by_dataset_bytes", num(row["usedbydataset"]), lbl,
                "Space consumed by the dataset's own data.")
        reg.add("zfs_dataset_compression_ratio", num(row["compressratio"]), lbl,
                "Achieved compression ratio; 1 means no saving.")
        reg.add("zfs_dataset_quota_bytes", num(row["quota"]), lbl,
                "Configured quota in bytes; 0 means unlimited.")
        reg.add("zfs_dataset_refquota_bytes", num(row["refquota"]), lbl,
                "Configured refquota in bytes; 0 means unlimited.")
        reg.add("zfs_dataset_recordsize_bytes", num(row["recordsize"]), lbl,
                "Configured record size.")
        reg.add("zfs_dataset_mounted", ONOFF.get(row["mounted"]), lbl,
                "Whether the dataset is currently mounted.")
        reg.add("zfs_dataset_readonly", ONOFF.get(row["readonly"]), lbl,
                "Whether the dataset is read-only.")

        autosnap = ONOFF.get(row["com.sun:auto-snapshot"])
        reg.add("zfs_dataset_autosnapshot_enabled", autosnap, lbl,
                "Whether com.sun:auto-snapshot is set true on the dataset.")

        # Info-style metrics: the value is always 1, the detail lives in the labels.
        # The label is named after the property rather than a generic "algorithm" so
        # that dashboards can attach it to a numeric series with group_left() without
        # two of these colliding on the same label name.
        reg.add("zfs_dataset_compression", 1,
                {**lbl, "compression": row["compression"]},
                "Configured compression algorithm, 1 for the active setting.")
        reg.add("zfs_dataset_encryption", 1,
                {**lbl, "encryption": row["encryption"]},
                "Configured encryption suite, 1 for the active setting.")
        reg.add("zfs_dataset_encrypted", int(row["encryption"] not in ("off", "-")),
                lbl, "Whether the dataset is encrypted.")
        reg.add("zfs_dataset_key_available", int(row["keystatus"] == "available"),
                lbl, "Whether the encryption key is loaded (unencrypted datasets "
                     "report 0).")
        # keystatus is a three-state property -- "none" on an unencrypted dataset is
        # not the same thing as "unavailable" on an encrypted one. Collapsing it to a
        # boolean loses that distinction and makes every unencrypted dataset look
        # locked, so publish the three states as a code too.
        reg.add("zfs_dataset_keystatus_code", KEYSTATUS_CODES.get(row["keystatus"], 3),
                lbl, "Encryption key status: 0 not encrypted, 1 key loaded, "
                     "2 key not loaded, 3 unknown.")
        reg.add("zfs_dataset_type", 1, {**lbl, "type": row["type"]},
                "Dataset type, 1 for the actual type.")


# --------------------------------------------------------------------------------------
# Stage: snapshots
# --------------------------------------------------------------------------------------

def collect_snapshots(reg: Registry) -> None:
    out = run(["zfs", "list", "-Hp", "-t", "snapshot", "-o", "name,creation,used"])
    if out is None:
        return

    # (dataset, label) -> [count, newest, oldest, used]
    per_dataset: dict[tuple[str, str], list] = defaultdict(
        lambda: [0, None, None, 0.0]
    )
    per_pool: dict[tuple[str, str], list] = defaultdict(lambda: [0, None, None, 0.0])

    for line in out.splitlines():
        fields = line.split("\t")
        if len(fields) != 3:
            continue
        full, creation, used = fields
        if "@" not in full:
            continue
        dataset, snapname = full.split("@", 1)
        pool, _ = split_dataset(dataset)

        created = num(creation)
        if created is None:
            continue
        size = num(used) or 0

        match = AUTOSNAP_RE.match(snapname)
        label = match.group(1) if match else "manual"

        for bucket, key in (
            (per_dataset, (dataset, label)),
            (per_pool, (pool, label)),
        ):
            entry = bucket[key]
            entry[0] += 1
            entry[1] = created if entry[1] is None else max(entry[1], created)
            entry[2] = created if entry[2] is None else min(entry[2], created)
            entry[3] += size

    for (dataset, label), (count, newest, oldest, used) in per_dataset.items():
        pool, _ = split_dataset(dataset)
        lbl = {"pool": pool, "dataset": dataset, "label": label}
        reg.add("zfs_snapshot_count", count, lbl,
                "Number of snapshots of this dataset with this label.")
        reg.add("zfs_snapshot_newest_timestamp_seconds", newest, lbl,
                "Creation time of the most recent snapshot.")
        reg.add("zfs_snapshot_oldest_timestamp_seconds", oldest, lbl,
                "Creation time of the oldest snapshot; with the newest one this is "
                "the retention window actually on disk.")
        reg.add("zfs_snapshot_used_bytes", used, lbl,
                "Space consumed by these snapshots.")

    for (pool, label), (count, newest, oldest, used) in per_pool.items():
        lbl = {"pool": pool, "label": label}
        reg.add("zfs_pool_snapshot_count", count, lbl,
                "Number of snapshots in the pool with this label.")
        reg.add("zfs_pool_snapshot_newest_timestamp_seconds", newest, lbl,
                "Creation time of the pool's most recent snapshot with this label.")
        reg.add("zfs_pool_snapshot_oldest_timestamp_seconds", oldest, lbl,
                "Creation time of the pool's oldest snapshot with this label.")
        reg.add("zfs_pool_snapshot_used_bytes", used, lbl,
                "Space consumed by the pool's snapshots with this label.")


def collect_retention(reg: Registry, retention: dict[str, int]) -> None:
    """Publish the configured auto-snapshot policy so dashboards can compare
    observed snapshot counts and ages against what is supposed to exist."""
    for label, keep in sorted(retention.items()):
        reg.add("zfs_snapshot_retention_target", keep, {"label": label},
                "Configured number of snapshots to keep for this label.")
    for label, seconds in sorted(LABEL_INTERVALS.items()):
        if label in retention:
            reg.add("zfs_snapshot_interval_seconds", seconds, {"label": label},
                    "Expected interval between snapshots of this label.")


# --------------------------------------------------------------------------------------


def parse_retention(values: list[str]) -> dict[str, int]:
    out = {}
    for item in values:
        label, _, count = item.partition("=")
        try:
            out[label] = int(count)
        except ValueError:
            warn(f"ignoring malformed --retention {item!r}")
    return out


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--output", required=True,
                        help="path of the .prom file to write")
    parser.add_argument("--retention", action="append", default=[], metavar="LABEL=N",
                        help="configured auto-snapshot retention, repeatable")
    args = parser.parse_args()

    reg = Registry()
    started = time.time()

    for stage in (collect_pools, collect_status, collect_datasets, collect_snapshots):
        try:
            stage(reg)
        except Exception as exc:  # noqa: BLE001 - one bad stage must not lose the rest
            warn(f"{stage.__name__} failed: {exc!r}")

    try:
        collect_retention(reg, parse_retention(args.retention))
    except Exception as exc:  # noqa: BLE001
        warn(f"collect_retention failed: {exc!r}")

    reg.add("zfs_metrics_collection_duration_seconds", time.time() - started, {},
            "Time this collector took to gather ZFS metrics.")

    tmp = f"{args.output}.tmp"
    with open(tmp, "w") as handle:
        handle.write(reg.render())
    os.replace(tmp, args.output)
    return 0


if __name__ == "__main__":
    sys.exit(main())