#!/bin/bash

# Enable strict error handling
set -euo pipefail
trap 'handle_error $? $LINENO' ERR

# Configuration
BACKUP_SCRIPT="/home/litespeed/mb-backups/backup_all.sh"
LOG_DIR="/home/litespeed/mb-backups/logs/auto"
ACTION_LOG_FILE="${LOG_DIR}/schedule_actions.log"
BACKUP_LOG_PREFIX="${LOG_DIR}/backup_"
PASSWORD_FILE="/etc/restic-password"
LOCK_FILE="/tmp/backup_schedule.lock"
MAX_RETRIES=3
RETRY_DELAY=5

# Error handling function
handle_error() {
    local exit_code=$1
    local line_no=$2
    log_action "ERROR: Command failed at line ${line_no} with exit code ${exit_code}"
    cleanup_and_exit 1
}

# Cleanup function
cleanup_and_exit() {
    local exit_code=$1
    [ -f "$LOCK_FILE" ] && rm -f "$LOCK_FILE"
    exit "${exit_code}"
}

# Ensure single instance
ensure_single_instance() {
    if [ -f "$LOCK_FILE" ]; then
        if kill -0 "$(cat "$LOCK_FILE")" 2>/dev/null; then
            log_action "ERROR: Another instance is running"
            exit 1
        fi
    fi
    echo $$ > "$LOCK_FILE"
}

# Enhanced logging function
log_action() {
    local timestamp=$(date +'%Y-%m-%d %H:%M:%S')
    local log_msg="[$timestamp] $1"
    echo "$log_msg" | tee -a "$ACTION_LOG_FILE"
    
    # Log critical errors to system log
    if [[ "$1" == *"ERROR"* ]]; then
        logger -t "backup-schedule" "$1"
    fi
}

# Validate system requirements
check_system_requirements() {
    local required_space=1048576  # 1GB in KB
    local available_space=$(df -k "$LOG_DIR" | awk 'NR==2 {print $4}')
    
    if [ ! -x "$BACKUP_SCRIPT" ]; then
        log_action "ERROR: Backup script not executable or not found"
        return 1
    fi
    
    if [ "$available_space" -lt "$required_space" ]; then
        log_action "ERROR: Insufficient disk space"
        return 1
    }
    
    return 0
}

# Validate cron schedule
validate_schedule() {
    local schedule="$1"
    if ! [[ $schedule =~ ^[0-9,\*/-]+ [0-9,\*/-]+ [0-9,\*/-]+ [0-9,\*/-]+ [0-9,\*/-]+$ ]]; then
        log_action "ERROR: Invalid cron schedule format: $schedule"
        return 1
    fi
    return 0
}

# Get Restic password with retry mechanism
get_restic_password() {
    local retry_count=0
    while [ $retry_count -lt $MAX_RETRIES ]; do
        if [ -f "$PASSWORD_FILE" ] && [ -s "$PASSWORD_FILE" ]; then
            cat "$PASSWORD_FILE"
            return 0
        fi
        retry_count=$((retry_count + 1))
        sleep $RETRY_DELAY
    done
    log_action "ERROR: Failed to retrieve Restic password after $MAX_RETRIES attempts"
    return 1
}

# Enhanced cron job management
add_update_cron_job() {
    local action="$1"
    local schedule="$2"
    local restic_password="$3"
    local temp_crontab="/tmp/temp_crontab.$$"
    
    # Validate inputs
    if ! validate_schedule "$schedule"; then
        return 1
    fi
    
    # Prepare cron command with environment variables and error handling
    local CMD="RESTIC_PASSWORD=\"$restic_password\" \
        LOG_FILE=\"${BACKUP_LOG_PREFIX}\$(date +\\%Y-\\%m-\\%d_\\%H-\\%M-\\%S).log\" \
        $BACKUP_SCRIPT > \"\$LOG_FILE\" 2>&1 || echo \"\$(date) Backup failed\" >> \"$ACTION_LOG_FILE\""
    
    # Safely update crontab
    crontab -l 2>/dev/null | grep -v "$BACKUP_SCRIPT" > "$temp_crontab" || true
    echo "$schedule $CMD" >> "$temp_crontab"
    
    if crontab "$temp_crontab"; then
        log_action "Successfully ${action}d backup schedule: $schedule"
        rm -f "$temp_crontab"
        return 0
    else
        log_action "ERROR: Failed to ${action} backup schedule"
        rm -f "$temp_crontab"
        return 1
    fi
}

# Enhanced cron job removal
remove_cron_job() {
    local temp_crontab="/tmp/temp_crontab.$$"
    
    if crontab -l 2>/dev/null | grep -q "$BACKUP_SCRIPT"; then
        crontab -l | grep -v "$BACKUP_SCRIPT" > "$temp_crontab"
        if crontab "$temp_crontab"; then
            log_action "Backup schedule successfully removed"
            rm -f "$temp_crontab"
            return 0
        else
            log_action "ERROR: Failed to remove backup schedule"
            rm -f "$temp_crontab"
            return 1
        fi
    else
        log_action "No backup schedule found to remove"
        return 0
    fi
}

# Main execution
main() {
    # Create log directory with proper permissions
    mkdir -p "$LOG_DIR"
    chmod 750 "$LOG_DIR"
    
    # Ensure single instance
    ensure_single_instance
    
    # Check system requirements
    if ! check_system_requirements; then
        cleanup_and_exit 1
    fi
    
    case $1 in
        add|update)
            if [ "$#" -ne 3 ]; then
                log_action "Usage: $0 {add|update} 'schedule' 'restic_password'"
                log_action "Example: $0 add '0 1 * * *' 'secret_password'"
                cleanup_and_exit 1
            fi
            
            local restic_password="${3:-$(get_restic_password)}"
            if [ -z "$restic_password" ]; then
                cleanup_and_exit 1
            fi
            
            add_update_cron_job "$1" "$2" "$restic_password"
            ;;
        remove)
            remove_cron_job
            ;;
        status)
            crontab -l | grep "$BACKUP_SCRIPT" || echo "No backup schedule found"
            ;;
        *)
            log_action "Invalid action: $1. Use add, update, remove, or status"
            cleanup_and_exit 1
            ;;
    esac
    
    cleanup_and_exit 0
}

# Execute main function with all arguments
main "$@"