#!/bin/bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "$0")" && pwd)"
BACKUP_DIR="$SCRIPT_DIR/pve-backups"
SSH_KEY=""
REMOTE_HOST=""
REMOTE_USER="root"
SUDO_CMD=""
HAS_RSYNC=0
SCRIPT_NAME=$(basename "$0")

usage() {
    cat <<EOF
Usage: $SCRIPT_NAME <command> [options]

Commands:
  backup        Create a backup of Proxmox configuration
  restore       Restore Proxmox configuration from a backup
  verify        Verify a backup archive's integrity (SHA256)
  list          List available backups

Backup options:
  -d, --dir DIR         Local backup directory (default: $SCRIPT_DIR/pve-backups)
  -h, --host HOST       Remote Proxmox host (SSH)
  -u, --user USER       SSH user (default: root)
  -i, --identity FILE   SSH identity file

Restore / Verify options:
  -d, --dir DIR         Directory containing backup archive
  -f, --file FILE       Specific backup archive to restore/verify
  -y, --yes             Skip confirmation prompts (DANGEROUS)

List options:
  -d, --dir DIR         Directory to list backups from (default: $BACKUP_DIR)

Examples:
  $SCRIPT_NAME backup
  $SCRIPT_NAME backup -d /mnt/backup
  $SCRIPT_NAME backup -h 192.168.1.100 -i ~/.ssh/id_rsa
  $SCRIPT_NAME restore -f /mnt/backup/pve-backup-20250101_120000.tar.gz
  $SCRIPT_NAME verify -f /mnt/backup/pve-backup-20250101_120000.tar.gz
  $SCRIPT_NAME list
EOF
    exit 1
}

log() { echo "[*] $*"; }
warn() { echo "[!] $*" >&2; }
die() { echo "[FATAL] $*" >&2; exit 1; }

# Parse common args (consume them, return remaining via ${PARSE_REMAINING[@]})
parse_common_args() {
    local parsed=()
    PARSE_REMAINING=()
    while [[ $# -gt 0 ]]; do
        case "$1" in
            -d|--dir) BACKUP_DIR="$2"; shift 2 ;;
            -h|--host) REMOTE_HOST="$2"; shift 2 ;;
            -u|--user) REMOTE_USER="$2"; shift 2 ;;
            -i|--identity) SSH_KEY="$2"; shift 2 ;;
            -y|--yes) FORCE_YES=1; shift ;;
            -f|--file) RESTORE_FILE="$2"; shift 2 ;;
            --) shift; parsed+=("$@"); break ;;
            -*)
                PARSE_REMAINING=("$@")
                return
                ;;
            *) parsed+=("$1"); shift ;;
        esac
    done
    PARSE_REMAINING=("${parsed[@]}")
}

build_ssh_cmd() {
    local cmd=()
    if [[ -n "$REMOTE_HOST" ]]; then
        cmd=(ssh)
        [[ -n "$SSH_KEY" ]] && cmd+=(-i "$SSH_KEY")
        cmd+=("${REMOTE_USER}@${REMOTE_HOST}")
    fi
    echo "${cmd[@]}"
}

build_ssh_rsync_cmd() {
    local cmd="ssh"
    [[ -n "$SSH_KEY" ]] && cmd+=" -i $SSH_KEY"
    echo "$cmd"
}

run_remote() {
    local ssh_cmd
    ssh_cmd=$(build_ssh_cmd)
    if [[ -n "$ssh_cmd" ]]; then
        $ssh_cmd $SUDO_CMD "$@"
    else
        ${SUDO_CMD}bash -c "$*"
    fi
}

run_remote_raw() {
    local ssh_cmd
    ssh_cmd=$(build_ssh_cmd)
    if [[ -n "$ssh_cmd" ]]; then
        $ssh_cmd "${SUDO_CMD}bash -s" <<< "$1"
    else
        ${SUDO_CMD}bash -c "$1"
    fi
}

check_proxmox() {
    if [[ -z "$REMOTE_HOST" ]]; then
        if [[ "$(id -u)" -ne 0 ]]; then
            SUDO_CMD="sudo -n "
        fi
    elif [[ "$REMOTE_USER" != "root" ]]; then
        SUDO_CMD="sudo -n "
    fi

    if ! run_remote "test -f /etc/pve/datacenter.cfg" 2>/dev/null; then
        local where="${REMOTE_HOST:-locally}"
        local who="${REMOTE_USER:-$(whoami)}"
        if [[ -n "$SUDO_CMD" ]]; then
            die "Cannot read /etc/pve $where — user '$who' needs passwordless sudo"
        else
            die "Cannot read /etc/pve $where — this script must be run as root"
        fi
    fi
}

check_rsync() {
    HAS_RSYNC=0
    if [[ -z "$REMOTE_HOST" ]]; then
        return
    fi
    if run_remote "command -v rsync" &>/dev/null; then
        HAS_RSYNC=1
        log "rsync available on remote — using for incremental transfers"
    else
        log "rsync not found on remote — falling back to tar/cat transfer"
    fi
}

# rsync wrapper: copies remote path(s) into dest dir.
# Handles SSH, sudo (via --rsync-path), and both file and directory sources.
# Prefix src with 'dir:' to treat it as a directory (--delete).
rsync_pull() {
    local src="$1"
    local dest="$2"
    local is_dir=0

    if [[ "$HAS_RSYNC" -ne 1 || -z "$REMOTE_HOST" ]]; then
        return 1
    fi

    if [[ "$src" == dir:* ]]; then
        is_dir=1
        src="${src:4}"
    fi

    local rsync_ssh
    rsync_ssh=$(build_ssh_rsync_cmd)

    local rsync_opts=(-a)
    [[ "$is_dir" -eq 1 ]] && rsync_opts+=(--delete)
    [[ -n "$SUDO_CMD" ]] && rsync_opts+=(--rsync-path="${SUDO_CMD}rsync")

    mkdir -p "$dest"
    rsync "${rsync_opts[@]}" -e "$rsync_ssh" \
        "${REMOTE_USER}@${REMOTE_HOST}:$src" "$dest" &>/dev/null
}

collect_pve_configs() {
    local temp_dir="$1"
    local dest="$temp_dir/etc"

    mkdir -p "$dest/pve" "$dest/apt" "$dest/system"

    # ---- /etc/pve ----
    log "Backing up /etc/pve configuration..."
    if rsync_pull "dir:/etc/pve/" "$dest/pve/" && [[ -f "$dest/pve/datacenter.cfg" ]]; then
        # Clean up rsync'd noise
        rm -rf "$dest/pve/.rrd" "$dest/pve/nodes/"*/qemu-server/*.lock \
               "$dest/pve/nodes/"*/lxc/*.lock "$dest/pve/priv/known_hosts" 2>/dev/null || true
    else
        # Fallback: tar-based snapshot via SSH
        rm -rf "$dest/pve" 2>/dev/null; mkdir "$dest/pve"
        run_remote_raw '
        set -euo pipefail
        tmp=$(mktemp -d)
        mkdir "$tmp/pve"
        cd /etc/pve
        tar cf - \
            --exclude="./nodes/*/qemu-server/*.lock" \
            --exclude="./nodes/*/lxc/*.lock" \
            --exclude="./priv/known_hosts" \
            --exclude=".rrd" \
            . 2>/dev/null | tar xf - -C "$tmp/pve"
        tar czf - -C "$tmp" pve
        rm -rf "$tmp"
        ' > "$dest/pve/pve-config.tar.gz" 2>/dev/null
    fi

    # ---- Network & system files ----
    log "Backing up network and system configs..."
    local sysfiles=(
        /etc/network/interfaces /etc/hostname /etc/hosts
        /etc/resolv.conf /etc/mailname /etc/fstab
    )
    local rsync_ok=1
    for f in "${sysfiles[@]}"; do
        if ! rsync_pull "$f" "$dest/system/"; then
            rsync_ok=0
            break
        fi
    done
    if [[ "$rsync_ok" -eq 0 ]]; then
        for f in "${sysfiles[@]}"; do
            run_remote "cat $f 2>/dev/null" > "$dest/system/$(basename "$f")" 2>/dev/null || true
        done
    fi

    # ---- APT sources ----
    log "Backing up APT sources..."
    if ! rsync_pull "dir:/etc/apt/sources.list.d/" "$dest/apt/" || \
       ! rsync_pull "/etc/apt/sources.list" "$dest/apt/" || \
       ! rsync_pull "/etc/apt/auth.conf" "$dest/apt/"; then
        # Fallback
        rm -rf "$dest/apt" 2>/dev/null; mkdir "$dest/apt"
        run_remote_raw '
        set -euo pipefail
        tmp=$(mktemp -d)
        mkdir "$tmp/apt"
        cp -a /etc/apt/sources.list "$tmp/apt/" 2>/dev/null || true
        cp -a /etc/apt/sources.list.d "$tmp/apt/" 2>/dev/null || true
        cp -a /etc/apt/auth.conf "$tmp/apt/" 2>/dev/null || true
        tar czf - -C "$tmp" apt
        rm -rf "$tmp"
        ' > "$dest/apt/apt-sources.tar.gz" 2>/dev/null
    fi

    # ---- Corosync ----
    log "Backing up corosync config..."
    if ! rsync_pull "/etc/corosync/corosync.conf" "$dest/system/"; then
        run_remote "cat /etc/corosync/corosync.conf 2>/dev/null" > "$dest/system/corosync.conf" 2>/dev/null || true
    fi

    # ---- Ceph ----
    log "Backing up Ceph config..."
    if ! rsync_pull "/etc/ceph/ceph.conf" "$dest/system/"; then
        run_remote "cat /etc/ceph/ceph.conf 2>/dev/null" > "$dest/system/ceph.conf" 2>/dev/null || true
    fi

    # ---- sysctl / cron ----
    log "Backing up sysctl and cron..."
    if ! rsync_pull "dir:/etc/sysctl.d/" "$dest/system/sysctl.d/" || \
       ! rsync_pull "dir:/etc/cron.d/" "$dest/system/cron.d/" || \
       ! rsync_pull "dir:/etc/cron.daily/" "$dest/system/cron.daily/"; then
        # Fallback
        rm -rf "$dest/system/sysctl.d" "$dest/system/cron.d" "$dest/system/cron.daily" 2>/dev/null
        run_remote_raw '
        tmp=$(mktemp -d)
        mkdir "$tmp/etc"
        cp -a /etc/sysctl.d "$tmp/etc/" 2>/dev/null || true
        cp -a /etc/cron.d "$tmp/etc/" 2>/dev/null || true
        cp -a /etc/cron.daily "$tmp/etc/" 2>/dev/null || true
        tar czf - -C "$tmp" etc
        rm -rf "$tmp"
        ' > "$dest/system/sysconfig.tar.gz" 2>/dev/null
    fi

    # ---- Package manifest (always via command) ----
    log "Recording package manifest..."
    run_remote "dpkg-query -W 2>/dev/null || true" > "$dest/system/packages.txt" 2>/dev/null || true
    run_remote "apt-mark showmanual 2>/dev/null || true" > "$dest/system/packages-manual.txt" 2>/dev/null || true
    run_remote "pveversion 2>/dev/null || true" > "$dest/system/pveversion.txt" 2>/dev/null || true

    # ---- Firewall rules (always via command) ----
    log "Backing up firewall rules..."
    run_remote "iptables-save 2>/dev/null || true" > "$dest/system/iptables.rules" 2>/dev/null || true
    run_remote "ip6tables-save 2>/dev/null || true" > "$dest/system/ip6tables.rules" 2>/dev/null || true

    log "Config collection complete."
}

do_backup() {
    check_proxmox
    check_rsync

    local stamp=$(run_remote "date +%Y%m%d_%H%M%S 2>/dev/null || date +%Y%m%d_%H%M%S")
    local hostname=$(run_remote "hostname -s 2>/dev/null || echo unknown")

    local temp_dir=$(mktemp -d)
    local archive_name="pve-backup-${hostname}-${stamp}.tar.gz"
    local archive_path="$BACKUP_DIR/$archive_name"

    mkdir -p "$BACKUP_DIR" "$temp_dir"

    collect_pve_configs "$temp_dir"

    # Add metadata
    cat > "$temp_dir/META" <<EOM
hostname: $hostname
created:  $(date -u +%Y-%m-%dT%H:%M:%SZ)
pveversion: $(run_remote "pveversion 2>/dev/null || echo unknown")
pve_nodes: $(run_remote 'ls /etc/pve/nodes/ 2>/dev/null | tr "\n" " "' 2>/dev/null)
EOM

    tar czf "$archive_path" -C "$temp_dir" $(ls -A "$temp_dir")
    rm -rf "$temp_dir"

    log "Backup written to: $archive_path"
    echo "  Size: $(du -h "$archive_path" | cut -f1)"

    # Generate and store SHA256 checksum
    (cd "$BACKUP_DIR" && sha256sum "$archive_name" > "$archive_path.sha256")
    log "Checksum: $(cat "$archive_path.sha256")"
}

# Remote rollback + restore of /etc/pve via rsync or tar pipe.
# Handles both rsync'd directories and tarball archives.
_restore_pve() {
    local src_dir="$1"
    local ssh_cmd
    ssh_cmd=$(build_ssh_cmd)

    run_remote "tar czf /etc/pve-rollback.tar.gz -C /etc/pve . 2>/dev/null" 2>/dev/null || true

    if [[ -f "$src_dir/pve-config.tar.gz" ]]; then
        log "  Restoring /etc/pve (from tarball)..."
        run_remote "tar xzf - -C /etc/pve" < "$src_dir/pve-config.tar.gz" 2>/dev/null || \
            warn "Failed to restore some /etc/pve files"
    elif [[ -f "$src_dir/datacenter.cfg" ]]; then
        log "  Restoring /etc/pve (from directory)..."
        local tried_rsync=0
        if [[ -n "$ssh_cmd" ]] && command -v rsync &>/dev/null && run_remote "command -v rsync" &>/dev/null; then
            local rsync_ssh
            rsync_ssh=$(build_ssh_rsync_cmd)
            local rsync_opts=(-a --delete)
            [[ -n "$SUDO_CMD" ]] && rsync_opts+=(--rsync-path="${SUDO_CMD}rsync")
            rsync "${rsync_opts[@]}" -e "$rsync_ssh" "$src_dir/" "${REMOTE_USER}@${REMOTE_HOST}:/etc/pve/" &>/dev/null && \
                tried_rsync=1
        fi
        if [[ "$tried_rsync" -eq 0 ]]; then
            tar czf - -C "$src_dir" . | run_remote "tar xzf - -C /etc/pve" 2>/dev/null || \
                warn "Failed to restore /etc/pve files"
        fi
    fi
}

_restore_dir() {
    local src_dir="$1" remote_path="$2" label="$3"
    if [[ ! -d "$src_dir" ]] || ! ls -A "$src_dir" &>/dev/null; then
        return
    fi
    log "  Restoring $label..."
    local ssh_cmd
    ssh_cmd=$(build_ssh_cmd)
    local tried_rsync=0
    if [[ -n "$ssh_cmd" ]] && command -v rsync &>/dev/null && run_remote "command -v rsync" &>/dev/null; then
        local rsync_ssh
        rsync_ssh=$(build_ssh_rsync_cmd)
        local rsync_opts=(-a)
        [[ -n "$SUDO_CMD" ]] && rsync_opts+=(--rsync-path="${SUDO_CMD}rsync")
        rsync "${rsync_opts[@]}" -e "$rsync_ssh" "$src_dir/" "${REMOTE_USER}@${REMOTE_HOST}:$remote_path" &>/dev/null && \
            tried_rsync=1
    fi
    if [[ "$tried_rsync" -eq 0 ]]; then
        tar czf - -C "$src_dir" . | run_remote "tar xzf - -C $remote_path" 2>/dev/null || \
            warn "Failed to restore $label"
    fi
}

do_restore() {
    check_proxmox

    local archive_file=""

    if [[ -n "${RESTORE_FILE:-}" ]]; then
        archive_file="$RESTORE_FILE"
    else
        local latest
        latest=$(ls -t "$BACKUP_DIR"/pve-backup-*.tar.gz 2>/dev/null | head -1) || true
        if [[ -z "$latest" ]]; then
            die "No backup archives found in $BACKUP_DIR and no --file specified"
        fi
        archive_file="$latest"
    fi

    if [[ ! -f "$archive_file" ]]; then
        die "Backup file not found: $archive_file"
    fi

    log "Restoring from: $archive_file"

    if [[ "${FORCE_YES:-0}" -ne 1 ]]; then
        echo ""
        warn "WARNING: This will overwrite Proxmox configuration files!"
        warn "Your current configs will be backed up to /var/tmp/pve-restore-rollback/"
        echo ""
        read -r -p "Are you sure you want to proceed? [y/N] " reply
        if [[ ! "$reply" =~ ^[Yy]$ ]]; then
            die "Restore cancelled."
        fi
    fi

    local rollback_dir="/var/tmp/pve-restore-rollback-$(date +%Y%m%d_%H%M%S)"
    log "Creating rollback snapshot at $rollback_dir..."
    mkdir -p "$rollback_dir"

    local temp_dir
    temp_dir=$(mktemp -d)
    tar xzf "$archive_file" -C "$temp_dir"

    # Restore /etc/pve
    _restore_pve "$temp_dir/etc/pve"

    # Restore individual system files
    for f in "$temp_dir"/etc/system/*; do
        local basename_f
        basename_f=$(basename "$f")
        [[ -f "$f" ]] || continue
        local target=""
        case "$basename_f" in
            interfaces) target="/etc/network/interfaces" ;;
            hostname) target="/etc/hostname" ;;
            hosts) target="/etc/hosts" ;;
            resolv.conf) target="/etc/resolv.conf" ;;
            fstab) target="/etc/fstab" ;;
            corosync.conf) target="/etc/corosync/corosync.conf" ;;
            ceph.conf) target="/etc/ceph/ceph.conf" ;;
            mailname) target="/etc/mailname" ;;
            *) continue ;;
        esac
        if [[ -s "$f" ]]; then
            run_remote "mkdir -p $(dirname "$target") && cat > $target" < "$f" 2>/dev/null || warn "Failed to restore $target"
        fi
    done

    # Restore APT sources
    if [[ -f "$temp_dir/etc/apt/apt-sources.tar.gz" ]]; then
        log "Restoring APT sources..."
        run_remote "tar xzf - -C /" < "$temp_dir/etc/apt/apt-sources.tar.gz" 2>/dev/null || warn "Failed to restore APT sources"
    elif [[ -d "$temp_dir/etc/apt" ]]; then
        _restore_dir "$temp_dir/etc/apt" "/etc/apt" "APT sources"
    fi

    # Restore sysconfig
    if [[ -f "$temp_dir/etc/system/sysconfig.tar.gz" ]]; then
        log "Restoring sysctl/cron configs..."
        run_remote "tar xzf - -C /" < "$temp_dir/etc/system/sysconfig.tar.gz" 2>/dev/null || warn "Failed to restore sysconfig"
    elif [[ -d "$temp_dir/etc/system/sysctl.d" ]]; then
        _restore_dir "$temp_dir/etc/system/sysctl.d" "/etc/sysctl.d" "sysctl"
        _restore_dir "$temp_dir/etc/system/cron.d" "/etc/cron.d" "cron"
        _restore_dir "$temp_dir/etc/system/cron.daily" "/etc/cron.daily" "cron.daily"
    fi

    rm -rf "$temp_dir"

    log "Restore complete!"
    log "Rollback snapshot: $rollback_dir"
    echo ""
    warn "NOTE: Some changes (e.g., network, hostname) require a reboot to take full effect."
    warn "NOTE: Restored /etc/pve files sync via pmxcfs — restart pveproxy if web UI acts odd."
}

do_verify() {
    local archive_file=""

    if [[ -n "${RESTORE_FILE:-}" ]]; then
        archive_file="$RESTORE_FILE"
    else
        die "Specify an archive with -f/--file"
    fi

    if [[ ! -f "$archive_file" ]]; then
        die "File not found: $archive_file"
    fi

    local checksum_file="${archive_file}.sha256"

    log "Verifying: $archive_file"

    # Check tar integrity
    if tar tzf "$archive_file" &>/dev/null; then
        log "  Archive integrity: OK (valid tar.gz)"
    else
        die "  Archive integrity: FAILED — corrupt or truncated"
    fi

    # Check SHA256 if available
    local archive_dir archive_basename
    archive_dir=$(dirname "$archive_file")
    archive_basename=$(basename "$archive_file")
    if [[ -f "$checksum_file" ]]; then
        if cd "$archive_dir" && sha256sum -c "$archive_basename.sha256" &>/dev/null; then
            cd "$OLDPWD"
            log "  SHA256 checksum:  OK (matches stored checksum)"
        else
            die "  SHA256 checksum:  FAILED — archive modified or corrupted"
        fi
    else
        # Generate and save checksum if none exists
        warn "  SHA256 checksum:  none found — generating now"
        (cd "$archive_dir" && sha256sum "$archive_basename" > "$archive_basename.sha256")
        log "  Checksum saved to: $checksum_file"
    fi
}

do_list() {
    ls -lh "$BACKUP_DIR"/pve-backup-*.tar.gz 2>/dev/null || echo "No backups found in $BACKUP_DIR"
}

# ---- Main ----
COMMAND="${1:-}"
[[ -z "$COMMAND" ]] && usage
shift

FORCE_YES=0
RESTORE_FILE=""

parse_common_args "$@"
set -- "${PARSE_REMAINING[@]}"

case "$COMMAND" in
    backup)
        do_backup
        ;;
    restore)
        do_restore
        ;;
    verify)
        do_verify
        ;;
    list)
        do_list
        ;;
    *)
        usage
        ;;
esac