#!/bin/bash

# Exit on error
set -e

# Configuration
BACKUP_REPO_PATH="/mnt/backups"
PASSWORD_FILE="/etc/restic-password"
LOG_DIR="/home/litespeed/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 - simplified approach based on working check_backup_repo.sh
validate_repository() {
    log_message "Attempting to access Restic repository..."
    
    # Ensure the backup repository path exists
    mkdir -p "$BACKUP_REPO_PATH"
    
    # Get password using robust method, but fallback to simple method if needed
    local restic_password
    if ! restic_password=$(get_restic_password); then
        log_message "WARNING: Unable to access password from robust methods, trying simple fallback..."
        
        # Fallback to simple password access (like working script)
        if [ ! -f "$PASSWORD_FILE" ]; then
            log_message "Password file not found. Creating a new one with a default password."
            echo "default-password" > "$PASSWORD_FILE"
        fi
        restic_password=$(cat "$PASSWORD_FILE")
    fi
    
    # Export password and repository (like working script)
    export RESTIC_PASSWORD="$restic_password"
    export RESTIC_REPOSITORY="$BACKUP_REPO_PATH"
    
    # Use the same logic as the working check_backup_repo.sh
    # Check if the repository contains files
    if [ "$(find "$BACKUP_REPO_PATH" -mindepth 1 2>/dev/null | wc -l)" -gt 0 ]; then
        log_message "Repository directory contains files. Checking if it's a valid Restic repository..."
        
        # Check if this looks like a Restic repository structure
        if [ -f "$BACKUP_REPO_PATH/config" ] && [ -d "$BACKUP_REPO_PATH/data" ]; then
            log_message "Detected valid Restic repository structure. Testing access..."
            
            # First try to unlock in case there are stale locks
            log_message "Checking for stale locks before testing access..."
            restic unlock 2>/dev/null || true
            
            # Test if repository is accessible with current password
            if restic snapshots &>/dev/null; then
                log_message "Repository access successful with current password."
            else
                log_message "ERROR: Repository structure is valid but access failed."
                log_message "This usually indicates a password mismatch."
                log_message "Repository config file permissions: $(ls -la "$BACKUP_REPO_PATH/config" 2>/dev/null)"
                log_message "Trying to unlock repository first..."
                
                # Try unlocking again with verbose output
                if ! restic unlock 2>/dev/null; then
                    log_message "Warning: Could not unlock repository, but continuing..."
                fi
                
                # Try accessing again after unlock
                if restic snapshots &>/dev/null; then
                    log_message "Repository access successful after unlock."
                else
                    log_message "ERROR: Repository access still failed after unlock."
                    log_message "Please verify the password in $PASSWORD_FILE matches the repository."
                    
                    # Show more diagnostic information
                    log_message "Diagnostic information:"
                    log_message "- Repository path: $BACKUP_REPO_PATH"
                    log_message "- Password file: $PASSWORD_FILE"
                    log_message "- Password file exists: $([ -f "$PASSWORD_FILE" ] && echo 'Yes' || echo 'No')"
                    log_message "- Password file readable: $([ -r "$PASSWORD_FILE" ] && echo 'Yes' || echo 'No')"
                    log_message "- Current user: $(whoami) (UID: $EUID)"
                    
                    # Try one more time with the working script's simple approach
                    log_message "Attempting simple password access method..."
                    if [ -f "$PASSWORD_FILE" ]; then
                        simple_password=$(cat "$PASSWORD_FILE" 2>/dev/null || echo "")
                        if [ -n "$simple_password" ]; then
                            export RESTIC_PASSWORD="$simple_password"
                            if restic snapshots &>/dev/null; then
                                log_message "SUCCESS: Repository accessible with simple password method."
                            else
                                log_message "ERROR: Repository still not accessible. Password mismatch likely."
                                exit 1
                            fi
                        else
                            log_message "ERROR: Could not read password from file."
                            exit 1
                        fi
                    else
                        log_message "ERROR: Password file not found."
                        exit 1
                    fi
                fi
            fi
        else
            # Directory has files but doesn't look like Restic repository
            log_message "ERROR: Repository directory contains files but is not a valid Restic repository structure."
            log_message "Repository contents: $(ls -la "$BACKUP_REPO_PATH" 2>/dev/null | head -10)"
            exit 1
        fi
        log_message "Valid Restic repository found and accessible."
    else
        # Initialize a new Restic repository if empty (same as working script)
        log_message "No files found in the backup repository. Initializing a new repository..."
        if ! restic init; then
            log_message "ERROR: Failed to initialize the backup repository."
            exit 1
        fi
        log_message "Backup repository initialized successfully."
    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 (simplified approach)
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