#!/usr/bin/env python3
"""Thin out the auto-backup_ snapshots of one received dataset, independent of the source.

Usage: prune-snapshots [--dry-run] DATASET

Keeps the newest snapshot per day, ISO week and month for the most recent days, weeks and months
in /etc/backup-server/retention.json, and always the newest snapshot, which the next incremental
receive builds on. Other snapshots of the dataset are never touched. Called by the clients'
backup_path_via_zfs after each successful receive, so it runs while this server is awake.
Exit codes: 0 ok, 1 zfs failed, 2 usage.
"""
from __future__ import annotations

import datetime
import json
import re
import subprocess
import sys

PREFIX = "auto-backup_"
CONFIG = "/etc/backup-server/retention.json"
# only datasets received from clients (tank/<client uuid>/...), as the backup user calls this via sudo
DATASET_RE = re.compile(r"\Atank/[0-9a-f-]{36}/[A-Za-z0-9._/-]+\Z")
BUCKETS = {
    "daily": lambda d: d.strftime("%Y-%m-%d"),
    "weekly": lambda d: "%d-W%02d" % d.isocalendar()[:2],
    "monthly": lambda d: d.strftime("%Y-%m"),
}


def zfs(*args: str) -> str:
    return subprocess.run(["zfs", *args], capture_output=True, text=True, check=True, timeout=120).stdout


def keep(snapshots: list[tuple[str, int]], retention: dict[str, int]) -> set[str]:
    """Names to keep from (name, creation epoch) pairs: per granularity the newest snapshot of each
    of the most recent buckets, plus the newest snapshot overall."""
    newest_first = sorted(snapshots, key=lambda s: s[1], reverse=True)
    kept = {newest_first[0][0]} if newest_first else set()
    for granularity, count in retention.items():
        seen: set[str] = set()
        for name, creation in newest_first:
            bucket = BUCKETS[granularity](datetime.datetime.fromtimestamp(creation))
            if bucket not in seen and len(seen) < count:
                seen.add(bucket)
                kept.add(name)
    return kept


def main(argv: list[str]) -> int:
    dry_run = "--dry-run" in argv
    args = [a for a in argv if a != "--dry-run"]
    if len(args) != 1 or not DATASET_RE.match(args[0]):
        print(__doc__, file=sys.stderr)
        return 2
    dataset = args[0]
    with open(CONFIG, encoding="utf-8") as fh:
        retention = {k: int(v) for k, v in json.load(fh).items() if k in BUCKETS}
    try:
        lines = zfs("list", "-H", "-p", "-t", "snapshot", "-d", "1", "-o", "name,creation", dataset).splitlines()
    except subprocess.CalledProcessError as e:
        print(f"prune-snapshots: {e.stderr.strip()}", file=sys.stderr)
        return 1
    snapshots = []
    for line in lines:
        name, creation = line.split("\t")
        if name.split("@", 1)[1].startswith(PREFIX):
            snapshots.append((name, int(creation)))
    kept = keep(snapshots, retention)
    for name, _ in snapshots:
        if name in kept:
            continue
        print(f"{'would destroy' if dry_run else 'destroying'} {name}")
        if not dry_run:
            try:
                zfs("destroy", name)
            except subprocess.CalledProcessError as e:
                print(f"prune-snapshots: {e.stderr.strip()}", file=sys.stderr)
                return 1
    print(f"{dataset}: {len(kept)} of {len(snapshots)} {PREFIX} snapshots kept")
    return 0


if __name__ == "__main__":
    sys.exit(main(sys.argv[1:]))