#!/usr/bin/env bash
set -u
set -o pipefail

VERSION="2.0.0"

COLOR_RESET="\e[0m"
COLOR_ORANGE="\e[38;5;208m"
COLOR_GREEN="\e[32m"
COLOR_RED="\e[31m"
COLOR_CYAN="\e[36m"
COLOR_BOLD="\e[1m"

DEFAULT_ACTION="backup"
TMP_RESTORE_DIR="/tmp/pve-restore.$$"
SCRIPT_NAME="$(basename "$0")"

ACTION=""
TARGET_FILE=""
BACKUP_GUESTS="yes"
GUEST_BACKUP_TARGET=""
GUEST_BACKUP_MODE="snapshot"
GUEST_BACKUP_COMPRESS="zstd"

BACKUP_ITEMS=(
    "/etc/pve"
    "/etc/network/interfaces"
    "/etc/hosts"
    "/etc/hostname"
    "/etc/resolv.conf"
    "/etc/vzdump.conf"
    "/etc/modules"
    "/etc/modules-load.d"
    "/etc/modprobe.d"
    "/etc/sysctl.conf"
    "/etc/sysctl.d"
    "/etc/cron.d"
    "/etc/cron.daily"
    "/etc/cron.weekly"
    "/etc/cron.monthly"
    "/etc/aliases"
    "/root/.ssh"
)

print_banner() {
    clear
    echo -e "${COLOR_CYAN}${COLOR_BOLD}"
    echo "=============================================================="
    echo "               Backup - Restore Proxmox VE"
    echo "=============================================================="
    echo -e "${COLOR_RESET}"
}

msg_process() {
    echo -e "${COLOR_ORANGE}[PROCESS]${COLOR_RESET} $1"
}

msg_success() {
    echo -e "${COLOR_GREEN}[SUCCESS]${COLOR_RESET} $1"
}

msg_failed() {
    echo -e "${COLOR_RED}[FAILED]${COLOR_RESET} $1" >&2
}

msg_info() {
    echo -e "${COLOR_CYAN}[INFO]${COLOR_RESET} $1"
}

die() {
    msg_failed "$1"
    cleanup
    exit 1
}

cleanup() {
    if [[ -d "$TMP_RESTORE_DIR" ]]; then
        rm -rf "$TMP_RESTORE_DIR"
    fi
}

run_cmd() {
    local description="$1"
    shift

    msg_process "$description"
    if "$@"; then
        msg_success "$description"
        return 0
    else
        local rc=$?
        msg_failed "$description (exit code: $rc)"
        return $rc
    fi
}

run_cmd_quiet() {
    local description="$1"
    shift

    msg_process "$description"
    if "$@" >/dev/null 2>&1; then
        msg_success "$description"
        return 0
    else
        local rc=$?
        msg_failed "$description (exit code: $rc)"
        return $rc
    fi
}

require_root() {
    if [[ "${EUID}" -ne 0 ]]; then
        die "This script must be run as root."
    fi
    msg_success "Running as root"
}

require_commands() {
    local required_cmds=("tar" "hostname" "date" "systemctl" "cp" "mkdir" "rm" "grep" "awk" "pvesm")
    local missing=()

    msg_process "Checking required commands"
    for cmd in "${required_cmds[@]}"; do
        if ! command -v "$cmd" >/dev/null 2>&1; then
            missing+=("$cmd")
        fi
    done

    if [[ "$ACTION" == "backup" ]]; then
        if ! command -v vzdump >/dev/null 2>&1; then
            missing+=("vzdump")
        fi
    fi

    if [[ ${#missing[@]} -gt 0 ]]; then
        msg_failed "Missing required commands: ${missing[*]}"
        exit 1
    fi

    msg_success "All required commands are available"
}

detect_proxmox() {
    msg_process "Checking if this host is Proxmox VE"
    if [[ -f /etc/pve/.version ]] || command -v pveversion >/dev/null 2>&1; then
        msg_success "Proxmox VE host detected"
    else
        die "This does not appear to be a Proxmox VE host."
    fi
}

validate_backup_filename() {
    local filename="$1"

    [[ -z "$filename" ]] && return 1
    [[ "$filename" =~ [[:space:]] ]] && return 1
    [[ "$filename" == */ ]] && return 1

    case "$filename" in
        *.tar.gz|*.tgz) return 0 ;;
        *) return 1 ;;
    esac
}

prompt_action() {
    local choice
    echo
    read -r -p "Choose action [Backup/Restore] (default: Backup): " choice
    choice="${choice,,}"

    if [[ -z "$choice" ]]; then
        ACTION="$DEFAULT_ACTION"
    elif [[ "$choice" == "backup" || "$choice" == "b" ]]; then
        ACTION="backup"
    elif [[ "$choice" == "restore" || "$choice" == "r" ]]; then
        ACTION="restore"
    else
        msg_failed "Invalid choice. Defaulting to Backup."
        ACTION="$DEFAULT_ACTION"
    fi

    msg_info "Selected action: ${ACTION^}"
}

prompt_filename() {
    local prompt_text="$1"
    local filename

    while true; do
        echo
        read -r -p "$prompt_text" filename
        if validate_backup_filename "$filename"; then
            TARGET_FILE="$filename"
            msg_success "Filename accepted: $TARGET_FILE"
            break
        else
            msg_failed "Invalid filename. Use a filename ending with .tar.gz or .tgz and no spaces."
        fi
    done
}

prompt_backup_guests() {
    local answer
    echo
    read -r -p "Do you also want to backup all VMs and LXCs? [Y/n]: " answer
    answer="${answer,,}"

    if [[ -z "$answer" || "$answer" == "y" || "$answer" == "yes" ]]; then
        BACKUP_GUESTS="yes"
        msg_info "Guest backup selected: Yes"
    elif [[ "$answer" == "n" || "$answer" == "no" ]]; then
        BACKUP_GUESTS="no"
        msg_info "Guest backup selected: No"
    else
        msg_failed "Invalid choice. Defaulting to Yes."
        BACKUP_GUESTS="yes"
        msg_info "Guest backup selected: Yes"
    fi
}

list_vzdump_targets() {
    msg_info "Available Proxmox storages that support backup content:"
    pvesm status -content backup 2>/dev/null | awk 'NR==1 || NF {print}'
}

prompt_guest_backup_target() {
    local target
    echo
    list_vzdump_targets
    echo
    read -r -p "Enter Proxmox backup storage name for VMs/LXCs (example: local, backup-nas, pbs-store): " target

    if [[ -z "$target" ]]; then
        die "Guest backup storage name cannot be empty."
    fi

    if ! pvesm status -storage "$target" >/dev/null 2>&1; then
        die "Storage '$target' was not found in Proxmox storage configuration."
    fi

    if ! pvesm status -content backup 2>/dev/null | awk 'NR>1 {print $1}' | grep -Fxq "$target"; then
        die "Storage '$target' does not appear to support VZDump backup content."
    fi

    GUEST_BACKUP_TARGET="$target"
    msg_success "Guest backup storage selected: $GUEST_BACKUP_TARGET"
}

show_backup_content() {
    msg_info "The following host paths will be processed if they exist:"
    local item
    for item in "${BACKUP_ITEMS[@]}"; do
        echo "  - $item"
    done
}

build_existing_items_list() {
    EXISTING_ITEMS=()
    local item
    for item in "${BACKUP_ITEMS[@]}"; do
        if [[ -e "$item" ]]; then
            EXISTING_ITEMS+=("$item")
        else
            msg_info "Skipping missing path: $item"
        fi
    done

    if [[ ${#EXISTING_ITEMS[@]} -eq 0 ]]; then
        die "No backup items were found on this host."
    fi
}

backup_host_config() {
    show_backup_content
    build_existing_items_list

    if [[ -e "$TARGET_FILE" ]]; then
        echo
        read -r -p "File '$TARGET_FILE' already exists. Overwrite? [y/N]: " overwrite
        overwrite="${overwrite,,}"
        if [[ "$overwrite" != "y" && "$overwrite" != "yes" ]]; then
            die "Backup aborted by user."
        fi
    fi

    msg_info "Host configuration backup file: $TARGET_FILE"

    run_cmd "Creating Proxmox VE host configuration backup archive" \
        tar -czpf "$TARGET_FILE" --absolute-names "${EXISTING_ITEMS[@]}" \
        || die "Host configuration backup failed."

    run_cmd "Verifying host configuration backup archive integrity" \
        tar -tzf "$TARGET_FILE" >/dev/null \
        || die "Host configuration archive verification failed."

    run_cmd "Displaying host configuration backup archive details" \
        ls -lh "$TARGET_FILE" \
        || die "Host configuration backup file was created but details could not be displayed."

    msg_success "Host configuration backup completed successfully: $TARGET_FILE"
}

backup_guests() {
    if [[ "$BACKUP_GUESTS" != "yes" ]]; then
        msg_info "Skipping VM/LXC backup because user selected No"
        return 0
    fi

    prompt_guest_backup_target

    msg_info "Guest backup target storage: $GUEST_BACKUP_TARGET"
    msg_info "Guest backup mode: $GUEST_BACKUP_MODE"
    msg_info "Guest backup compression: $GUEST_BACKUP_COMPRESS"

    run_cmd "Backing up all VMs and LXCs with vzdump" \
        vzdump --all 1 --storage "$GUEST_BACKUP_TARGET" --mode "$GUEST_BACKUP_MODE" --compress "$GUEST_BACKUP_COMPRESS" \
        || die "VM/LXC backup failed."

    msg_success "All selected VM/LXC backups completed successfully"
}

pre_restore_checks() {
    [[ -f "$TARGET_FILE" ]] || die "Restore file does not exist: $TARGET_FILE"

    run_cmd "Verifying restore archive is readable" \
        tar -tzf "$TARGET_FILE" >/dev/null \
        || die "Restore archive is invalid or corrupted."

    run_cmd "Creating temporary restore directory" \
        mkdir -p "$TMP_RESTORE_DIR" \
        || die "Could not create temporary restore directory."

    run_cmd "Extracting restore archive to temporary directory" \
        tar -xzpf "$TARGET_FILE" -C "$TMP_RESTORE_DIR" \
        || die "Could not extract restore archive."

    if [[ ! -d "$TMP_RESTORE_DIR/etc" ]]; then
        die "Restore archive does not contain expected /etc content."
    fi

    msg_success "Restore archive passed validation"
}

restore_item_if_exists() {
    local relative_path="$1"
    local src="$TMP_RESTORE_DIR/$relative_path"
    local dst="/$relative_path"

    if [[ -e "$src" ]]; then
        msg_process "Restoring /$relative_path"
        mkdir -p "$(dirname "$dst")" || { msg_failed "Failed creating parent directory for /$relative_path"; return 1; }
        cp -a "$src" "$dst" || { msg_failed "Failed restoring /$relative_path"; return 1; }
        msg_success "Restored /$relative_path"
    else
        msg_info "Skipping /$relative_path because it was not found in the archive"
    fi
}

restore_config() {
    pre_restore_checks

    echo
    read -r -p "Restore will overwrite system configuration files. Continue? [y/N]: " confirm
    confirm="${confirm,,}"
    if [[ "$confirm" != "y" && "$confirm" != "yes" ]]; then
        die "Restore aborted by user."
    fi

    msg_info "Restore file: $TARGET_FILE"

    restore_item_if_exists "etc/hosts" || die "Restore failed"
    restore_item_if_exists "etc/hostname" || die "Restore failed"
    restore_item_if_exists "etc/resolv.conf" || die "Restore failed"
    restore_item_if_exists "etc/network/interfaces" || die "Restore failed"
    restore_item_if_exists "etc/vzdump.conf" || die "Restore failed"
    restore_item_if_exists "etc/modules" || die "Restore failed"
    restore_item_if_exists "etc/modules-load.d" || die "Restore failed"
    restore_item_if_exists "etc/modprobe.d" || die "Restore failed"
    restore_item_if_exists "etc/sysctl.conf" || die "Restore failed"
    restore_item_if_exists "etc/sysctl.d" || die "Restore failed"
    restore_item_if_exists "etc/cron.d" || die "Restore failed"
    restore_item_if_exists "etc/cron.daily" || die "Restore failed"
    restore_item_if_exists "etc/cron.weekly" || die "Restore failed"
    restore_item_if_exists "etc/cron.monthly" || die "Restore failed"
    restore_item_if_exists "etc/aliases" || die "Restore failed"
    restore_item_if_exists "root/.ssh" || die "Restore failed"

    if [[ -d "$TMP_RESTORE_DIR/etc/pve" ]]; then
        run_cmd "Restoring /etc/pve" cp -a "$TMP_RESTORE_DIR/etc/pve/." /etc/pve/ \
            || die "Failed restoring /etc/pve"
    else
        msg_info "Skipping /etc/pve because it was not found in the archive"
    fi

    run_cmd_quiet "Reloading systemd daemon" systemctl daemon-reload || true
    run_cmd_quiet "Restarting pve-cluster service" systemctl restart pve-cluster || true
    run_cmd_quiet "Restarting pvedaemon service" systemctl restart pvedaemon || true
    run_cmd_quiet "Restarting pveproxy service" systemctl restart pveproxy || true
    run_cmd_quiet "Restarting pvestatd service" systemctl restart pvestatd || true

    msg_success "Restore completed successfully"
    msg_info "A reboot is strongly recommended after restore."
}

main() {
    trap cleanup EXIT

    print_banner
    require_root
    prompt_action
    require_commands
    detect_proxmox

    if [[ "$ACTION" == "backup" ]]; then
        prompt_filename "Enter host configuration backup filename (.tar.gz or .tgz): "
        prompt_backup_guests
        backup_host_config
        backup_guests
    else
        prompt_filename "Enter restore filename (.tar.gz or .tgz): "
        restore_config
    fi
}

main "$@"