#!/bin/bash

# Exit on error
set -e

# Configuration
BACKUP_REPO_PATH="/mnt/backups"
PASSWORD_FILE="/etc/restic-password"
LOG_DIR="/home/jelastic/mb-backups/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 with proper permissions
mkdir -p "$LOG_DIR"
chmod 755 "$LOG_DIR"

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

# Function: Get Restic Password with robust access handling
get_restic_password() {
    local password=""
    
    # Method 1: Try direct file access (works for root and litespeed users)
    if [ -f "$PASSWORD_FILE" ] && [ -r "$PASSWORD_FILE" ]; then
        password=$(cat "$PASSWORD_FILE" 2>/dev/null || echo "")
        if [ -n "$password" ]; then
            log_message "Password accessed directly from $PASSWORD_FILE"
            echo "$password"
            return 0
        fi
    fi
    
    # Method 2: Try sudo access (for non-root users who can sudo)
    if [ "$EUID" -ne 0 ] && command -v sudo >/dev/null 2>&1; then
        if sudo -n test -r "$PASSWORD_FILE" 2>/dev/null; then
            password=$(sudo cat "$PASSWORD_FILE" 2>/dev/null || echo "")
            if [ -n "$password" ]; then
                log_message "Password accessed via sudo from $PASSWORD_FILE"
                echo "$password"
                return 0
            fi
        fi
    fi
    
    # Method 3: Check if running as root but password file has wrong permissions
    if [ "$EUID" -eq 0 ] && [ -f "$PASSWORD_FILE" ]; then
        password=$(cat "$PASSWORD_FILE" 2>/dev/null || echo "")
        if [ -n "$password" ]; then
            log_message "Password accessed as root from $PASSWORD_FILE"
            echo "$password"
            return 0
        fi
    fi
    
    # Method 4: Environment variable fallback
    if [ -n "$RESTIC_PASSWORD" ]; then
        log_message "Using RESTIC_PASSWORD environment variable"
        echo "$RESTIC_PASSWORD"
        return 0
    fi
    
    # Method 5: Check alternative locations
    local alt_locations=(
        "/home/jelastic/.restic-password"
        "/home/litespeed/.restic-password"
        "/home/nginx/.restic-password"
        "/root/.restic-password"
    )
    
    for alt_file in "${alt_locations[@]}"; do
        if [ -f "$alt_file" ] && [ -r "$alt_file" ]; then
            password=$(cat "$alt_file" 2>/dev/null || echo "")
            if [ -n "$password" ]; then
                log_message "Password found at alternative location: $alt_file"
                echo "$password"
                return 0
            fi
        fi
    done
    
    return 1
}

# 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 with robust password handling
validate_repository() {
    log_message "Attempting to access Restic repository..."
    
    # Get password using robust method
    local restic_password
    if ! restic_password=$(get_restic_password); then
        log_message "ERROR: Unable to access Restic password from any source."
        log_message "Checked locations:"
        log_message "  - Primary: $PASSWORD_FILE"
        log_message "  - Environment: RESTIC_PASSWORD"
        log_message "  - Alternative locations in home directories"
        log_message "Current user: $(whoami) (UID: $EUID)"
        log_message "Password file permissions: $(ls -la $PASSWORD_FILE 2>/dev/null || echo 'File not found')"
        exit 1
    fi
    
    # Export password and repository
    export RESTIC_PASSWORD="$restic_password"
    export RESTIC_REPOSITORY="$BACKUP_REPO_PATH"

    # Set the password and repository
    export RESTIC_PASSWORD=$(cat /etc/restic-password)
    export RESTIC_REPOSITORY="/mnt/backups"

    # Initialize the repository
    restic init

    # Verify it worked
    restic snapshots

    # Test repository access
    if ! restic snapshots &>/dev/null; then
        log_message "ERROR: Unable to access the Restic repository at $BACKUP_REPO_PATH"
        log_message "Repository path exists: $([ -d "$BACKUP_REPO_PATH" ] && echo 'Yes' || echo 'No')"
        log_message "Repository contents: $(ls -la "$BACKUP_REPO_PATH" 2>/dev/null | wc -l || echo '0') items"
        
        # Check if repository needs initialization
        if [ ! -d "$BACKUP_REPO_PATH" ] || [ -z "$(ls -A "$BACKUP_REPO_PATH" 2>/dev/null)" ]; then
            log_message "Repository appears to be empty or non-existent. Initializing repository..."
            if restic init; then
                log_message "Repository initialized successfully."
            else
                log_message "ERROR: Failed to initialize repository."
                exit 1
            fi
        else
            # Repository exists but might not be a valid Restic repository
            log_message "Repository directory exists but may not be a valid Restic repository."
            log_message "Attempting to initialize as new repository..."
            if restic init; then
                log_message "Repository initialized successfully."
            else
                log_message "ERROR: Repository exists but is not accessible. Manual intervention required."
                log_message "Directory contents:"
                ls -la "$BACKUP_REPO_PATH" 2>/dev/null || echo "Cannot list directory"
                exit 1
            fi
        fi
        
        # Test access again after initialization
        if ! restic snapshots &>/dev/null; then
            log_message "ERROR: Repository initialization succeeded but still cannot access repository."
            exit 1
        fi
    fi
    
    log_message "Repository access validated successfully."
}

# 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..."
    log_message "Running as user: $(whoami) (UID: $EUID)"
    log_message "Repository path: $BACKUP_REPO_PATH"
    
    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