#!/bin/bash

# Exit on errors and enable error tracing
set -e
trap 'log "ERROR: An error occurred during the restoration process."' ERR

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

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

# Ensure the log directory exists
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"
}

# Validate required dependencies
validate_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
}

# Set up Restic environment
setup_restic_environment() {
    if [ ! -f "$RESTIC_PASSWORD_FILE" ]; then
        log "ERROR: Restic password file not found at $RESTIC_PASSWORD_FILE."
        exit 1
    fi

    export RESTIC_PASSWORD=$(cat "$RESTIC_PASSWORD_FILE")
    export RESTIC_REPOSITORY="$RESTIC_REPOSITORY"
    log "Restic environment set up successfully."
}

# Ensure no stale locks exist in the repository
remove_stale_locks() {
    log "Checking for stale locks in the repository..."
    if restic list locks | grep -q "lock"; then
        log "Stale lock detected. Unlocking repository..."
        restic unlock || {
            log "ERROR: Failed to remove stale locks."
            exit 1
        }
        log "Repository unlocked successfully."
    else
        log "No stale locks found. Proceeding with restoration."
    fi
}

# Extract database credentials from wp-config.php
extract_db_credentials() {
    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")
    DB_HOST=$(awk -F"'" '/define\( *'"'"'DB_HOST'"'"'/{print $4}' "$WP_CONFIG")

    # Set default DB_HOST if empty
    if [ -z "$DB_HOST" ]; then
        DB_HOST="localhost"
    fi

    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
}

# Retrieve snapshot metadata
retrieve_snapshot_metadata() {
    log "Retrieving metadata for snapshot ID: $SNAPSHOT_ID..."
    SNAPSHOT_DATA=$(restic snapshots --json | jq -r ".[] | select(.short_id == \"$SNAPSHOT_ID\")")

    if [ -z "$SNAPSHOT_DATA" ]; then
        log "ERROR: Snapshot ID $SNAPSHOT_ID not found in the 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 metadata retrieved: Path=$SNAPSHOT_PATH, Tags=$SNAPSHOT_TAGS"
}

# Perform the restoration based on snapshot tags
restore_snapshot() {
    if [[ "$SNAPSHOT_TAGS" == *"wordpress_db"* ]] && [[ "$SNAPSHOT_PATH" == *".sql"* ]]; then
        log "Detected database backup. Restoring database $DB_NAME..."
        log "Using DB credentials: User=$DB_USER, Host=$DB_HOST, DB=$DB_NAME"
        
        if restic dump "$SNAPSHOT_ID" "$SNAPSHOT_PATH" | mysql -h "$DB_HOST" -u "$DB_USER" -p"$DB_PASSWORD" "$DB_NAME"; then
            log "Database restoration completed successfully for $DB_NAME."
        else
            log "ERROR: Database restoration failed. Verify MySQL credentials or permissions."
            log "Check that 'DB_USER=$DB_USER' has necessary privileges on 'DB_NAME=$DB_NAME'."
            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 restore "$SNAPSHOT_ID" --target "$RESTORE_DIR"; then
            log "Core files restoration completed successfully."
        else
            log "ERROR: Core files restoration failed."
            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 restore "$SNAPSHOT_ID" --target "$RESTORE_DIR"; then
            log "Media files restoration completed successfully."
        else
            log "ERROR: Media files restoration failed."
            exit 1
        fi
    else
        log "Unknown snapshot type. Attempting general restoration to a temporary directory..."
        TEMP_RESTORE_DIR="/tmp/restoration_$SNAPSHOT_ID"
        mkdir -p "$TEMP_RESTORE_DIR"
        if restic restore "$SNAPSHOT_ID" --target "$TEMP_RESTORE_DIR"; then
            log "General restoration completed successfully to $TEMP_RESTORE_DIR."
        else
            log "ERROR: General restoration failed."
            exit 1
        fi
    fi
}

# Main execution flow
validate_dependencies
setup_restic_environment
remove_stale_locks
extract_db_credentials
retrieve_snapshot_metadata
restore_snapshot

log "Restoration process completed successfully."