#!/bin/bash

set -e

# Restic Service Startup Script for HuggingFace Environment
# Supports automatic backup and restore functionality

# Logging functions
log() {
    echo "[$(date '+%Y-%m-%d %H:%M:%S')] [RESTIC] $*"
}

log_success() {
    echo -e "\033[32m[$(date '+%Y-%m-%d %H:%M:%S')] [SUCCESS] $*\033[0m"
}

log_error() {
    echo -e "\033[31m[$(date '+%Y-%m-%d %H:%M:%S')] [ERROR] $*\033[0m"
}

log_warn() {
    echo -e "\033[33m[$(date '+%Y-%m-%d %H:%M:%S')] [WARN] $*\033[0m"
}

# Cleanup function
cleanup() {
    log "Performing cleanup operations..."
    # Kill background processes
    if [[ -n "$BACKUP_PID" ]]; then
        kill "$BACKUP_PID" 2>/dev/null || true
    fi
}

trap cleanup EXIT

# Prepare configuration file
prepare_config() {
    USER_HOME="${HOME:-/home/user}"
    local config_file="$USER_HOME/config/restic.conf"
    
    if [[ ! -f "$config_file" ]]; then
        log "Configuration file does not exist, creating from template..."
        mkdir -p "$(dirname "$config_file")"
        
        if [[ -f "$USER_HOME/config/restic.conf" ]]; then
            # Copy project configuration file and adapt paths
            cp "$USER_HOME/config/restic.conf" "$config_file"
            
            # Simple path replacement (adapt to different user environments)
            if [[ "$USER_HOME" != "/home/user" ]]; then
                sed -i "s|/home/user|$USER_HOME|g" "$config_file"
            fi
            
            log_success "Created configuration from template: $USER_HOME/config/restic.conf"
        else
            # Create basic configuration
            cat > "$config_file" << EOF
# Basic restic configuration
RESTIC_ENABLED="false"
RESTIC_REPOSITORY_PATH="$USER_HOME/data/restic-repo"
RESTIC_PASSWORD=""
RESTIC_BACKUP_PATHS="$USER_HOME"
RESTIC_BACKUP_INTERVAL="3600"
RESTIC_MAINTENANCE_INTERVAL="86400"
RESTIC_LOG_FILE="$USER_HOME/log/restic.log"
RESTIC_CACHE_DIR="$USER_HOME/.cache/restic"
RESTIC_BACKUP_ON_STARTUP="true"
RESTIC_AUTO_PRUNE="true"
RESTIC_KEEP_DAILY="7"
RESTIC_KEEP_WEEKLY="4"
RESTIC_KEEP_MONTHLY="6"
EOF
            log_success "Created basic configuration file"
        fi
    else
        log "Configuration file already exists: $config_file"
    fi
}

# Load configuration file (environment variables take precedence)
load_config() {
    local config_file="${1:-${HOME:-/home/user}/config/restic.conf}"
    
    if [[ -f "$config_file" ]]; then
        log "Loading configuration file: $config_file"
        
        # Use safer configuration loading method, remove inline comments
        local temp_config="/tmp/restic_config_$$"
        
        # Process configuration file: remove inline comments, keep valid variable assignments
        grep -E '^[^#]*=' "$config_file" | \
        sed 's/\([^"]*"[^"]*"\)[[:space:]]*#.*/\1/' | \
        sed 's/^\([^#]*\)=\(.*\)$/\1=${\1:-\2}/' > "$temp_config"
        
        source "$temp_config"
        rm -f "$temp_config"
        
        log_success "Configuration loaded (environment variables take precedence)"
    else
        log_warn "Configuration file does not exist: $config_file"
        log "Using environment variables and default values"
    fi
}

# Create necessary directories
ensure_directories() {
    local dirs=(
        "$(dirname "$RESTIC_REPOSITORY_PATH")" 
        "$(dirname "$RESTIC_LOG_FILE")" 
        "$RESTIC_CACHE_DIR"
    )
    
    for dir in "${dirs[@]}"; do
        mkdir -p "$dir" 2>/dev/null || true
    done
}

# Setup password
setup_password() {
    if [[ -z "${RESTIC_PASSWORD:-}" ]]; then
        # Auto-generate password
        RESTIC_PASSWORD="$(openssl rand -base64 32)"
        log_success "Password auto-generated"
    fi
    
    export RESTIC_REPOSITORY="$RESTIC_REPOSITORY_PATH"
    export RESTIC_PASSWORD
    export RESTIC_CACHE_DIR
}

# Check restic command
check_restic() {
    if ! command -v restic &> /dev/null; then
        log_error "Restic command not found"
        return 1
    fi
    
    local version
    version=$(restic version | head -n1)
    log_success "Found: $version"
}

# Initialize repository
init_repository() {
    log "Checking repository..."
    
    if restic snapshots > /dev/null 2>&1; then
        log_success "Repository already exists"
    else
        log "Initializing repository..."
        if restic init; then
            log_success "Repository initialized successfully"
        else
            log_error "Repository initialization failed"
            return 1
        fi
    fi
}

# Find backup script path
find_backup_script() {
    local script_paths=(
        "$(dirname "$0")/../utils/restic-backup.sh"
        "./script/utils/restic-backup.sh"
        "./restic-backup.sh"
    )
    
    for script_path in "${script_paths[@]}"; do
        if [[ -f "$script_path" ]]; then
            echo "$script_path"
            return 0
        fi
    done
    
    log_error "restic-backup.sh script not found"
    return 1
}

# Perform backup (using independent backup script)
perform_backup() {
    local backup_script
    if ! backup_script=$(find_backup_script); then
        log_error "Cannot find backup script, falling back to built-in backup functionality"
        return 1
    fi
    
    log "Using backup script: $backup_script"
    
    # Use independent backup script
    bash "$backup_script" -c "${config_file:-${HOME:-/home/user}/config/restic.conf}"
    return $?
}

# Maintenance tasks
run_maintenance() {
    log "Starting maintenance tasks..."
    
    # Check repository
    log "Checking repository integrity..."
    if restic check; then
        log_success "Repository check passed"
    else
        log_error "Repository check failed"
        return 1
    fi
    
    # Clean old snapshots
    log "Cleaning old snapshots..."
    local args=("forget")
    
    [[ -n "${RESTIC_KEEP_HOURLY:-}" ]] && args+=("--keep-hourly" "${RESTIC_KEEP_HOURLY}")
    [[ -n "${RESTIC_KEEP_DAILY:-}" ]] && args+=("--keep-daily" "${RESTIC_KEEP_DAILY}")
    [[ -n "${RESTIC_KEEP_WEEKLY:-}" ]] && args+=("--keep-weekly" "${RESTIC_KEEP_WEEKLY}")
    [[ -n "${RESTIC_KEEP_MONTHLY:-}" ]] && args+=("--keep-monthly" "${RESTIC_KEEP_MONTHLY}")
    [[ "${RESTIC_AUTO_PRUNE:-true}" == "true" ]] && args+=("--prune")
    
    if restic "${args[@]}"; then
        log_success "Cleanup completed"
    else
        log_warn "Cleanup failed"
    fi
    
    log_success "Maintenance tasks completed"
}

# Main service loop
main_service() {
    log "Starting restic service..."
    
    # Backup on startup
    if [[ "${RESTIC_BACKUP_ON_STARTUP:-true}" == "true" ]]; then
        perform_backup
    fi
    
    # Main loop
    local last_backup=$(date +%s)
    local last_maintenance=$(date +%s)
    
    while true; do
        local current_time=$(date +%s)
        
        # Check if backup is needed
        if [[ $((current_time - last_backup)) -ge "${RESTIC_BACKUP_INTERVAL:-3600}" ]]; then
            if perform_backup; then
                last_backup=$current_time
            fi
        fi
        
        # Check if maintenance is needed
        if [[ $((current_time - last_maintenance)) -ge "${RESTIC_MAINTENANCE_INTERVAL:-86400}" ]]; then
            if run_maintenance; then
                last_maintenance=$current_time
            fi
        fi
        
        # Wait 60 seconds
        sleep 60
    done
}

# Main program
main() {
    log "======= Restic Service Startup ======="
    
    # 1. Prepare configuration file (if it doesn't exist)
    prepare_config
    
    # 2. Load configuration file
    local config_file="${1:-${HOME:-/home/user}/config/restic.conf}"
    load_config "$config_file"
    
    # 3. Create necessary directories
    ensure_directories
    
    # 4. Check if enabled
    if [[ "${RESTIC_ENABLED:-false}" != "true" ]]; then
        log_warn "Restic service is disabled"
        log "To enable: set RESTIC_ENABLED=true in environment variables or configuration file"
        exit 0
    fi
    
    # 5. Check dependencies and setup environment
    check_restic || exit 1
    setup_password
    
    # Display configuration
    log "Backup service configuration:"
    log "  Repository path: $RESTIC_REPOSITORY_PATH"
    log "  Backup paths: $RESTIC_BACKUP_PATHS"
    log "  Backup interval: ${RESTIC_BACKUP_INTERVAL:-3600} seconds"
    log "  Maintenance interval: ${RESTIC_MAINTENANCE_INTERVAL:-86400} seconds"
    log "  Backup on startup: ${RESTIC_BACKUP_ON_STARTUP:-true}"
    
    # Initialize repository
    init_repository || exit 1
    
    # Start service
    main_service
}

# Run main program
main "$@"