#!/bin/bash
set -euo pipefail

# ─── Color helpers ────────────────────────────────────────────────────────────
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'
BLUE='\033[0;34m'; CYAN='\033[0;36m'; BOLD='\033[1m'; NC='\033[0m'

info()    { echo -e "${CYAN}[INFO]${NC} $*"; }
success() { echo -e "${GREEN}[OK]${NC}   $*"; }
warn()    { echo -e "${YELLOW}[WARN]${NC} $*"; }
error()   { echo -e "${RED}[ERR]${NC}  $*" >&2; }
step()    { echo -e "\n${BOLD}${BLUE}══ $* ${NC}"; }

# ─── Sudo Keepalive ───────────────────────────────────────────────────────────
ensure_sudo() {
    if ! sudo -v 2>/dev/null; then
        error "This script requires sudo access"
        exit 1
    fi
    ( while true; do sudo -n true 2>/dev/null; sleep 50; done ) &
    SUDO_KEEPALIVE_PID=$!
}

# ─── OS Detection ─────────────────────────────────────────────────────────────
detect_os() {
    if [ -f /etc/os-release ]; then
        . /etc/os-release
        OS_ID="${ID:-unknown}"
        OS_ID_LIKE="${ID_LIKE:-}"
    else
        error "Cannot detect OS (no /etc/os-release found)"
        exit 1
    fi

    case "$OS_ID" in
        aosc)   DISTRO="aosc" ;;
        debian) DISTRO="debian" ;;
        ubuntu) DISTRO="ubuntu" ;;
        fedora) DISTRO="fedora" ;;
        *)
            case "$OS_ID_LIKE" in
                *debian*) DISTRO="debian" ;;
                *fedora*|*rhel*) DISTRO="fedora" ;;
                *) error "Unsupported distro: $OS_ID"; exit 1 ;;
            esac
            ;;
    esac
    info "Detected distro: ${BOLD}$DISTRO${NC} (ID=$OS_ID)"
}

# ─── Install Snapper ──────────────────────────────────────────────────────────
install_snapper() {
    step "Installing Snapper"

    if command -v snapper &>/dev/null; then
        success "snapper already installed ($(snapper --version 2>&1 | head -1))"
        return
    fi

    info "Installing snapper ..."
    case "$DISTRO" in
        aosc)           sudo oma install -y snapper ;;
        debian|ubuntu)  sudo apt-get update -qq && sudo apt-get install -y snapper ;;
        fedora)         sudo dnf install -y snapper ;;
    esac
    success "snapper installed"
}

# ─── List Btrfs Subvolumes ────────────────────────────────────────────────────
list_btrfs_mounts() {
    # Returns sorted list of mount points that are btrfs
    # -l disables tree output to avoid box-drawing characters in paths
    findmnt -t btrfs -n -l -o TARGET 2>/dev/null | sort
}

# ─── Select Subvolumes Interactively ─────────────────────────────────────────
select_subvolumes() {
    step "Select Subvolumes to Snapshot"

    local mounts
    mapfile -t mounts < <(list_btrfs_mounts)

    if [ ${#mounts[@]} -eq 0 ]; then
        error "No btrfs subvolumes found mounted on this system"
        exit 1
    fi

    echo "Mounted btrfs filesystems:"
    echo ""
    for i in "${!mounts[@]}"; do
        local subvol device
        subvol=$(findmnt -n -l -o OPTIONS "${mounts[$i]}" 2>/dev/null \
            | tr ',' '\n' | grep '^subvol=' | head -1 | sed 's/subvol=//' || echo "(default)")
        device=$(findmnt -n -l -o SOURCE "${mounts[$i]}" 2>/dev/null || echo "?")
        printf "  %2d)  %-25s  subvol=%-20s  %s\n" \
            "$((i+1))" "${mounts[$i]}" "$subvol" "$device"
    done
    echo ""

    SELECTED_MOUNTS=()
    read -rp "Enter numbers to snapshot (e.g. 1 2 3, or 'all'): " selection

    if [ "$selection" = "all" ]; then
        SELECTED_MOUNTS=("${mounts[@]}")
    else
        for num in $selection; do
            if [[ "$num" =~ ^[0-9]+$ ]] && [ "$num" -ge 1 ] && [ "$num" -le "${#mounts[@]}" ]; then
                SELECTED_MOUNTS+=("${mounts[$((num-1))]}")
            else
                warn "Invalid selection: $num, skipping"
            fi
        done
    fi

    if [ ${#SELECTED_MOUNTS[@]} -eq 0 ]; then
        error "No valid subvolumes selected"
        exit 1
    fi

    info "Selected: ${SELECTED_MOUNTS[*]}"
}

# ─── Derive Config Name from Mount Point ─────────────────────────────────────
config_name_from_mount() {
    local mount="$1"
    if [ "$mount" = "/" ]; then
        echo "root"
    else
        # /home -> home, /var/data -> var-data
        echo "${mount#/}" | tr '/' '-'
    fi
}

# ─── Write Snapper Config ─────────────────────────────────────────────────────
write_snapper_config() {
    local mount="$1"
    local name="$2"
    local config_file="/etc/snapper/configs/$name"

    sudo tee "$config_file" > /dev/null << EOF
# snapper config for server (daily snapshots)

# subvolume to snapshot
SUBVOLUME="$mount"

# filesystem type
FSTYPE="btrfs"

# btrfs qgroup for space aware cleanup algorithms
QGROUP=""

# fraction or absolute size of the filesystem space snapshots may use
SPACE_LIMIT="0.3"

# fraction or absolute size of the filesystem space that should be free
FREE_LIMIT="0.15"

# users and groups allowed to work with config
ALLOW_USERS=""
ALLOW_GROUPS=""

# sync users and groups from ALLOW_USERS and ALLOW_GROUPS to .snapshots directory
SYNC_ACL="no"

# start comparing pre- and post-snapshot in background after creating post-snapshot
BACKGROUND_COMPARISON="yes"

# run daily number cleanup
NUMBER_CLEANUP="yes"

# minimum age in seconds before a snapshot may be deleted by cleanup
NUMBER_MIN_AGE="3600"

# keep at most this many numbered snapshots
NUMBER_LIMIT="20"

# keep at most this many important numbered snapshots
NUMBER_LIMIT_IMPORTANT="5"

# timeline snapshots (automatic snapshots)
TIMELINE_CREATE="yes"

# cleanup timeline snapshots after some time
TIMELINE_CLEANUP="yes"

# minimum age in seconds before timeline snapshots may be deleted
TIMELINE_MIN_AGE="3600"

# no hourly snapshots retained
TIMELINE_LIMIT_HOURLY="0"

# keep 14 daily snapshots (last 2 weeks)
TIMELINE_LIMIT_DAILY="14"

# keep 4 weekly snapshots (last month)
TIMELINE_LIMIT_WEEKLY="4"

# keep 6 monthly snapshots (last 6 months)
TIMELINE_LIMIT_MONTHLY="6"

# no quarterly snapshots
TIMELINE_LIMIT_QUARTERLY="0"

# keep 2 yearly snapshots
TIMELINE_LIMIT_YEARLY="2"

# cleanup empty pre-post pairs
EMPTY_PRE_POST_CLEANUP="yes"

# minimum age in seconds before empty pre-post pairs may be deleted
EMPTY_PRE_POST_MIN_AGE="3600"
EOF
}

# ─── Apply Snapper Config for a Mount Point ───────────────────────────────────
apply_snapper_config() {
    local mount="$1"
    local name
    name=$(config_name_from_mount "$mount")

    step "Configuring snapper for $mount (config: $name)"

    # Check if config already exists
    if sudo snapper list-configs 2>/dev/null | awk 'NR>1 {print $1}' | grep -qw "$name"; then
        warn "Config '$name' already exists, overwriting policy settings ..."
        write_snapper_config "$mount" "$name"
    else
        info "Creating snapper config '$name' for $mount ..."
        sudo snapper -c "$name" create-config "$mount"
        success "Config '$name' created"
        info "Applying snapshot policy ..."
        write_snapper_config "$mount" "$name"
    fi

    success "Policy applied to config '$name'"
}

# ─── Enable Snapper Systemd Timers ───────────────────────────────────────────
enable_timers() {
    step "Enabling Snapper Timers"

    local timers=("snapper-timeline.timer" "snapper-cleanup.timer")
    for timer in "${timers[@]}"; do
        if systemctl list-unit-files "$timer" 2>/dev/null | grep -q "$timer"; then
            sudo systemctl enable --now "$timer"
            success "$timer enabled and started"
        else
            warn "$timer not found — snapper timers may need manual setup on this distro"
        fi
    done
}

# ─── Snapshot Test ────────────────────────────────────────────────────────────
test_snapshot() {
    local mount="$1"
    local name
    name=$(config_name_from_mount "$mount")

    step "Snapshot Test — $mount (config: $name)"

    info "Creating test snapshot ..."
    local snap_num
    snap_num=$(sudo snapper -c "$name" create --description "snapper-setup-test" --print-number)
    success "Snapshot #$snap_num created"

    info "Current snapshots for '$name':"
    sudo snapper -c "$name" list

    info "Deleting test snapshot #$snap_num ..."
    sudo snapper -c "$name" delete "$snap_num"
    success "Test snapshot deleted — snapshot test passed"
}

# ─── Main ─────────────────────────────────────────────────────────────────────
main() {
    echo -e "${BOLD}${CYAN}"
    echo "╔══════════════════════════════════════════╗"
    echo "║        Snapper Auto Configuration        ║"
    echo "╚══════════════════════════════════════════╝"
    echo -e "${NC}"

    detect_os
    ensure_sudo
    install_snapper
    select_subvolumes

    for mount in "${SELECTED_MOUNTS[@]}"; do
        apply_snapper_config "$mount"
    done

    enable_timers

    for mount in "${SELECTED_MOUNTS[@]}"; do
        test_snapshot "$mount"
    done

    # Stop sudo keepalive
    if [ -n "${SUDO_KEEPALIVE_PID:-}" ]; then
        kill "$SUDO_KEEPALIVE_PID" 2>/dev/null || true
    fi

    echo ""
    echo -e "${BOLD}${GREEN}══ Snapper configuration complete! ══${NC}"
    info "Timeline snapshots will run automatically via systemd timers"
    info "List all configs:  sudo snapper list-configs"
    info "View snapshots:    sudo snapper -c <config> list"
    info "Manual snapshot:   sudo snapper -c <config> create --description 'my-snapshot'"
}

main "$@"