#!/bin/bash

declare -a pkgs
while read pkg; do
    pkgs+=("${pkg}")
done

declare conf_file
conf_file='/etc/pacman-zfs-snapshot.conf'

declare important_names snaps_trivial_keep snaps_important_keep snaps_trivial_suffix snaps_important_suffix
if [[ -r "${conf_file}" ]]; then
    source "${conf_file}"
fi

# User-defined
do_dry_run="${do_dry_run:-false}"
important_names="${important_names:-linux|systemd|zfs-(dkms|utils)}"
snaps_trivial_keep="${snaps_trivial_keep:-15}"
snaps_important_keep="${snaps_important_keep:-5}"
snaps_trivial_suffix="${snaps_trivial_suffix:-trv}"
snaps_important_suffix="${snaps_important_suffix:-imp}"
pkgs_list_max_length="${pkgs_list_max_length:-30}"
snap_only_local_datasets="${snap_only_local_datasets:-true}"
snap_field_separator="${snap_field_separator:-_}"
snap_name_prefix="${snap_name_prefix:-pacman}"
snap_date_format="${snap_date_format:-%F-%H%M}"
snap_timezone="${snap_timezone:-Etc/UTC}"
snap_op_installation_suffix="${snap_op_installation_suffix:-inst}"
snap_op_remove_suffix="${snap_op_remove_suffix:-rmvl}"
snap_op_upgrade_suffix="${snap_op_upgrade_suffix:-upgr}"

# Internal
declare zfs_prop pkg_separator max_zfs_snapshot_name_length color_reset color_lyellow color_red
zfs_prop='space.quico:auto-snapshot'
pkg_separator=':'
max_zfs_snapshot_name_length='255'
color_reset='\e[0m'
color_lyellow='\e[93m'
color_red='\e[31m'

declare operation conf_op_suffix
operation="${1}"
case "${operation}" in
    install)
        conf_op_suffix="${snap_op_installation_suffix}"
        ;;
    remove)
        conf_op_suffix="${snap_op_remove_suffix}"
        ;;
    upgrade)
        conf_op_suffix="${snap_op_upgrade_suffix}"
        ;;
esac

function pprint () {
    local style msg exit_code
    style="${1:?}"
    msg="${2:?}"
    exit_code="${3}"

    case "${style}" in
        warn)
            printf -- "${color_lyellow}"'[WARN]'"${color_reset}"' %s\n' "${msg}"
            ;;
        err)
            printf -- "${color_red}"'[ERR] '"${color_reset}"' %s\n' "${msg}"
            ;;
        info)
            printf -- '[INFO] %s\n' "${msg}"
            ;;
    esac

    [[ "${exit_code}" ]] && exit "${exit_code}"
}

function split_pkgs_by_importance () {
    local pkgs_in_transaction
    pkgs_in_transaction=("${@}")
    for pkg in "${pkgs_in_transaction[@]}"; do
        if grep -Piq -- '^('"${important_names}"')$' <<<"${pkg}"; then
            important_pkgs_in_transaction+=("${pkg}")
        else
            trivial_pkgs_in_transaction+=("${pkg}")
        fi
    done
}

function set_severity () {
    if [[ "${#important_pkgs_in_transaction[@]}" -ge '1' ]]; then
        severity="${snaps_important_suffix}"
    else
        severity="${snaps_trivial_suffix}"
    fi
}

function get_globally_snappable_datasets () {
    local datasets_list
    # For all datasets show their "${zfs_prop}" property; only print dataset
    # name in column 1 and property value in column 2. In awk limit this
    # list to datasets where tab-delimited column 2 has exact string
    # '^true$' then further limit output by eliminating snapshots from list,
    # i.e. dataset names that contain an '@' character.
    datasets_list="$(zfs get -H -o 'name,value' "${zfs_prop}" | \
                     awk -F'\t' '{if($2 ~ /^true$/ && $1 !~ /@/) print $1}')"
    while IFS= read -u10 -r dataset; do
        globally_snappable_datasets+=("${dataset}")
    done 10<<<"${datasets_list}"
}

function get_local_snappable_datasets () {
    local datasets_list
    datasets_list="$(findmnt --json --list --output 'fstype,source,target' | \
                     jq --raw-output '.[][] | select(.fstype=="zfs") | .source')"
    while IFS= read -u10 -r dataset; do
        local_snappable_datasets+=("${dataset}")
    done 10<<<"${datasets_list}"
}

function trim_globally_snappable_datasets () {
    for global_dataset in "${globally_snappable_datasets[@]}"; do
        for local_dataset in "${local_snappable_datasets[@]}"; do
            if grep -Piq -- '^'"${local_dataset}"'$' <<<"${global_dataset}"; then
                snappable_datasets+=("${global_dataset}")
            fi
        done
    done
}

function write_pkg_list_oneline () {
    if [[ "${severity}" == "${snaps_important_suffix}" ]]; then
        for pkg in "${important_pkgs_in_transaction[@]}"; do
            if [[ "${unabridged_pkg_list_oneline}" ]]; then
                unabridged_pkg_list_oneline="${unabridged_pkg_list_oneline}${pkg_separator}${pkg}"
            else
                unabridged_pkg_list_oneline="${pkg}"
            fi
        done
    fi
    if [[ "${#trivial_pkgs_in_transaction[@]}" -ge '1' ]]; then
        for pkg in "${trivial_pkgs_in_transaction[@]}"; do
            if [[ "${unabridged_pkg_list_oneline}" ]]; then
                unabridged_pkg_list_oneline="${unabridged_pkg_list_oneline}${pkg_separator}${pkg}"
            else
                unabridged_pkg_list_oneline="${pkg}"
            fi
        done
    fi
}

function find_max_dataset_name_length () {
    local longest_op_suffix op_suffix_string
    longest_op_suffix='0'
    for op_suffix in "${snap_op_installation_suffix}" "${snap_op_remove_suffix}" "${snap_op_upgrade_suffix}"; do
        if [[ "${#op_suffix}" -gt "${longest_op_suffix}" ]]; then
            longest_op_suffix="${#op_suffix}"
        fi
    done
    op_suffix_string="$(head -c "${longest_op_suffix}" '/dev/zero' | tr '\0' '_')"

    local longest_sev_suffix sev_suffix_string
    longest_sev_suffix='0'
    for sev_suffix in "${snaps_trivial_suffix}" "${snaps_important_suffix}"; do
        if [[ "${#sev_suffix}" -gt "${longest_sev_suffix}" ]]; then
            longest_sev_suffix="${#sev_suffix}"
        fi
    done
    sev_suffix_string="$(head -c "${longest_sev_suffix}" '/dev/zero' | tr '\0' '_')"

    local dataset_name_no_pkgs
    max_dataset_name_length='0'
    for dataset in "${snappable_datasets[@]}"; do
        dataset_name_no_pkgs="${dataset}"'@'"${snap_name_prefix}${snap_field_separator}${date_string}${snap_field_separator}"'op:'"${op_suffix_string}${snap_field_separator}"'sev:'"${sev_suffix_string}${snap_field_separator}"'pkgs:'
        if [[ "${#dataset_name_no_pkgs}" -gt "${max_dataset_name_length}" ]]; then
            max_dataset_name_length="${#dataset_name_no_pkgs}"
        fi
    done

    if [[ "${max_dataset_name_length}" -gt "${max_zfs_snapshot_name_length}" ]]; then
        pprint 'warn' 'Snapshot name would exceed ZFS '"${max_zfs_snapshot_name_length}"' chars limit. Skipping snapshots ...' '0'
    fi
}

function trim_single_remaining_package_name () {
    local pkg_name
    pkg_name="${shorter_pkg_list}"
    case 1 in
        # Trim to 1 to 3 characters, no trailing ellipsis (...)
        $(( 1<=pkgs_list_max_length && pkgs_list_max_length<=3 )))
            pkg_name="${pkg_name::${pkgs_list_max_length}}"
            ;;
        # Show as many pkg name characters as we can while also
        # fitting an ellipsis into the name (...) to indicate
        # that we've cut the pkg name off at the end.
        $(( pkgs_list_max_length>=4 )))
            pkg_name="${pkg_name::$(( pkgs_list_max_length - 3 ))}"'...'
            ;;
    esac
    shorter_pkg_list="${pkg_name}"
}

function trim_pkg_list_oneline () {
    local available_pkg_list_length
    available_pkg_list_length="$((${max_zfs_snapshot_name_length} - ${max_dataset_name_length}))"
    if [[ "${available_pkg_list_length}" -lt "${pkgs_list_max_length}" ]]; then
        # If we have fewer characters available before hitting the
        # ZFS internal maximum snapshot name length than the user
        # wants limit package list length.
        pkgs_list_max_length="${available_pkg_list_length}"
    fi

    local shorter_pkg_list
    if [[ "${pkgs_list_max_length}" -le '0' ]]; then
        # User wants zero characters of pkg names in snapshot name,
        # no need to even find an appropriate pkg name string. Just
        # set to empty string and we're done here.
        shorter_pkg_list=''
    else
        shorter_pkg_list="${unabridged_pkg_list_oneline}"
        while [[ "${#shorter_pkg_list}" -gt "${pkgs_list_max_length}" ]]; do
            shorter_pkg_list="${shorter_pkg_list%${pkg_separator}*}"
            if ! grep -Piq "${pkg_separator}" <<<"${shorter_pkg_list}"; then
                # Only one package remains in package list, no need to continue
                break
            fi
        done
        # If pkg name is still too long trim it. If there's enough
        # space for an ellipsis (...) we add that to indicate we've
        # trimmed the name, otherwise we just take however many
        # characters of the pkg name we can get.
        if [[ "${#shorter_pkg_list}" -gt "${pkgs_list_max_length}" ]]; then
            trim_single_remaining_package_name
        fi
    fi

    trimmed_pkg_list_oneline="${shorter_pkg_list}"
}

function omit_duplicate_snaps () {
    local existing_snaps
    local -a unneeded_snaps
    existing_snaps="$(zfs list -t all -oname -H)"

    for planned_snap in "${planned_snaps[@]}"; do
        if grep -Piq -- '^'"${planned_snap}"'$' <<<"${existing_snaps}"; then
            unneeded_snaps+=("${planned_snap}")
        else
            needed_snaps+=("${planned_snap}")
        fi
    done

    if [[ "${#unneeded_snaps[@]}" -gt '0' ]]; then
        if [[ "${do_dry_run}" == 'true' ]]; then
            pprint 'warn' 'Dry-run, ZFS snapshot skipped (same operation exists at '"${date_string}"'):'
        else
            pprint 'warn' 'ZFS snapshot skipped (same operation exists at '"${date_string}"'):'
        fi
        for unneeded_snap in "${unneeded_snaps[@]}"; do
            pprint 'warn' '  '"${unneeded_snap}"
        done
    fi
}

function do_snaps () {
    local snap_name snap_return_code
    local -a planned_snaps
    for snappable_dataset_id in "${!snappable_datasets[@]}"; do
        snap_name="${snappable_datasets[${snappable_dataset_id}]}"'@'"${snap_name_prefix}${snap_field_separator}${date_string}${snap_field_separator}"'op:'"${conf_op_suffix}${snap_field_separator}"'sev:'"${severity}"
        # If we have at least one pkg name character to append we do
        # so now but if we're not even allowed to append a single
        # character we might as well skip the 'pkgs' field
        # altogether.
        if [[ "${pkgs_list_max_length}" -ge '1' ]]; then
            snap_name="${snap_name}${snap_field_separator}"'pkgs:'"${trimmed_pkg_list_oneline}"
        fi
        planned_snaps["${snappable_dataset_id}"]="${snap_name}"
    done
    local -a needed_snaps
    omit_duplicate_snaps
    if [[ "${#needed_snaps[@]}" -gt '0' ]]; then
        if [[ "${do_dry_run}" == 'true' ]]; then
            pprint 'info' 'Dry-run, pretending to atomically do ZFS snapshot:'
            for needed_snap in "${needed_snaps[@]}"; do
                pprint 'info' '  '"${needed_snap}"
            done
        else
            zfs snapshot "${needed_snaps[@]}"
            snap_return_code="${?}"
            if [[ "${snap_return_code}" -eq '0' ]]; then
                successfully_snapped_datasets=("${snappable_datasets[@]}")
                pprint 'info' 'ZFS snapshot atomically done:'
                for needed_snap in "${needed_snaps[@]}"; do
                    pprint 'info' '  '"${needed_snap}"
                done
            else
                pprint 'warn' 'ZFS snapshot failed:'
                for needed_snap in "${needed_snaps[@]}"; do
                    pprint 'warn' '  '"${needed_snap}"
                done
            fi
        fi
    else
        if [[ "${do_dry_run}" == 'true' ]]; then
            pprint 'warn' 'Dry-run, no ZFS snapshot left to do after accounting for identical operations at '"${date_string}"'.'
        else
            pprint 'warn' 'No ZFS snapshot left to do after accounting for identical operations at '"${date_string}"'.'
        fi
    fi
}

function get_snaps_in_cur_sev () {
    local dataset_to_query
    dataset_to_query="${1:?}"
    snap_list="$(zfs list -H -o 'name' -t snapshot "${dataset_to_query}")"
    snaps_done_by_us="$(grep -Pi -- '@'"${snap_name_prefix}${snap_field_separator}" <<<"${snap_list}")"
    snaps_in_cur_sev="$(grep -Pi -- "${snap_field_separator}"'sev:'"${severity}${snap_field_separator}" <<<"${snaps_done_by_us}")"
    printf -- '%s\n' "${snaps_in_cur_sev}"
}

function do_retention () {
    local snap_list snaps_done_by_us snaps_in_cur_sev snaps_limit oldest_snap snap_return_code
    local -a destroyed_snaps failed_to_destroy_snaps
    if [[ "${do_dry_run}" == 'true' ]]; then
        pprint 'info' 'Dry-run, skipping potential ZFS destroy operations ...'
    else
        for successfully_snapped_dataset in "${successfully_snapped_datasets[@]}"; do
            snaps_in_cur_sev="$(get_snaps_in_cur_sev "${successfully_snapped_dataset}")"
            if [[ "${severity}" == "${snaps_important_suffix}" ]]; then
                snaps_limit="${snaps_important_keep}"
            else
                snaps_limit="${snaps_trivial_keep}"
            fi
            while [[ "$(get_snaps_in_cur_sev "${successfully_snapped_dataset}" | wc -l)" -gt "${snaps_limit}" ]]; do
                oldest_snap="$(get_snaps_in_cur_sev "${successfully_snapped_dataset}" | head -n1)"
                zfs destroy "${oldest_snap}"
                snap_return_code="${?}"
                if [[ "${snap_return_code}" -eq '0' ]]; then
                    destroyed_snaps+=("${oldest_snap}")
                else
                    failed_to_destroy_snaps+=("${oldest_snap}")
                fi
            done
            if [[ "${#destroyed_snaps[@]}" -gt '0' ]]; then
                pprint 'info' 'Oldest ZFS snapshot'"$([[ "${#failed_to_destroy_snaps[@]}" -gt '1' ]] && printf -- '%s' 's')"' in chain '"'"'sev:'"${severity}"''"'"' destroyed:'
                for destroyed_snap in "${destroyed_snaps[@]}"; do
                    pprint 'info' '  '"${destroyed_snap}"
                done
            fi
            if [[ "${#failed_to_destroy_snaps[@]}" -gt '0' ]]; then
                pprint 'warn' 'Failed to prune ZFS snapshot'"$([[ "${#failed_to_destroy_snaps[@]}" -gt '1' ]] && printf -- '%s' 's')"' in chain '"'"'sev:'"${severity}"''"'"':'
                for failed_to_destroy_snap in "${failed_to_destroy_snaps[@]}"; do
                    pprint 'warn' '  '"${failed_to_destroy_snap}"
                done
            fi
        done
    fi
    return 0
}

function main () {
    local pkgs_in_transaction
    pkgs_in_transaction=("${@}")

    local -a important_pkgs_in_transaction trivial_pkgs_in_transaction
    split_pkgs_by_importance "${pkgs_in_transaction[@]}"

    local severity
    set_severity

    local -a globally_snappable_datasets
    get_globally_snappable_datasets

    local -a snappable_datasets
    if [[ "${snap_only_local_datasets}" == 'true' ]]; then
        local local_snappable_datasets
        get_local_snappable_datasets
        trim_globally_snappable_datasets
        if [[ "${#snappable_datasets[@]}" -eq '0' ]]; then
            pprint 'info' 'ZFS snapshot skipped, no local (= currently mounted) dataset has'
            pprint 'info' 'property '"'"''"${zfs_prop}"''"'"' set to '"'"'true'"'"'. At the same'
            pprint 'info' 'time option '"'"'snap_only_local_datasets'"'"' equals '"'"'true'"'"' so'
            pprint 'info' 'we must only snapshot local datasets. Nothing to do here while'
            pprint 'info' 'none of them have '"'"''"${zfs_prop}"''"'"' set to '"'"'true'"'"'.' '0'
        fi
    else
        snappable_datasets=("${globally_snappable_datasets}")
        if [[ "${#snappable_datasets[@]}" -eq '0' ]]; then
            pprint 'info' 'ZFS snapshot skipped, no dataset has property '"'"''"${zfs_prop}"''"'"' set to '"'"'true'"'"'.' '0'
        fi
    fi

    local unabridged_pkg_list_oneline
    write_pkg_list_oneline

    local date_string max_dataset_name_length
    date_string="$($([[ "${snap_timezone}" ]] && printf -- 'export TZ='"${snap_timezone}"); date +"${snap_date_format}")"
    find_max_dataset_name_length

    local trimmed_pkg_list_oneline
    trim_pkg_list_oneline

    local -a successfully_snapped_datasets
    do_snaps
    do_retention
}

main "${pkgs[@]}"