#!/usr/bin/env python3
"""Exhaustion-trajectory disk watcher with safe /tmp remediation."""

import argparse
import errno
import fnmatch
import os
import stat
import subprocess
import sys
import time
from collections import defaultdict, deque


def env_int(key, default):
    try:
        return int(os.environ.get(key, default))
    except (TypeError, ValueError):
        return default


# Keep the fallback watcher responsive without polling the filesystem every few seconds.
POLL = env_int("POLL", 60)
RUNWAY_WINDOW = env_int("RUNWAY_WINDOW", 1800)
RATE_WINDOW = env_int("RATE_WINDOW", 600)
NOTICE_ETA_S = env_int("NOTICE_ETA_S", 360)
CRIT_ETA_S = env_int("CRIT_ETA_S", 90)
SCAN_ETA_S = env_int("SCAN_ETA_S", 1800)
SCAN_FREE_PCT = env_int("SCAN_FREE_PCT", 20)
BYTE_FLOOR_PCT = env_int("BYTE_FLOOR_PCT", 3)
INODE_FLOOR_PCT = env_int("INODE_FLOOR_PCT", 3)
GROW_WINDOW = env_int("GROW_WINDOW", 1800)
TMP_STALE_HOURS = env_int("TMP_STALE_HOURS", 6)
MAX_CANDIDATES = env_int("MAX_CANDIDATES", 48)
RECENT_CHILDREN = env_int("RECENT_CHILDREN", 12)
MIN_GROW_BYTES = env_int("MIN_GROW_BYTES", 268435456)
MIN_GROW_INODES = env_int("MIN_GROW_INODES", 20000)
MIN_SCAN_AGE = env_int("MIN_SCAN_AGE", 60)

UID_ME = os.getuid()
HOME = os.path.expanduser("~")
STATE = os.path.join(os.environ.get("XDG_STATE_HOME", os.path.expanduser("~/.local/state")), "mem-guard")
LOG = os.path.join(STATE, "events.log")
os.makedirs(STATE, exist_ok=True)

WATCH_TARGETS = ["/", HOME, "/tmp", "/boot", "/boot/efi"]
TMP_PROTECT_PATTERNS = {
    ".X11-unix",
    ".XIM-unix",
    ".font-unix",
    ".ICE-unix",
    ".Test-unix",
    "systemd-private-*",
    "ssh-*",
    "tmux-*",
    "pulse-*",
}

HAVE_NOTIFY = subprocess.run(["sh", "-c", "command -v notify-send"], capture_output=True).returncode == 0
HAVE_WAIT = False
if HAVE_NOTIFY:
    help_run = subprocess.run(["notify-send", "--help"], capture_output=True, text=True)
    HAVE_WAIT = "--wait" in (help_run.stdout + help_run.stderr)


def logline(message):
    with open(LOG, "a", encoding="utf-8") as handle:
        handle.write(time.strftime("%F %T") + "  " + message + "\n")


def parse_mountinfo():
    mounts = []
    with open("/proc/self/mountinfo", encoding="utf-8") as handle:
        for raw in handle:
            raw = raw.rstrip("\n")
            if " - " not in raw:
                continue
            left, right = raw.split(" - ", 1)
            parts = left.split()
            tail = right.split()
            if len(parts) < 5 or len(tail) < 3:
                continue
            mountpoint = parts[4].replace("\\040", " ")
            mounts.append({
                "mountpoint": mountpoint,
                "fstype": tail[0],
                "source": tail[1],
            })
    mounts.sort(key=lambda item: len(item["mountpoint"]), reverse=True)
    return mounts


def resolve_mount(path, mountinfo):
    real = os.path.realpath(path)
    for item in mountinfo:
        mountpoint = item["mountpoint"]
        if mountpoint == "/":
            matched = real.startswith("/")
        else:
            matched = real == mountpoint or real.startswith(mountpoint.rstrip("/") + "/")
        if matched:
            try:
                st = os.stat(mountpoint)
            except OSError:
                return None
            info = dict(item)
            info["st_dev"] = st.st_dev
            info["is_tmpfs"] = item["fstype"] in {"tmpfs", "ramfs"}
            return info
    return None


def watch_mounts():
    mountinfo = parse_mountinfo()
    selected = {}
    for target in WATCH_TARGETS:
        info = resolve_mount(target, mountinfo)
        if info is None:
            continue
        selected[info["mountpoint"]] = info
    return [selected[key] for key in sorted(selected)]


def stat_mount(mount):
    stats = os.statvfs(mount["mountpoint"])
    total_bytes = stats.f_blocks * stats.f_frsize
    free_bytes = stats.f_bavail * stats.f_frsize
    total_inodes = getattr(stats, "f_files", 0)
    free_inodes = getattr(stats, "f_favail", 0) or getattr(stats, "f_ffree", 0)
    return {
        "bytes_total": total_bytes,
        "bytes_free": free_bytes,
        "bytes_floor": total_bytes * BYTE_FLOOR_PCT // 100,
        "bytes_pct": pct(free_bytes, total_bytes),
        "inodes_total": total_inodes,
        "inodes_free": free_inodes,
        "inodes_floor": total_inodes * INODE_FLOOR_PCT // 100,
        "inodes_pct": pct(free_inodes, total_inodes),
    }


def pct(value, total):
    if total <= 0:
        return 100
    return int((100 * value) / total)


def fmt_bytes(value):
    units = ["B", "K", "M", "G", "T"]
    number = float(max(0, value))
    for unit in units:
        if number < 1024.0 or unit == units[-1]:
            if unit == "B":
                return f"{int(number)}{unit}"
            return f"{number:.1f}{unit}"
        number /= 1024.0
    return f"{number:.1f}P"


def fmt_count(value):
    if value >= 1000000:
        return f"{value / 1000000:.1f}M"
    if value >= 1000:
        return f"{value / 1000:.1f}K"
    return str(value)


def runway_eta(history, now, free_value, floor_value):
    while history and now - history[0][0] > RUNWAY_WINDOW:
        history.popleft()
    ref = None
    for ts, value in history:
        if now - ts <= RATE_WINDOW:
            ref = (ts, value)
            break
    if ref is None and history:
        ref = history[0]
    if ref is None or now - ref[0] < MIN_SCAN_AGE:
        return None
    burn = (ref[1] - free_value) / (now - ref[0])
    if burn <= 0:
        return None
    if free_value <= floor_value:
        return 0.0
    return max(0.0, (free_value - floor_value)) / burn


def recent_children(base, mount_dev, limit, nested=False):
    items = []
    try:
        with os.scandir(base) as scan:
            for entry in scan:
                try:
                    st = entry.stat(follow_symlinks=False)
                except OSError:
                    continue
                if st.st_dev != mount_dev:
                    continue
                if entry.is_symlink():
                    continue
                items.append((st.st_mtime, entry.path))
    except OSError:
        return []
    items.sort(reverse=True)
    out = [path for _, path in items[:limit]]
    if not nested:
        return out
    extra = []
    for path in out[: max(4, limit // 2)]:
        if not os.path.isdir(path):
            continue
        for name in ("node_modules", ".cache", "log", "logs", "dist", "build"):
            candidate = os.path.join(path, name)
            try:
                st = os.lstat(candidate)
            except OSError:
                continue
            if st.st_dev == mount_dev and not stat.S_ISLNK(st.st_mode):
                extra.append(candidate)
    return out + extra


def static_candidates(mount):
    mp = mount["mountpoint"]
    mount_dev = mount["st_dev"]
    candidates = []
    if mp == "/tmp":
        candidates.extend(recent_children("/tmp", mount_dev, MAX_CANDIDATES))
    elif mp == "/boot":
        candidates.extend(recent_children("/boot", mount_dev, MAX_CANDIDATES))
    elif mp == "/":
        for path in ("/var/tmp", "/var/log", "/var/cache", "/tmp", os.path.join(HOME, ".cache")):
            if path_exists_on_mount(path, mount_dev):
                candidates.append(path)
    if mp == os.path.realpath(HOME) or HOME == mp or HOME.startswith(mp.rstrip("/") + "/"):
        for path in (
            os.path.join(HOME, ".cache"),
            os.path.join(HOME, ".local", "state"),
            os.path.join(HOME, ".local", "share"),
            os.path.join(HOME, ".local", "share", "Trash"),
            os.path.join(HOME, ".npm"),
            os.path.join(HOME, ".cargo"),
        ):
            if path_exists_on_mount(path, mount_dev):
                candidates.append(path)
        candidates.extend(recent_children(HOME, mount_dev, RECENT_CHILDREN, nested=True))
    if mp not in {"/tmp", "/boot"}:
        candidates.extend(recent_children(mp, mount_dev, min(RECENT_CHILDREN, 8)))
    return dedupe_paths(candidates)[:MAX_CANDIDATES]


def dedupe_paths(paths):
    seen = set()
    out = []
    for path in paths:
        real = os.path.realpath(path)
        if real in seen:
            continue
        seen.add(real)
        out.append(real)
    return out


def path_exists_on_mount(path, mount_dev):
    try:
        st = os.lstat(path)
    except OSError:
        return False
    return st.st_dev == mount_dev and not stat.S_ISLNK(st.st_mode)


def same_mount(path, mount_dev):
    try:
        return os.lstat(path).st_dev == mount_dev
    except OSError:
        return False


def disk_bytes_for_stat(st):
    blocks = getattr(st, "st_blocks", 0)
    if blocks:
        return blocks * 512
    return st.st_size


def measure_path(path, mount_dev):
    total_bytes = 0
    total_entries = 0
    newest_mtime = 0.0
    stack = [path]
    seen = set()
    while stack:
        current = stack.pop()
        try:
            st = os.lstat(current)
        except OSError:
            continue
        if st.st_dev != mount_dev:
            continue
        key = (st.st_dev, st.st_ino)
        if key in seen:
            continue
        seen.add(key)
        total_entries += 1
        total_bytes += disk_bytes_for_stat(st)
        newest_mtime = max(newest_mtime, st.st_mtime)
        if not stat.S_ISDIR(st.st_mode):
            continue
        try:
            with os.scandir(current) as scan:
                for entry in scan:
                    if entry.is_symlink():
                        continue
                    stack.append(entry.path)
        except OSError:
            continue
    return {
        "path": path,
        "bytes": total_bytes,
        "entries": total_entries,
        "mtime": newest_mtime,
    }


def sample_culprits(now, mount, metric_key, suspect_hist):
    candidates = static_candidates(mount)
    samples = []
    for path in candidates:
        sample = measure_path(path, mount["st_dev"])
        if sample["bytes"] <= 0 and sample["entries"] <= 1:
            continue
        suspect_hist[path].append((now, sample["bytes"], sample["entries"]))
        trim_suspect_history(suspect_hist[path], now)
        samples.append(sample)
    if not samples:
        return []
    ranked = []
    for sample in samples:
        hist = suspect_hist[sample["path"]]
        min_bytes = min(item[1] for item in hist) if hist else sample["bytes"]
        min_entries = min(item[2] for item in hist) if hist else sample["entries"]
        growth_bytes = sample["bytes"] - min_bytes
        growth_entries = sample["entries"] - min_entries
        if metric_key == "bytes":
            score = growth_bytes if growth_bytes >= MIN_GROW_BYTES else sample["bytes"]
            fallback = sample["bytes"]
        else:
            score = growth_entries if growth_entries >= MIN_GROW_INODES else sample["entries"]
            fallback = sample["entries"]
        ranked.append((score, fallback, growth_bytes, growth_entries, sample))
    ranked.sort(key=lambda item: (-item[0], -item[1], item[4]["path"]))
    return ranked


def trim_suspect_history(history, now):
    while history and now - history[0][0] > GROW_WINDOW:
        history.popleft()


def metric_reason(mount, metric_key, stats, eta):
    place = mount["mountpoint"]
    tmpfs_note = " tmpfs" if mount["is_tmpfs"] else ""
    if metric_key == "bytes":
        free_text = f"{fmt_bytes(stats['bytes_free'])} free of {fmt_bytes(stats['bytes_total'])}"
        what = f"{place}{tmpfs_note} bytes"
    else:
        free_text = f"{fmt_count(stats['inodes_free'])} free of {fmt_count(stats['inodes_total'])} inodes"
        what = f"{place} inodes"
    if eta is not None and eta < CRIT_ETA_S:
        return "CRITICAL", "critical", f"{what} nearly exhausted — ~{int(eta)}s ETA ({free_text})."
    if eta is not None and eta < NOTICE_ETA_S:
        return "notice", "normal", f"{what} falling — ~{int(max(1, eta // 60))} min ETA ({free_text})."
    return None, None, None


def should_sample(metric_key, free_pct, eta, mount):
    if mount["mountpoint"] == "/tmp":
        return True
    if eta is not None and eta < SCAN_ETA_S:
        return True
    if free_pct <= SCAN_FREE_PCT:
        return True
    if metric_key == "inodes" and free_pct <= max(SCAN_FREE_PCT, 30):
        return True
    return False


def culprit_lines(metric_key, ranked):
    lines = []
    for _, _, growth_bytes, growth_entries, sample in ranked[:3]:
        label = sample["path"]
        if metric_key == "bytes":
            growth = f"+{fmt_bytes(growth_bytes)}" if growth_bytes > 0 else "flat"
            now_text = fmt_bytes(sample["bytes"])
        else:
            growth = f"+{fmt_count(growth_entries)}" if growth_entries > 0 else "flat"
            now_text = fmt_count(sample["entries"])
        lines.append(f"{label} ({growth}, now {now_text})")
    return lines


def notify_alert(tier, urgency, mount, metric_key, reason, ranked, action_kind):
    culprit = ranked[0][4]["path"] if ranked else mount["mountpoint"]
    summary = f"Disk {tier} — {mount['mountpoint']} {metric_key}"
    body = [reason]
    if ranked:
        body.append("Culprit candidates:")
        body.extend(culprit_lines(metric_key, ranked))
    logline(f"{tier}  mount={mount['mountpoint']} metric={metric_key} culprit={culprit}  {reason}")
    if not HAVE_NOTIFY:
        return None
    if HAVE_WAIT and action_kind == "tmp-clean":
        run = subprocess.run(
            [
                "notify-send",
                "--wait",
                "-u",
                urgency,
                "-t",
                "25000",
                "-A",
                "clean=Clean stale /tmp junk",
                "-A",
                "dismiss=Dismiss",
                summary,
                "\n".join(body),
            ],
            capture_output=True,
            text=True,
        )
        return run.stdout.strip() or None
    subprocess.run(["notify-send", "-u", urgency, summary, "\n".join(body)])
    return None


def protected_tmp_name(name):
    return any(fnmatch.fnmatchcase(name, pattern) for pattern in TMP_PROTECT_PATTERNS)


def open_paths_under_tmp(tmp_dev):
    paths = []
    for proc_name in os.listdir("/proc"):
        if not proc_name.isdigit():
            continue
        fd_dir = os.path.join("/proc", proc_name, "fd")
        try:
            entries = os.listdir(fd_dir)
        except OSError:
            continue
        for entry in entries:
            target_path = os.path.join(fd_dir, entry)
            try:
                target = os.readlink(target_path)
            except OSError:
                continue
            if " (deleted)" in target:
                target = target.split(" (deleted)", 1)[0]
            if not target.startswith("/tmp/") and target != "/tmp":
                continue
            try:
                st = os.stat(target)
            except OSError:
                continue
            if st.st_dev == tmp_dev:
                paths.append(os.path.realpath(target))
    return paths


def has_open_descendant(path, open_paths):
    real = os.path.realpath(path)
    prefix = real.rstrip("/") + "/"
    for target in open_paths:
        if target == real or target.startswith(prefix):
            return True
    return False


def eligible_tmp_entry(path, tmp_dev, cutoff, open_paths):
    name = os.path.basename(path)
    if protected_tmp_name(name):
        return False
    try:
        st = os.lstat(path)
    except OSError:
        return False
    if st.st_dev != tmp_dev or st.st_uid != UID_ME:
        return False
    if stat.S_ISLNK(st.st_mode) or stat.S_ISSOCK(st.st_mode) or stat.S_ISFIFO(st.st_mode):
        return False
    if stat.S_ISCHR(st.st_mode) or stat.S_ISBLK(st.st_mode):
        return False
    newest_touch = max(st.st_mtime, st.st_atime, st.st_ctime)
    if newest_touch > cutoff:
        return False
    if has_open_descendant(path, open_paths):
        return False
    return True


def remove_tree(path, mount_dev):
    removed_bytes = 0
    removed_entries = 0
    pending_dirs = []
    stack = [path]
    while stack:
        current = stack.pop()
        try:
            st = os.lstat(current)
        except OSError:
            continue
        if st.st_dev != mount_dev or stat.S_ISLNK(st.st_mode):
            continue
        if stat.S_ISSOCK(st.st_mode) or stat.S_ISFIFO(st.st_mode) or stat.S_ISCHR(st.st_mode) or stat.S_ISBLK(st.st_mode):
            continue
        if stat.S_ISDIR(st.st_mode):
            pending_dirs.append((current, st))
            try:
                with os.scandir(current) as scan:
                    for entry in scan:
                        stack.append(entry.path)
            except OSError:
                continue
            continue
        try:
            os.unlink(current)
            removed_bytes += disk_bytes_for_stat(st)
            removed_entries += 1
        except OSError:
            continue
    pending_dirs.sort(key=lambda item: item[0].count(os.sep), reverse=True)
    for directory, st in pending_dirs:
        try:
            os.rmdir(directory)
            removed_bytes += disk_bytes_for_stat(st)
            removed_entries += 1
        except OSError:
            continue
    return removed_bytes, removed_entries


def clean_tmp():
    try:
        tmp_st = os.stat("/tmp")
    except OSError as exc:
        return False, f"/tmp unavailable: {exc.strerror}"
    tmp_dev = tmp_st.st_dev
    cutoff = time.time() - (TMP_STALE_HOURS * 3600)
    open_paths = open_paths_under_tmp(tmp_dev)
    removed_paths = 0
    removed_bytes = 0
    removed_entries = 0
    try:
        with os.scandir("/tmp") as scan:
            for entry in scan:
                path = entry.path
                if not eligible_tmp_entry(path, tmp_dev, cutoff, open_paths):
                    continue
                bytes_out, entries_out = remove_tree(path, tmp_dev)
                if bytes_out or entries_out:
                    removed_paths += 1
                    removed_bytes += bytes_out
                    removed_entries += entries_out
    except OSError as exc:
        return False, f"/tmp scan failed: {exc.strerror}"
    if not removed_paths:
        return False, f"Nothing safe to delete in /tmp older than {TMP_STALE_HOURS}h."
    return True, f"Removed {removed_paths} stale /tmp entries, {fmt_bytes(removed_bytes)}, {fmt_count(removed_entries)} filesystem entries."


metric_history = defaultdict(deque)
suspect_history = defaultdict(deque)
alerted_tier = {}
TIER_RANK = {"notice": 1, "CRITICAL": 2}


def evaluate_mount(now, mount):
    stats = stat_mount(mount)
    alerts = []

    metric_history[(mount["mountpoint"], "bytes")].append((now, stats["bytes_free"]))
    eta_bytes = runway_eta(metric_history[(mount["mountpoint"], "bytes")], now, stats["bytes_free"], stats["bytes_floor"])
    tier, urgency, reason = metric_reason(mount, "bytes", stats, eta_bytes)
    if tier:
        alerts.append(("bytes", stats["bytes_pct"], eta_bytes, tier, urgency, reason))

    if stats["inodes_total"] > 0:
        metric_history[(mount["mountpoint"], "inodes")].append((now, stats["inodes_free"]))
        eta_inodes = runway_eta(metric_history[(mount["mountpoint"], "inodes")], now, stats["inodes_free"], stats["inodes_floor"])
        tier, urgency, reason = metric_reason(mount, "inodes", stats, eta_inodes)
        if tier:
            alerts.append(("inodes", stats["inodes_pct"], eta_inodes, tier, urgency, reason))

    alerted_metrics = set()
    for metric_key, free_pct, eta, tier, urgency, reason in alerts:
        alerted_metrics.add(metric_key)
        alert_key = (mount["mountpoint"], metric_key)
        if TIER_RANK[tier] <= TIER_RANK.get(alerted_tier.get(alert_key), 0):
            continue
        alerted_tier[alert_key] = tier
        ranked = []
        if should_sample(metric_key, free_pct, eta, mount):
            ranked = sample_culprits(now, mount, metric_key, suspect_history)
        action_kind = "tmp-clean" if mount["mountpoint"] == "/tmp" else None
        action = notify_alert(tier, urgency, mount, metric_key, reason, ranked, action_kind)
        if action == "clean":
            ok, message = clean_tmp()
            state = "CLEANED" if ok else "CLEANUP-SKIP"
            logline(f"{state}  {message}")
            if HAVE_NOTIFY:
                subprocess.run(["notify-send", "-u", "normal", "Disk guard cleanup", message])

    for metric_key in ("bytes", "inodes"):
        if metric_key not in alerted_metrics:
            alerted_tier.pop((mount["mountpoint"], metric_key), None)


def once():
    now = time.time()
    for mount in watch_mounts():
        try:
            evaluate_mount(now, mount)
        except OSError as exc:
            if exc.errno not in {errno.ENOENT, errno.EACCES}:
                logline(f"ERROR  mount={mount['mountpoint']}  {exc}")


def main():
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--once", action="store_true", help="Run one evaluation pass and exit.")
    parser.add_argument("--cleanup-tmp", action="store_true", help="Run the safe /tmp cleanup action and exit.")
    args = parser.parse_args()

    if args.cleanup_tmp:
        ok, message = clean_tmp()
        print(message)
        return 0 if ok else 1
    if args.once:
        once()
        return 0

    while True:
        once()
        time.sleep(POLL)


if __name__ == "__main__":
    try:
        raise SystemExit(main())
    except KeyboardInterrupt:
        sys.exit(0)