#!/bin/bash

# Exit on error
set -e

# Configuration
BACKUP_REPO_PATH="/mnt/backups"
PASSWORD_FILE="/etc/restic-password"
LOG_DIR="/home/litespeed/logs"
LOG_FILE="${LOG_DIR}/repo_stats_$(date +'%Y-%m-%d').log"
RETENTION_POLICY="--keep-last 7 --prune" # Modify retention policy as needed

# Ensure the log directory exists
mkdir -p "$LOG_DIR"

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

# Function: Validate dependencies
validate_dependencies() {
    for cmd in restic; do
        if ! command -v "$cmd" &>/dev/null; then
            log_message "ERROR: Required command '$cmd' not found."
            exit 1
        fi
    done
}

# Function: Validate repository access
validate_repository() {
    if [ ! -f "$PASSWORD_FILE" ]; then
        log_message "ERROR: Password file not found at $PASSWORD_FILE."
        exit 1
    fi
    export RESTIC_PASSWORD=$(cat "$PASSWORD_FILE")
    export RESTIC_REPOSITORY="$BACKUP_REPO_PATH"

    if ! restic snapshots &>/dev/null; then
        log_message "ERROR: Unable to access the Restic repository. Check password and repository path."
        exit 1
    fi
    log_message "Repository access validated."
}

# Function: Check repository stats
check_repository_stats() {
    log_message "Fetching repository stats..."
    if restic stats --mode restore-size >> "$LOG_FILE" 2>&1; then
        log_message "Repository stats fetched successfully."
    else
        log_message "ERROR: Failed to fetch repository stats."
        exit 1
    fi
}

# Function: Apply retention policy
apply_retention_policy() {
    log_message "Applying retention policy: $RETENTION_POLICY"
    if restic forget $RETENTION_POLICY >> "$LOG_FILE" 2>&1; then
        log_message "Retention policy applied successfully. Unnecessary snapshots pruned."
    else
        log_message "ERROR: Failed to apply retention policy."
        exit 1
    fi
}

# Function: Check for stale locks
check_and_remove_stale_locks() {
    log_message "Checking for stale locks in the repository..."
    if [ -d "${BACKUP_REPO_PATH}/locks" ] && [ "$(ls -A "${BACKUP_REPO_PATH}/locks" 2>/dev/null)" ]; then
        log_message "Stale locks detected. Removing locks..."
        if restic unlock; then
            log_message "Stale locks removed successfully."
        else
            log_message "ERROR: Failed to remove stale locks."
            exit 1
        fi
    else
        log_message "No stale locks found."
    fi
}

# Function: Perform integrity check
perform_integrity_check() {
    log_message "Performing repository integrity check..."
    if restic check --read-data-subset=5% >> "$LOG_FILE" 2>&1; then
        log_message "Repository integrity check passed."
    else
        log_message "ERROR: Repository integrity check failed."
        exit 1
    fi
}

# Main execution
main() {
    log_message "Starting repository stats check and maintenance..."
    validate_dependencies
    validate_repository
    check_and_remove_stale_locks
    check_repository_stats
    apply_retention_policy
    perform_integrity_check
    log_message "Repository stats check and maintenance completed successfully."
}

# Execute main function
main