#!/bin/bash

set -o pipefail

declare -r config_file='/etc/apt-zfs-snapshot.conf'
declare -r zfs_property='space.quico:auto-snapshot'
declare -r zfs_property_value='true'
declare -r zfs_max_name_length=255

if [[ -r "${config_file}" ]]; then
    # shellcheck disable=SC1090
    source "${config_file}"
fi

dry_run="${dry_run:-false}"
important_names="${important_names:-linux-image(-.*)?|linux-headers(-.*)?|systemd|zfs-(dkms|utils)}"
snapshots_trivial_keep="${snapshots_trivial_keep:-25}"
snapshots_important_keep="${snapshots_important_keep:-10}"
trivial_suffix="${trivial_suffix:-trv}"
important_suffix="${important_suffix:-imp}"
packages_max_length="${packages_max_length:-30}"
snap_only_local_datasets="${snap_only_local_datasets:-true}"
snapshot_roots="${snapshot_roots:-}"
field_separator="${field_separator:-_}"
snapshot_prefix="${snapshot_prefix:-apt}"
date_format="${date_format:-%F-%H%M}"
timezone="${timezone:-Etc/UTC}"
install_suffix="${install_suffix:-inst}"
remove_suffix="${remove_suffix:-rmvl}"
upgrade_suffix="${upgrade_suffix:-upgr}"
downgrade_suffix="${downgrade_suffix:-down}"
mixed_suffix="${mixed_suffix:-mixd}"
package_separator="${package_separator:-:}"

declare -a packages operations successful_datasets
declare existing_snapshots

print_msg() {
    local level="$1" message="$2"
    printf '[%s] %s\n' "${level}" "${message}" >&2
}

add_package() {
    local package="$1" operation="$2"
    [[ -n "${package}" ]] || return 0
    packages+=("${package//+/_}")
    operations+=("${operation}")
}

parse_apt_input() {
    local fd="${APT_HOOK_INFO_FD:-0}" line package old_version direction new_version action
    local protocol='1' in_records='false'

    while IFS= read -r line <&"${fd}"; do
        if [[ "${line}" == VERSION\ * ]]; then
            protocol="${line#VERSION }"
            continue
        fi
        if [[ "${protocol}" == '1' && -n "${line}" ]]; then
            package="${line##*/}"
            package="${package%.deb}"
            package="${package%_*_*}"
            add_package "${package}" 'install'
            continue
        fi
        if [[ "${in_records}" != 'true' ]]; then
            [[ -z "${line}" ]] && in_records='true'
            continue
        fi

        if [[ "${protocol}" == '1' ]]; then
            package="${line##*/}"
            package="${package%.deb}"
            package="${package%_*_*}"
            add_package "${package}" 'install'
            continue
        fi

        read -r package old_version _old_arch _old_multiarch direction new_version _new_arch _new_multiarch action <<<"${line}"
        [[ -n "${package}" ]] || continue
        if [[ "${action}" == '**CONFIGURE**' ]]; then
            continue
        elif [[ "${action}" == '**REMOVE**' ]]; then
            add_package "${package}" 'remove'
        elif [[ "${direction}" == '>' ]]; then
            add_package "${package}" 'downgrade'
        elif [[ "${old_version}" == '-' ]]; then
            add_package "${package}" 'install'
        elif [[ "${direction}" == '<' ]]; then
            add_package "${package}" 'upgrade'
        fi
    done
}

get_datasets() {
    local dataset value candidate parent
    local -a candidates
    local mounted_datasets
    if [[ -n "${snapshot_roots}" ]]; then
        read -r -a candidates <<<"${snapshot_roots}"
        snappable_datasets=("${candidates[@]}")
        return 0
    fi

    mounted_datasets="$(findmnt -rn -t zfs -o SOURCE)"
    while IFS=$'\t' read -r dataset value; do
        [[ "${value}" == "${zfs_property_value}" && "${dataset}" != *@* ]] || continue
        if [[ "${snap_only_local_datasets}" != 'true' ]] || grep -Fxq -- "${dataset}" <<<"${mounted_datasets}"; then
            candidates+=("${dataset}")
        fi
    done < <(zfs get -H -t filesystem,volume -o name,value "${zfs_property}")

    # An inherited property marks every descendant. Keep only the highest
    # selected dataset because the snapshot operation is recursive.
    for dataset in "${candidates[@]}"; do
        for parent in "${candidates[@]}"; do
            [[ "${dataset}" == "${parent}"/* ]] && continue 2
        done
        snappable_datasets+=("${dataset}")
    done
}

operation_suffix() {
    local -A seen=()
    local operation
    for operation in "${operations[@]}"; do seen["${operation}"]=1; done
    if (( ${#seen[@]} > 1 )); then printf '%s' "${mixed_suffix}"; return; fi
    case "${operations[0]:-upgrade}" in
        install) printf '%s' "${install_suffix}" ;;
        remove) printf '%s' "${remove_suffix}" ;;
        downgrade) printf '%s' "${downgrade_suffix}" ;;
        *) printf '%s' "${upgrade_suffix}" ;;
    esac
}

severity_and_packages() {
    local package important='false'
    local -a important_packages trivial_packages
    for package in "${packages[@]}"; do
        if grep -Piq -- "^(${important_names})$" <<<"${package}"; then
            important_packages+=("${package}")
            important='true'
        else
            trivial_packages+=("${package}")
        fi
    done
    if [[ "${important}" == 'true' ]]; then
        severity="${important_suffix}"
        package_list="$(IFS="${package_separator}"; printf '%s' "${important_packages[*]}")"
        [[ -n "${trivial_packages[*]}" ]] && package_list+="${package_separator}$(IFS="${package_separator}"; printf '%s' "${trivial_packages[*]}")"
    else
        severity="${trivial_suffix}"
        package_list="$(IFS="${package_separator}"; printf '%s' "${trivial_packages[*]}")"
    fi
}

shorten_packages() {
    local value="${package_list}" limit="${packages_max_length}" dataset fixed_length available
    for dataset in "${snappable_datasets[@]}"; do
        fixed_length="${#dataset}"
        fixed_length=$((fixed_length + 1 + ${#snapshot_prefix} + ${#field_separator} + ${#date_string} + ${#field_separator} + 10 + ${#field_separator} + 4 + ${#severity} + ${#field_separator} + 5))
        available=$((zfs_max_name_length - fixed_length))
        (( available < limit )) && limit="${available}"
    done
    (( limit < 0 )) && limit=0
    (( limit > 0 )) || { package_list=''; return; }
    while (( ${#value} > limit )) && [[ "${value}" == *"${package_separator}"* ]]; do
        value="${value%"${package_separator}"*}"
    done
    if (( ${#value} > limit )); then
        if (( limit <= 3 )); then value="${value:0:limit}"; else value="${value:0:limit-3}..."; fi
    fi
    package_list="${value}"
}

existing_snapshot() { grep -Fxq -- "$1" <<<"${existing_snapshots}"; }

make_snapshot_name() {
    local counter="$1" dataset base
    names=()
    for dataset in "${snappable_datasets[@]}"; do
        base="${dataset}@${snapshot_prefix}${field_separator}${date_string}${field_separator}${counter}${field_separator}op:${operation_suffix_result}${field_separator}sev:${severity}${field_separator}pkgs:${package_list}"
        (( ${#base} <= zfs_max_name_length )) || { print_msg ERR "Snapshot name exceeds ${zfs_max_name_length} characters: ${base}"; return 2; }
        names+=("${base}")
        existing_snapshot "${base}" && return 1
    done
    return 0
}

prune_snapshots() {
    local dataset limit name count
    [[ "${dry_run}" == 'true' ]] && return 0
    limit="${snapshots_trivial_keep}"
    [[ "${severity}" == "${important_suffix}" ]] && limit="${snapshots_important_keep}"
    for dataset in "${successful_datasets[@]}"; do
        # Query only the root dataset. Destroying its recursive snapshot also
        # removes the matching descendant snapshots as one snapshot tree.
        mapfile -t snapshots < <(zfs list -H -t snapshot -o name -s creation "${dataset}" | grep -F "${dataset}@${snapshot_prefix}${field_separator}" | grep -F "${field_separator}sev:${severity}${field_separator}")
        count="${#snapshots[@]}"
        while (( count > limit )); do
            name="${snapshots[0]}"
            zfs destroy -r "${name}" || { print_msg WARN "Failed to destroy ${name}"; break; }
            snapshots=("${snapshots[@]:1}")
            ((count--))
            print_msg INFO "Destroyed old snapshot ${name}"
        done
    done
}

main() {
    local counter=0
    parse_apt_input
    (( ${#packages[@]} > 0 )) || { print_msg INFO 'No package actions received; skipping snapshot'; return 0; }
    get_datasets
    (( ${#snappable_datasets[@]} > 0 )) || { print_msg INFO "No eligible ZFS datasets; skipping snapshot"; return 0; }
    existing_snapshots="$(zfs list -H -t snapshot -o name)"
    date_string="$(TZ="${timezone}" date +"${date_format}")"
    severity_and_packages
    shorten_packages
    operation_suffix_result="$(operation_suffix)"
    while :; do
        make_snapshot_name "$((++counter))"
        case "$?" in
            0) break ;;
            1) ;;
            *) return 1 ;;
        esac
    done
    if [[ "${dry_run}" == 'true' ]]; then
        printf '[INFO] Dry-run, would create:\n' >&2
        printf '  %s\n' "${names[@]}" >&2
        return 0
    fi
    successful_datasets=()
    local name dataset snapshot_failed='false'
    for name in "${names[@]}"; do
        dataset="${name%%@*}"
        if zfs snapshot -r "${name}"; then
            successful_datasets+=("${dataset}")
            printf '[INFO] Created ZFS snapshots for %s:\n' "${dataset}" >&2
            printf '  %s\n' "${name}" >&2
        else
            snapshot_failed='true'
            print_msg WARN "ZFS snapshot failed for ${dataset}"
        fi
    done
    [[ "${snapshot_failed}" == 'false' ]] || print_msg WARN 'Snapshot set is incomplete'
    ((${#successful_datasets[@]} > 0)) && prune_snapshots
}

main