#!/usr/bin/env python3
"""Collect SMART self-test history and drive identity for node_exporter's textfile
collector.

Two things smartctl_exporter cannot give us:

* Self-test history. The exporter reads only `count` and `error_count_total` out of
  smartctl's self-test log and discards the log table itself, so it cannot tell you
  when a drive last ran a short or an extended test, or whether that test passed.

* A ready-to-display drive identity string. Grafana can only put a *label* in a row
  title, and PromQL has no way to turn a number (capacity, rotation rate) into one.
  So the formatting happens here, in `smart_drive_label`.

Device names are emitted without the /dev/ prefix so they join against
smartctl_exporter's `device` label, which is the bare name ("sda").

Run with --output PATH. Reading the self-test log does not spin an idle disk up, but
there is no reason to poll it often -- a 15 minute timer is plenty.
"""

from __future__ import annotations

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

# smartctl spells these differently per transport and firmware; normalise to the names
# people actually use.
TEST_TYPES = [
    (re.compile(r"\bshort\b", re.I), "short"),
    (re.compile(r"\b(extended|long)\b", re.I), "extended"),
    (re.compile(r"\bconveyance\b", re.I), "conveyance"),
    (re.compile(r"\bselective\b", re.I), "selective"),
    (re.compile(r"\boffline\b", re.I), "offline"),
]


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


def render_labels(labels: dict) -> str:
    if not labels:
        return ""
    parts = [f'{k}="{_clean(str(v))}"' for k, v in sorted(labels.items())]
    return "{" + ",".join(parts) + "}"


def _clean(value: str) -> str:
    return value.translate(str.maketrans("", "", '"\\\n'))


class Registry:
    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.setdefault(name, {"help": help, "type": type,
                                               "samples": []})
        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"]:
                if isinstance(value, float):
                    rendered = f"{value:.6f}".rstrip("0").rstrip(".") or "0"
                else:
                    rendered = str(value)
                out.append(f"{name}{render_labels(labels)} {rendered}")
        return "\n".join(out) + "\n"


def run_json(args: list[str]) -> dict | None:
    """Run smartctl with JSON output.

    smartctl uses its exit status as a bitfield rather than a simple success flag --
    bit 3 means "a self-test failed", which is data we specifically want, not an error.
    So parse stdout regardless of exit status and only give up if there is no JSON.
    """
    try:
        proc = subprocess.run(args, capture_output=True, text=True, timeout=120,
                              check=False)
    except (OSError, subprocess.TimeoutExpired) as exc:
        warn(f"{' '.join(args)}: {exc}")
        return None
    if not proc.stdout.strip():
        err = proc.stderr.strip().splitlines()
        warn(f"{' '.join(args)} produced no output: {err[0] if err else ''}")
        return None
    try:
        return json.loads(proc.stdout)
    except json.JSONDecodeError as exc:
        warn(f"{' '.join(args)}: could not parse JSON: {exc}")
        return None


def scan_devices() -> list[str]:
    data = run_json(["smartctl", "--scan", "-j"])
    if not data:
        return []
    return [d["name"] for d in data.get("devices", []) if d.get("name")]


def normalise_type(text: str) -> str:
    for pattern, name in TEST_TYPES:
        if pattern.search(text):
            return name
    return "other"


def iter_entries(data: dict):
    """Yield self-test log entries, whatever shape smartctl used.

    ATA groups them into tables:
        ata_smart_self_test_log: {standard: {table: [...]}, extended: {table: [...]}}

    SCSI does not have a table at all -- smartctl emits one flat top-level key per
    entry (scsiprint.cpp: `snprintf(st, "scsi_self_test_%d", k)`):
        scsi_self_test_0: {code: {...}, result: {...}, power_on_time: {hours: N}}

    Handling only the ATA shape would silently produce nothing for SAS drives.
    """
    for key, value in sorted(data.items()):
        if "self_test" not in key or not isinstance(value, dict):
            continue
        # SCSI: the value *is* one entry.
        if "code" in value or "result" in value:
            yield value
            continue
        # ATA: a table here, or one level down under standard/extended.
        if isinstance(value.get("table"), list):
            yield from (e for e in value["table"] if isinstance(e, dict))
        for sub in value.values():
            if isinstance(sub, dict) and isinstance(sub.get("table"), list):
                yield from (e for e in sub["table"] if isinstance(e, dict))


def _field(entry: dict, *names):
    for name in names:
        if name in entry:
            return entry[name]
    return {}


def entry_fields(entry: dict) -> tuple[str, str, bool | None, float | None]:
    """Pull (test type, status text, passed, lifetime hours) out of a log entry.

    All three transports name these differently:
        ATA   type / status{passed} / lifetime_hours
        SCSI  code / result        / power_on_time{hours}
        NVMe  self_test_code / self_test_result / power_on_hours
    """
    type_field = _field(entry, "type", "code", "self_test_code")
    if isinstance(type_field, dict):
        type_text = type_field.get("string") or str(type_field.get("value", ""))
    else:
        type_text = str(type_field)

    status = _field(entry, "status", "result", "self_test_result")
    if isinstance(status, dict):
        status_text = status.get("string") or str(status.get("value", ""))
        passed = status.get("passed")
    else:
        status_text = str(status)
        passed = None

    if passed is None and status_text:
        # SCSI logs carry no "passed" boolean, only a result string. Check for failure
        # before checking for "Completed", because SCSI reports a failed long test as
        # "Completed, segment failed" -- a naive startswith("completed") reads that as
        # a pass.
        low = status_text.lower()
        if "in progress" in low or low.startswith("aborted"):
            passed = None
        elif "fail" in low or "error" in low:
            passed = "without error" in low
        elif low.startswith("completed"):
            passed = True

    hours = entry.get("lifetime_hours")
    if hours is None and isinstance(entry.get("power_on_time"), dict):
        hours = entry["power_on_time"].get("hours")
    if hours is None and isinstance(entry.get("power_on_hours"), (int, float)):
        hours = entry["power_on_hours"]

    return normalise_type(type_text), status_text, passed, hours


def human_capacity(nbytes) -> str:
    """Format a capacity the way the label on the drive does: decimal, not binary."""
    if not isinstance(nbytes, (int, float)) or nbytes <= 0:
        return ""
    for unit, scale in (("TB", 10 ** 12), ("GB", 10 ** 9), ("MB", 10 ** 6)):
        if nbytes >= scale:
            value = nbytes / scale
            return f"{value:.1f} {unit}" if value < 100 else f"{value:.0f} {unit}"
    return f"{int(nbytes)} B"


def technology(data: dict) -> str:
    """A short "what kind of drive is this" string: bus plus spindle speed."""
    protocol = (data.get("device") or {}).get("protocol", "")
    bus = {"ATA": "SATA", "SCSI": "SAS", "NVMe": "NVMe"}.get(protocol, protocol)

    rpm = data.get("rotation_rate")
    if rpm is None or rpm == 0:
        # smartctl reports 0 (or omits it) for anything without platters.
        return f"{bus} SSD".strip()
    # 7200 -> "7.2K"; the round numbers people actually quote.
    spin = f"{rpm / 1000:.1f}K".replace(".0K", "K")
    return f"{bus} {spin}".strip()


def collect_identity(reg: Registry, name: str, data: dict) -> None:
    """Emit one info metric whose label is a preformatted identity line.

    This exists to be read by a human in a Grafana row title, so it is deliberately a
    single rendered string rather than separate labels -- Grafana can interpolate a
    variable's display text into a title, but cannot assemble one from several labels.

    Emits nothing unless smartctl actually returned identifying data. When it cannot
    open the device -- no privileges, a controller that needs an explicit -d type -- it
    still exits 0 with a well-formed JSON document that simply has no identity fields,
    and filling those in with placeholders would put "unknown model · no serial" in a
    dashboard row title. Emitting nothing instead lets the `smartctl:ident` recording
    rule fall back to "identity pending", which is the truth.
    """
    model = (data.get("model_name") or data.get("scsi_model_name")
             or data.get("scsi_product") or "")
    serial = data.get("serial_number") or ""
    capacity = human_capacity((data.get("user_capacity") or {}).get("bytes"))

    if not (model or serial or capacity):
        warn(f"{name}: no identity data; is smartctl able to open the device?")
        return

    parts = [p for p in (capacity, technology(data), model or "unknown model",
                         serial or "no serial") if p]
    reg.add("smart_drive_label", 1, {"device": name, "label": " · ".join(parts)},
            "Drive identity rendered for display; the value is always 1.")


def collect_device(reg: Registry, dev: str) -> bool:
    # -i rides along on the same invocation, so identity costs no extra smartctl run
    # and no extra chance of waking an idle disk.
    data = run_json(["smartctl", "-j", "-i", "-l", "selftest", dev])
    if data is None:
        return False

    name = os.path.basename(dev)
    collect_identity(reg, name, data)
    counts: dict[tuple[str, str], int] = defaultdict(int)
    totals: dict[str, int] = defaultdict(int)
    # smartctl lists newest first, so the first entry of each type is the latest.
    latest: dict[str, tuple[str, bool | None, float | None]] = {}

    entries = 0
    for entry in iter_entries(data):
        entries += 1
        ttype, status_text, passed, hours = entry_fields(entry)
        counts[(ttype, status_text)] += 1
        totals[ttype] += 1
        latest.setdefault(ttype, (status_text, passed, hours))

    for (ttype, status_text), n in sorted(counts.items()):
        reg.add("smart_selftest_result_count", n,
                {"device": name, "test_type": ttype, "result": status_text},
                "Self-test log entries by type and reported result.")

    for ttype, n in sorted(totals.items()):
        reg.add("smart_selftest_count", n, {"device": name, "test_type": ttype},
                "Self-test log entries of this type.")

    for ttype, (status_text, passed, hours) in sorted(latest.items()):
        lbl = {"device": name, "test_type": ttype}
        if passed is not None:
            reg.add("smart_selftest_last_passed", int(bool(passed)), lbl,
                    "Whether the most recent self-test of this type passed.")
        reg.add("smart_selftest_last_lifetime_hours", hours, lbl,
                "Drive powered-on hours at which the most recent self-test of this "
                "type ran. Subtract from the drive's current power-on hours to get "
                "how long ago that was.")
        reg.add("smart_selftest_last_status", 1, {**lbl, "status": status_text},
                "Reported status text of the most recent self-test, 1 for the "
                "current value.")

    reg.add("smart_selftest_log_entries", entries, {"device": name},
            "Total entries found in the drive's self-test log.")
    return True


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("--device", action="append", default=[],
                        help="device to inspect; repeatable. Default: smartctl --scan")
    args = parser.parse_args()

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

    devices = args.device or scan_devices()
    if not devices:
        warn("no devices found")

    seen = 0
    for dev in devices:
        try:
            if collect_device(reg, dev):
                seen += 1
        except Exception as exc:  # noqa: BLE001 - one bad disk must not lose the rest
            warn(f"{dev}: {exc!r}")

    reg.add("smart_selftest_devices", seen, {},
            "Devices successfully queried for self-test history.")
    reg.add("smart_selftest_collection_duration_seconds", time.time() - started, {},
            "Time this collector took to gather self-test history.")

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