#!/usr/bin/env python3
"""disk-usage-report.py —— 按挂载点输出磁盘使用率

用法：
    python3 disk-usage-report.py [-w 85] [--warn-root 80] [--warn-data 90]

退出码：0 = 均未超阈值；1 = 有挂载点超阈值；2 = 执行出错

说明：按原样提供，不附带任何保证。阈值应按挂载点用途分别设置，
      日志盘和数据盘的安全余量不同，全局单一阈值没有意义。
"""

import argparse
import shutil
import sys

# 挂载点的常用别名：--warn-log 匹配路径里含 /log 或 /logs 的挂载点
ALIAS_KEYS = {
    "root": ("/",),
    "log": ("log",),
    "data": ("data",),
    "home": ("home",),
    "var": ("var",),
}


def collect_mounts():
    """返回 [(mountpoint, total, used, free, percent)]，跳过伪文件系统。"""
    rows = []
    seen = set()
    with open("/proc/mounts", encoding="utf-8", errors="replace") as fh:
        for line in fh:
            parts = line.split()
            if len(parts) < 3:
                continue
            dev, mount, fstype = parts[0], parts[1], parts[2]
            if fstype in {
                "proc", "sysfs", "devtmpfs", "devpts", "tmpfs", "cgroup",
                "cgroup2", "overlay", "squashfs", "efivarfs", "securityfs",
                "debugfs", "tracefs", "pstore", "bpf", "autofs", "mqueue",
                "hugetlbfs", "configfs", "fusectl", "ramfs",
            }:
                continue
            if mount in seen:
                continue
            seen.add(mount)
            try:
                usage = shutil.disk_usage(mount)
            except OSError:
                continue
            if usage.total == 0:
                continue
            pct = usage.used * 100.0 / usage.total
            rows.append((mount, usage.total, usage.used, usage.free, pct))
    rows.sort(key=lambda r: r[0])
    return rows


def threshold_for(mount, default, overrides):
    for key, needles in ALIAS_KEYS.items():
        if key not in overrides:
            continue
        if any(n == mount or (n != "/" and n in mount) for n in needles):
            return overrides[key]
    return default


def human(n):
    units = ["B", "KiB", "MiB", "GiB", "TiB", "PiB"]
    value = float(n)
    for unit in units:
        if value < 1024 or unit == units[-1]:
            return f"{value:.1f} {unit}"
        value /= 1024
    return f"{n} B"


def main():
    ap = argparse.ArgumentParser(description="按挂载点输出磁盘使用率")
    ap.add_argument("-w", "--warn", type=float, default=85.0,
                    help="默认告警阈值（百分比，默认 85）")
    ap.add_argument("--warn-root", type=float, default=None, help="根分区阈值")
    ap.add_argument("--warn-log", type=float, default=None, help="日志盘阈值")
    ap.add_argument("--warn-data", type=float, default=None, help="数据盘阈值")
    ap.add_argument("--warn-home", type=float, default=None, help="home 盘阈值")
    ap.add_argument("--warn-var", type=float, default=None, help="var 盘阈值")
    args = ap.parse_args()

    overrides = {
        "root": args.warn_root,
        "log": args.warn_log,
        "data": args.warn_data,
        "home": args.warn_home,
        "var": args.warn_var,
    }
    overrides = {k: v for k, v in overrides.items() if v is not None}

    rows = collect_mounts()
    if not rows:
        print("错误：没有可检查的挂载点", file=sys.stderr)
        return 2

    print(f"{'挂载点':<28}{'容量':>12}{'已用':>12}{'可用':>12}{'使用率':>9}{'阈值':>8}")
    bad = 0
    for mount, total, used, free, pct in rows:
        limit = threshold_for(mount, args.warn, overrides)
        flag = "  << 超阈值" if pct >= limit else ""
        if flag:
            bad += 1
        print(f"{mount:<28}{human(total):>12}{human(used):>12}"
              f"{human(free):>12}{pct:>8.1f}%{limit:>7.0f}%{flag}")

    print(f"\n共 {len(rows)} 个挂载点，超阈值 {bad} 个。")
    return 1 if bad else 0


if __name__ == "__main__":
    try:
        sys.exit(main())
    except KeyboardInterrupt:
        sys.exit(130)
    except Exception as exc:  # noqa: BLE001 - 顶层兜底，保证退出码可判
        print(f"错误：{exc}", file=sys.stderr)
        sys.exit(2)