#!/bin/bash

# Exit immediately on errors
set -e

# Validate input parameters
if [ "$#" -ne 2 ]; then
    echo "Usage: $0 <snapshot_id> <restic_password>"
    exit 1
fi

# Assign input to variables
SNAPSHOT_ID=$1
RESTIC_PASSWORD=$2
RESTIC_REPOSITORY="/mnt/backups"
LOG_DIR="/home/litespeed/mb-backups/logs/restore"
WP_CONFIG="/var/www/webroot/ROOT/wp-config.php"

# Set up the environment
export RESTIC_PASSWORD
mkdir -p "$LOG_DIR"
LOG_FILE="${LOG_DIR}/restore_$(date +'%Y-%m-%d_%H-%M-%S').log"

# Logging function
log() {
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] $1" | tee -a "$LOG_FILE"
}

# Check for required dependencies
for cmd in restic jq mysql; do
    if ! command -v $cmd >/dev/null 2>&1; then
        log "Error: '$cmd' command not found. Please install $cmd."
        exit 1
    fi
done

# Ensure there are no stale locks in the repository
log "Checking for stale locks in the repository..."
if restic -r "$RESTIC_REPOSITORY" list locks | grep -q "lock"; then
    log "Stale lock detected. Unlocking repository..."
    restic -r "$RESTIC_REPOSITORY" unlock
    log "Repository unlocked successfully."
else
    log "No stale locks found. Proceeding with restoration."
fi

# Extract database credentials from wp-config.php
if [ ! -f "$WP_CONFIG" ]; then
    log "Error: wp-config.php not found at $WP_CONFIG. Ensure WordPress is installed."
    exit 1
fi

DB_NAME=$(awk -F"'" '/define\( *'"'"'DB_NAME'"'"'/{print $4}' "$WP_CONFIG")
DB_USER=$(awk -F"'" '/define\( *'"'"'DB_USER'"'"'/{print $4}' "$WP_CONFIG")
DB_PASSWORD=$(awk -F"'" '/define\( *'"'"'DB_PASSWORD'"'"'/{print $4}' "$WP_CONFIG")

if [ -z "$DB_NAME" ] || [ -z "$DB_USER" ] || [ -z "$DB_PASSWORD" ]; then
    log "Error: Could not extract database credentials from wp-config.php."
    exit 1
fi

log "Starting restoration for snapshot ID: $SNAPSHOT_ID..."

# Retrieve snapshot data
SNAPSHOT_DATA=$(restic -r "$RESTIC_REPOSITORY" snapshots --json | jq -r ".[] | select(.short_id == \"$SNAPSHOT_ID\")")
if [ -z "$SNAPSHOT_DATA" ]; then
    log "Error: Snapshot ID $SNAPSHOT_ID not found in repository."
    exit 1
fi

SNAPSHOT_PATH=$(echo "$SNAPSHOT_DATA" | jq -r ".paths[] // empty")
SNAPSHOT_TAGS=$(echo "$SNAPSHOT_DATA" | jq -r ".tags[] // empty")

if [ -z "$SNAPSHOT_PATH" ]; then
    log "Error: Snapshot $SNAPSHOT_ID does not contain valid paths."
    exit 1
fi

log "Snapshot Data Retrieved: Path=$SNAPSHOT_PATH, Tags=$SNAPSHOT_TAGS"

# Determine the restoration type based on the tags
if [[ "$SNAPSHOT_TAGS" == *"wordpress_db"* ]] && [[ "$SNAPSHOT_PATH" == "/stdin" ]]; then
    log "Detected database backup. Restoring database $DB_NAME..."
    if restic -r "$RESTIC_REPOSITORY" dump "$SNAPSHOT_ID" stdin | mysql -u "$DB_USER" -p"$DB_PASSWORD" "$DB_NAME"; then
        log "Database restoration completed successfully for $DB_NAME from Snapshot ID: $SNAPSHOT_ID."
    else
        log "Error: Database restoration failed for Snapshot ID: $SNAPSHOT_ID."
        exit 1
    fi
elif [[ "$SNAPSHOT_TAGS" == *"core_files"* ]]; then
    RESTORE_DIR="/var/www/webroot/ROOT"
    log "Detected core files backup. Restoring to $RESTORE_DIR..."
    if restic -r "$RESTIC_REPOSITORY" restore "$SNAPSHOT_ID" --target "$RESTORE_DIR"; then
        log "Core files restoration completed successfully for Snapshot ID: $SNAPSHOT_ID."
    else
        log "Error: Core files restoration failed for Snapshot ID: $SNAPSHOT_ID."
        exit 1
    fi
elif [[ "$SNAPSHOT_TAGS" == *"media_themes"* ]]; then
    RESTORE_DIR="/var/www/webroot/ROOT/wp-content/uploads"
    log "Detected media files backup. Restoring to $RESTORE_DIR..."
    if restic -r "$RESTIC_REPOSITORY" restore "$SNAPSHOT_ID" --target "$RESTORE_DIR"; then
        log "Media files restoration completed successfully for Snapshot ID: $SNAPSHOT_ID."
    else
        log "Error: Media files restoration failed for Snapshot ID: $SNAPSHOT_ID."
        exit 1
    fi
else
    log "Unknown snapshot type. Attempting general restoration to /..."
    TEMP_RESTORE_DIR="/tmp/restoration_$SNAPSHOT_ID"
    mkdir -p "$TEMP_RESTORE_DIR"
    if restic -r "$RESTIC_REPOSITORY" restore "$SNAPSHOT_ID" --target "$TEMP_RESTORE_DIR"; then
        log "General file restoration completed successfully to temporary directory: $TEMP_RESTORE_DIR."
    else
        log "Error: File restoration failed for Snapshot ID: $SNAPSHOT_ID."
        exit 1
    fi
fi

log "Restoration process completed successfully."