#!/bin/bash

# Exit on errors and enable comprehensive error tracing
set -e
trap 'log "CRITICAL ERROR: Full backup session restoration failed at line $LINENO. Aborting all operations."' ERR

# Validate input parameters
if [ "$#" -ne 1 ]; then
    echo "Usage: $0 <backup_session_tag>"
    echo ""
    echo "Example: $0 manual-backup-2024-01-15_10-30-45"
    echo "         $0 automated-backup-2024-01-15_02-00-30"
    echo ""
    echo "Use view_snapshots.sh to find available backup session tags."
    exit 1
fi

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

# Required component tags for a complete full backup
REQUIRED_CORE_TAG="core_files"
REQUIRED_MEDIA_TAG="media_themes"
REQUIRED_DB_TAG="wordpress_db"

# Ensure the log directory exists
mkdir -p "$LOG_DIR"
LOG_FILE="${LOG_DIR}/restore_full_backup_$(date +'%Y-%m-%d_%H-%M-%S').log"

# Create temporary validation directory
mkdir -p "$TEMP_VALIDATION_DIR"

# Cleanup function
cleanup() {
    log "Cleaning up temporary files..."
    rm -rf "$TEMP_VALIDATION_DIR" 2>/dev/null || true
}

# Set cleanup trap
trap 'cleanup; log "CRITICAL ERROR: Full backup session restoration failed. Cleanup completed."' ERR
trap 'cleanup' EXIT

# Logging function with enhanced formatting
log() {
    local level="INFO"
    local message="$1"
    
    if [[ "$message" == CRITICAL* ]] || [[ "$message" == ERROR* ]]; then
        level="ERROR"
    elif [[ "$message" == WARNING* ]]; then
        level="WARN"
    fi
    
    echo "[$(date +'%Y-%m-%d %H:%M:%S')] [$level] $message" | tee -a "$LOG_FILE"
}

# Validate required dependencies
validate_dependencies() {
    log "Validating system dependencies..."
    
    for cmd in restic jq mysql; do
        if ! command -v $cmd >/dev/null 2>&1; then
            log "CRITICAL ERROR: Required command '$cmd' not found. Please install $cmd."
            exit 1
        fi
    done
    
    log "All required dependencies validated successfully."
}

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

    if [ ! -d "$RESTIC_REPOSITORY" ]; then
        log "CRITICAL ERROR: Restic repository not found at $RESTIC_REPOSITORY."
        exit 1
    fi

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

# Remove stale locks from repository
remove_stale_locks() {
    log "Checking for stale locks in repository..."
    
    if restic list locks 2>/dev/null | grep -q "lock"; then
        log "WARNING: Stale locks detected. Attempting to unlock repository..."
        if restic unlock; then
            log "Repository unlocked successfully."
        else
            log "CRITICAL ERROR: Failed to remove stale locks from repository."
            exit 1
        fi
    else
        log "No stale locks found. Repository is clean."
    fi
}

# Validate backup session tag format
validate_session_tag_format() {
    log "Validating backup session tag format: '$BACKUP_SESSION_TAG'"
    
    # Check if tag matches expected patterns
    if [[ ! "$BACKUP_SESSION_TAG" =~ ^(manual|automated)-backup-[0-9]{4}-[0-9]{2}-[0-9]{2}_[0-9]{2}-[0-9]{2}-[0-9]{2}$ ]]; then
        log "CRITICAL ERROR: Invalid backup session tag format."
        log "Expected format: manual-backup-YYYY-MM-DD_HH-MM-SS or automated-backup-YYYY-MM-DD_HH-MM-SS"
        log "Received: '$BACKUP_SESSION_TAG'"
        exit 1
    fi
    
    log "Backup session tag format is valid."
}

# Retrieve and validate all snapshots for the session
retrieve_session_snapshots() {
    log "Retrieving all snapshots for backup session: '$BACKUP_SESSION_TAG'"
    
    # Get all snapshots with the session tag
    local snapshots_json
    snapshots_json=$(restic snapshots --tag "$BACKUP_SESSION_TAG" --json 2>/dev/null)
    
    if [ -z "$snapshots_json" ] || [ "$snapshots_json" = "null" ] || [ "$snapshots_json" = "[]" ]; then
        log "CRITICAL ERROR: No snapshots found with backup session tag '$BACKUP_SESSION_TAG'."
        log "Use 'view_snapshots.sh all' to see available backup sessions."
        exit 1
    fi
    
    # Save snapshots data for validation
    echo "$snapshots_json" > "$TEMP_VALIDATION_DIR/session_snapshots.json"
    
    # Count total snapshots in session
    local snapshot_count
    snapshot_count=$(echo "$snapshots_json" | jq '. | length')
    
    log "Found $snapshot_count snapshots with session tag '$BACKUP_SESSION_TAG'."
    
    # Validate we have exactly 3 snapshots (core, media, database)
    if [ "$snapshot_count" -ne 3 ]; then
        log "CRITICAL ERROR: Incomplete backup session detected."
        log "Expected exactly 3 snapshots (core_files, media_themes, wordpress_db), found $snapshot_count."
        log "This backup session is incomplete and cannot be safely restored."
        exit 1
    fi
    
    log "Snapshot count validation passed: Found expected 3 snapshots."
}

# Validate that all required components are present
validate_backup_completeness() {
    log "Performing comprehensive backup completeness validation..."
    
    local snapshots_json
    snapshots_json=$(cat "$TEMP_VALIDATION_DIR/session_snapshots.json")
    
    # Extract all tags from all snapshots in the session
    local all_tags
    all_tags=$(echo "$snapshots_json" | jq -r '.[].tags[]' | sort | uniq)
    
    # Check for required component tags
    local missing_components=()
    
    if ! echo "$all_tags" | grep -q "^${REQUIRED_CORE_TAG}$"; then
        missing_components+=("core_files")
    fi
    
    if ! echo "$all_tags" | grep -q "^${REQUIRED_MEDIA_TAG}$"; then
        missing_components+=("media_themes")
    fi
    
    if ! echo "$all_tags" | grep -q "^${REQUIRED_DB_TAG}$"; then
        missing_components+=("wordpress_db")
    fi
    
    # Report missing components
    if [ ${#missing_components[@]} -gt 0 ]; then
        log "CRITICAL ERROR: Backup session is missing required components:"
        for component in "${missing_components[@]}"; do
            log "  - Missing: $component"
        done
        log "Cannot proceed with restoration of incomplete backup session."
        exit 1
    fi
    
    log "Backup completeness validation passed: All required components present."
}

# Extract and validate individual component snapshots
extract_component_snapshots() {
    log "Extracting and validating individual component snapshots..."
    
    local snapshots_json
    snapshots_json=$(cat "$TEMP_VALIDATION_DIR/session_snapshots.json")
    
    # Extract core files snapshot
    CORE_SNAPSHOT=$(echo "$snapshots_json" | jq -r ".[] | select(.tags[] | contains(\"$REQUIRED_CORE_TAG\")) | .short_id")
    CORE_SNAPSHOT_FULL=$(echo "$snapshots_json" | jq -r ".[] | select(.tags[] | contains(\"$REQUIRED_CORE_TAG\"))")
    
    # Extract media files snapshot
    MEDIA_SNAPSHOT=$(echo "$snapshots_json" | jq -r ".[] | select(.tags[] | contains(\"$REQUIRED_MEDIA_TAG\")) | .short_id")
    MEDIA_SNAPSHOT_FULL=$(echo "$snapshots_json" | jq -r ".[] | select(.tags[] | contains(\"$REQUIRED_MEDIA_TAG\"))")
    
    # Extract database snapshot
    DB_SNAPSHOT=$(echo "$snapshots_json" | jq -r ".[] | select(.tags[] | contains(\"$REQUIRED_DB_TAG\")) | .short_id")
    DB_SNAPSHOT_FULL=$(echo "$snapshots_json" | jq -r ".[] | select(.tags[] | contains(\"$REQUIRED_DB_TAG\"))")
    
    # Validate all snapshots were found
    if [ -z "$CORE_SNAPSHOT" ] || [ "$CORE_SNAPSHOT" = "null" ]; then
        log "CRITICAL ERROR: Core files snapshot not found in backup session."
        exit 1
    fi
    
    if [ -z "$MEDIA_SNAPSHOT" ] || [ "$MEDIA_SNAPSHOT" = "null" ]; then
        log "CRITICAL ERROR: Media files snapshot not found in backup session."
        exit 1
    fi
    
    if [ -z "$DB_SNAPSHOT" ] || [ "$DB_SNAPSHOT" = "null" ]; then
        log "CRITICAL ERROR: Database snapshot not found in backup session."
        exit 1
    fi
    
    # Validate snapshot IDs are unique (no duplicates)
    if [ "$CORE_SNAPSHOT" = "$MEDIA_SNAPSHOT" ] || [ "$CORE_SNAPSHOT" = "$DB_SNAPSHOT" ] || [ "$MEDIA_SNAPSHOT" = "$DB_SNAPSHOT" ]; then
        log "CRITICAL ERROR: Duplicate snapshot IDs detected. Backup session integrity compromised."
        exit 1
    fi
    
    log "Component snapshot extraction successful:"
    log "  - Core Files Snapshot: $CORE_SNAPSHOT"
    log "  - Media Files Snapshot: $MEDIA_SNAPSHOT"
    log "  - Database Snapshot: $DB_SNAPSHOT"
}

# Validate session timestamp consistency
validate_session_timestamp_consistency() {
    log "Validating backup session timestamp consistency..."
    
    local snapshots_json
    snapshots_json=$(cat "$TEMP_VALIDATION_DIR/session_snapshots.json")
    
    # Extract timestamps from all snapshots
    local timestamps
    timestamps=$(echo "$snapshots_json" | jq -r '.[].time' | sort | uniq)
    
    # Count unique timestamps
    local timestamp_count
    timestamp_count=$(echo "$timestamps" | wc -l)
    
    # Allow for small time differences (up to 30 minutes) for backup session
    local first_timestamp
    local last_timestamp
    first_timestamp=$(echo "$timestamps" | head -n1)
    last_timestamp=$(echo "$timestamps" | tail -n1)
    
    # Convert to epoch for comparison
    local first_epoch
    local last_epoch
    first_epoch=$(date -d "$first_timestamp" +%s 2>/dev/null || echo "0")
    last_epoch=$(date -d "$last_timestamp" +%s 2>/dev/null || echo "0")
    
    # Calculate time difference (30 minutes = 1800 seconds)
    local time_diff=$((last_epoch - first_epoch))
    
    if [ $time_diff -gt 1800 ]; then
        log "WARNING: Large time difference detected between snapshots in session ($time_diff seconds)."
        log "First snapshot: $first_timestamp"
        log "Last snapshot: $last_timestamp"
        log "This may indicate snapshots from different backup sessions."
        
        # Ask for confirmation in interactive mode
        if [ -t 0 ]; then
            echo "Do you want to continue with restoration? (yes/no): "
            read -r confirmation
            if [ "$confirmation" != "yes" ]; then
                log "Restoration cancelled by user due to timestamp inconsistency."
                exit 1
            fi
        else
            log "CRITICAL ERROR: Timestamp inconsistency detected in non-interactive mode. Aborting."
            exit 1
        fi
    fi
    
    log "Session timestamp consistency validation completed."
}

# Extract database credentials from wp-config.php
extract_db_credentials() {
    log "Extracting database credentials from WordPress configuration..."
    
    if [ ! -f "$WP_CONFIG" ]; then
        log "CRITICAL ERROR: wp-config.php not found at $WP_CONFIG."
        log "Ensure WordPress is properly installed before restoration."
        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 "CRITICAL ERROR: Could not extract complete database credentials from wp-config.php."
        log "Extracted: DB_NAME='$DB_NAME', DB_USER='$DB_USER', DB_HOST='$DB_HOST'"
        exit 1
    fi
    
    log "Database credentials extracted successfully for database: $DB_NAME"
}

# Test database connectivity before restoration
test_database_connectivity() {
    log "Testing database connectivity before restoration..."
    
    if ! mysql -h "$DB_HOST" -u "$DB_USER" -p"$DB_PASSWORD" -e "SELECT 1;" >/dev/null 2>&1; then
        log "CRITICAL ERROR: Cannot connect to database."
        log "Please verify database credentials and ensure MySQL service is running."
        exit 1
    fi
    
    log "Database connectivity test passed."
}

# Create pre-restoration backup of current state
create_pre_restoration_backup() {
    log "Creating pre-restoration backup of current state..."
    
    local pre_backup_tag="pre-restore-$(date +'%Y-%m-%d_%H-%M-%S')"
    local pre_backup_dir="/tmp/pre_restore_backup_$$"
    
    mkdir -p "$pre_backup_dir"
    
    # Backup current database
    log "Backing up current database state..."
    if mysqldump -h "$DB_HOST" -u "$DB_USER" -p"$DB_PASSWORD" "$DB_NAME" > "$pre_backup_dir/current_database.sql" 2>/dev/null; then
        log "Current database backed up successfully."
    else
        log "WARNING: Could not backup current database state."
    fi
    
    # Create snapshot of current state
    if restic backup "$pre_backup_dir" --tag "$pre_backup_tag" --tag "pre-restoration-backup" >/dev/null 2>&1; then
        log "Pre-restoration backup created with tag: $pre_backup_tag"
        echo "$pre_backup_tag" > "$TEMP_VALIDATION_DIR/pre_restore_tag.txt"
    else
        log "WARNING: Could not create pre-restoration backup snapshot."
    fi
    
    rm -rf "$pre_backup_dir"
}

# Restore database component
restore_database() {
    log "Starting database restoration from snapshot: $DB_SNAPSHOT"
    
    # Get database snapshot details
    local db_snapshot_data
    db_snapshot_data=$(echo "$DB_SNAPSHOT_FULL")
    
    local db_paths
    db_paths=$(echo "$db_snapshot_data" | jq -r '.paths[]' 2>/dev/null || echo "")
    
    if [ -z "$db_paths" ]; then
        log "CRITICAL ERROR: Database snapshot contains no valid paths."
        exit 1
    fi
    
    log "Restoring database from path: $db_paths"
    log "Target database: $DB_NAME on $DB_HOST"
    
    # Restore database with error handling
    if restic dump "$DB_SNAPSHOT" "$db_paths" | mysql -h "$DB_HOST" -u "$DB_USER" -p"$DB_PASSWORD" "$DB_NAME"; then
        log "Database restoration completed successfully."
    else
        log "CRITICAL ERROR: Database restoration failed."
        log "You may need to restore from the pre-restoration backup if available."
        exit 1
    fi
}

# Restore core files component
restore_core_files() {
    log "Starting core files restoration from snapshot: $CORE_SNAPSHOT"
    
    local restore_target="/var/www/webroot/ROOT"
    
    log "Restoring core files to: $restore_target"
    
    if restic restore "$CORE_SNAPSHOT" --target "$restore_target"; then
        log "Core files restoration completed successfully."
    else
        log "CRITICAL ERROR: Core files restoration failed."
        exit 1
    fi
}

# Restore media files component
restore_media_files() {
    log "Starting media files restoration from snapshot: $MEDIA_SNAPSHOT"
    
    local restore_target="/var/www/webroot/ROOT"
    
    log "Restoring media files to: $restore_target"
    
    if restic restore "$MEDIA_SNAPSHOT" --target "$restore_target"; then
        log "Media files restoration completed successfully."
    else
        log "CRITICAL ERROR: Media files restoration failed."
        exit 1
    fi
}

# Verify restoration integrity
verify_restoration_integrity() {
    log "Performing post-restoration integrity verification..."
    
    # Check WordPress core files
    if [ ! -f "/var/www/webroot/ROOT/wp-config.php" ]; then
        log "WARNING: wp-config.php not found after restoration."
    fi
    
    if [ ! -d "/var/www/webroot/ROOT/wp-content" ]; then
        log "WARNING: wp-content directory not found after restoration."
    fi
    
    # Check database connectivity
    if mysql -h "$DB_HOST" -u "$DB_USER" -p"$DB_PASSWORD" -e "SHOW TABLES;" "$DB_NAME" >/dev/null 2>&1; then
        log "Database connectivity verified after restoration."
    else
        log "WARNING: Database connectivity issues detected after restoration."
    fi
    
    log "Restoration integrity verification completed."
}

# Display restoration summary
display_restoration_summary() {
    log "=== FULL BACKUP SESSION RESTORATION SUMMARY ==="
    log "Backup Session Tag: $BACKUP_SESSION_TAG"
    log "Components Restored:"
    log "  - Database (Snapshot: $DB_SNAPSHOT)"
    log "  - Core Files (Snapshot: $CORE_SNAPSHOT)" 
    log "  - Media Files (Snapshot: $MEDIA_SNAPSHOT)"
    log "Restoration completed at: $(date +'%Y-%m-%d %H:%M:%S')"
    
    if [ -f "$TEMP_VALIDATION_DIR/pre_restore_tag.txt" ]; then
        local pre_restore_tag
        pre_restore_tag=$(cat "$TEMP_VALIDATION_DIR/pre_restore_tag.txt")
        log "Pre-restoration backup available with tag: $pre_restore_tag"
    fi
    
    log "Full restoration log saved to: $LOG_FILE"
    log "=============================================="
}

# Main execution flow
main() {
    log "Starting full backup session restoration for: '$BACKUP_SESSION_TAG'"
    log "Process ID: $$"
    log "Log file: $LOG_FILE"
    
    # Phase 1: Environment Setup and Validation
    validate_dependencies
    setup_restic_environment
    remove_stale_locks
    validate_session_tag_format
    
    # Phase 2: Backup Session Discovery and Validation
    retrieve_session_snapshots
    validate_backup_completeness
    extract_component_snapshots
    validate_session_timestamp_consistency
    
    # Phase 3: Pre-restoration Preparation
    extract_db_credentials
    test_database_connectivity
    create_pre_restoration_backup
    
    # Phase 4: Restoration Execution (Order is critical)
    log "Beginning restoration sequence..."
    restore_database
    restore_core_files
    restore_media_files
    
    # Phase 5: Post-restoration Verification
    verify_restoration_integrity
    display_restoration_summary
    
    log "Full backup session restoration completed successfully!"
}

# Execute main function
main