#!/bin/bash

# Combined update script for docker-compose-nas
# - Updates .env file from .env.example while preserving values
# - Updates Authelia configuration from example file
# - Configures services with correct paths and API keys
# - Manages Authelia accounts
# - Controls service authentication requirements
# Created: April 26, 2025

set -e

# Color definitions
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[0;33m'
BLUE='\033[0;34m'
MAGENTA='\033[0;35m'
CYAN='\033[0;36m'
BOLD='\033[1m'
NC='\033[0m' # No Color

# Files
ENV_FILE=".env"
ENV_EXAMPLE=".env.example"
TIMESTAMP=$(date +"%Y%m%d-%H%M%S")
ENV_BACKUP=".env.${TIMESTAMP}.bak"
AUTHELIA_CONFIG="authelia/configuration.yml"
AUTHELIA_CONFIG_EXAMPLE="authelia/configuration.example.yml"
AUTHELIA_CONFIG_BACKUP="authelia/configuration.${TIMESTAMP}.bak"
COMPOSE_FILE="docker-compose.yml"
COMPOSE_BACKUP="docker-compose.${TIMESTAMP}.bak"

# Check if yq is installed
check_yq() {
    if ! command -v yq &> /dev/null; then
        echo -e "${YELLOW}Warning: 'yq' is not installed. While not required, it provides better YAML handling.${NC}"
        echo -e "${YELLOW}Installation instructions: https://github.com/mikefarah/yq#install${NC}"
        return 1
    fi
    return 0
}

# Print section header
print_header() {
    echo -e "\n${CYAN}${BOLD}$1${NC}"
    echo -e "${CYAN}$(printf '=%.0s' $(seq 1 ${#1}))${NC}"
}

# Check if a file exists
check_file() {
    if [ ! -f "$1" ]; then
        echo -e "${RED}Error: $1 doesn't exist${NC}"
        return 1
    fi
    return 0
}

# Function to create a backup
create_backup() {
    echo -e "${BLUE}Creating backup of $1 as $2...${NC}"
    cp "$1" "$2"
}


##################################################
# PART 1: Update .env file from .env.example
##################################################

update_env_file() {
    print_header "Environment File Update Tool"
    
    # Check if files exist
    if [ ! -f "$ENV_FILE" ]; then
        echo -e "${RED}Error: $ENV_FILE doesn't exist${NC}"
        echo -e "Creating a new $ENV_FILE from $ENV_EXAMPLE"
        cp "$ENV_EXAMPLE" "$ENV_FILE"
        echo -e "${GREEN}Done! Please review and fill in required values in $ENV_FILE${NC}"
        return 0
    fi

    if ! check_file "$ENV_EXAMPLE"; then
        return 1
    fi

    echo -e "${BLUE}This will update your $ENV_FILE based on the structure in $ENV_EXAMPLE${NC}"
    echo -e "${BLUE}Your existing values will be preserved where possible${NC}"
    echo -e "${BLUE}Backup will be created as: $ENV_BACKUP${NC}"
    echo -e "${YELLOW}Continue? [y/N]:${NC}"
    read -r answer
    if [[ ! "$answer" =~ ^[Yy]$ ]]; then
        echo -e "${RED}Environment update cancelled.${NC}"
        return 0
    fi

    # Create backup of current .env
    create_backup "$ENV_FILE" "$ENV_BACKUP"

    # Store current env values
    echo -e "${BLUE}Reading current environment values...${NC}"
    declare -A current_values
    declare -A current_keys_present
    while IFS='=' read -r key value; do
        # Skip comments and empty lines
        if [[ ! "$key" =~ ^#.*$ ]] && [[ ! -z "$key" ]]; then
            # Clean up any comments after the value
            value=$(echo "$value" | sed 's/[[:space:]]*#.*$//')
            # Trim leading/trailing whitespace
            key=$(echo "$key" | xargs)
            value=$(echo "$value" | xargs)
            # Store in associative array if key is not empty
            if [[ ! -z "$key" ]]; then
                current_values["$key"]="$value"
                # Track that this key existed in original file, regardless of value
                current_keys_present["$key"]=1
            fi
        fi
    done < "$ENV_FILE"

    # Create new env file from example
    echo -e "${BLUE}Creating new $ENV_FILE from $ENV_EXAMPLE...${NC}"
    cp "$ENV_EXAMPLE" "$ENV_FILE.new"

    # Track which keys from the current env have been used
    declare -A used_keys
    # Track new keys that need attention
    new_keys=()
    # Track keys with special warnings
    special_keys=()

    # Process the template and fill in values from current env
    while IFS= read -r line; do
        if [[ "$line" =~ ^([A-Za-z0-9_]+)=(.*)$ ]]; then
            key="${BASH_REMATCH[1]}"
            default_value="${BASH_REMATCH[2]}"
            
            # Mark the key as used if it exists in the original file
            if [[ -n "${current_keys_present[$key]}" ]]; then
                used_keys["$key"]=1
                
                # Replace the line with the current value if one exists
                if [[ -n "${current_values[$key]}" ]]; then
                    sed -i "s|^$key=.*$|$key=${current_values[$key]}|" "$ENV_FILE.new"
                fi
            # If key doesn't exist in original file and has empty/placeholder value
            elif [[ -z "$default_value" ]] || [[ "$default_value" == '""' ]] || [[ "$default_value" == "''" ]]; then
                new_keys+=("$key")
                
                # Special attention for Authelia keys
                if [[ "$key" == AUTHELIA_*_SECRET* ]] || [[ "$key" == AUTHELIA_*_KEY* ]]; then
                    special_keys+=("$key")
                fi
            fi
        fi
    done < "$ENV_FILE.new"

    # Create section for unused/deprecated keys at the bottom of the file
    echo -e "\n\n# --- DEPRECATED OR UNUSED KEYS (Kept for Reference) ---" >> "$ENV_FILE.new"
    echo -e "# Keys below were in your original .env but aren't in the current .env.example" >> "$ENV_FILE.new"
    echo -e "# They may be deprecated or renamed. Review and remove if no longer needed\n" >> "$ENV_FILE.new"

    unused_keys_count=0
    for key in "${!current_values[@]}"; do
        if [[ -z "${used_keys[$key]}" ]]; then
            echo "$key=${current_values[$key]} # DEPRECATED/UNUSED - Review" >> "$ENV_FILE.new"
            unused_keys_count=$((unused_keys_count + 1))
        fi
    done

    # Replace the old file with the new one
    mv "$ENV_FILE.new" "$ENV_FILE"

    # Generate summary 
    echo -e "\n${GREEN}${BOLD}Environment Update Complete!${NC}"
    echo -e "${BLUE}Summary:${NC}"
    echo -e "  - ${CYAN}Original config backed up to: $ENV_BACKUP${NC}"
    echo -e "  - ${CYAN}Updated .env structure to match .env.example${NC}"
    echo -e "  - ${CYAN}Preserved ${#used_keys[@]} existing values${NC}"

    if [[ $unused_keys_count -gt 0 ]]; then
        echo -e "  - ${YELLOW}Found $unused_keys_count deprecated/unused keys${NC}"
        echo -e "    ${YELLOW}These have been moved to the bottom of the file with warnings${NC}"
    fi

    if [[ ${#new_keys[@]} -gt 0 ]]; then
        echo -e "\n${YELLOW}${BOLD}NEW KEYS NEEDING ATTENTION:${NC}"
        echo -e "${YELLOW}The following keys are new and may need values set:${NC}"
        for key in "${new_keys[@]}"; do
            echo -e "  - ${MAGENTA}$key${NC}"
        done
    fi

    if [[ ${#special_keys[@]} -gt 0 ]]; then
        echo -e "\n${RED}${BOLD}IMPORTANT SECURITY KEYS:${NC}"
        echo -e "${RED}The following keys require secure values:${NC}"
        for key in "${special_keys[@]}"; do
            echo -e "  - ${MAGENTA}$key${NC}" 
            
            # Specific advice for Authelia keys
            if [[ "$key" == AUTHELIA_*_SECRET* ]] || [[ "$key" == AUTHELIA_*_KEY* ]]; then
                echo -e "    ${CYAN}Generate with: ${GREEN}openssl rand -hex 32${NC}"
            fi
        done
    fi

    echo -e "\n${BLUE}Review your updated $ENV_FILE file and adjust any values as needed.${NC}"
}


##################################################
# PART 2: Update Authelia configuration
##################################################

update_authelia_config() {
    print_header "Authelia Configuration Update Tool"

    # Check if files exist
    if ! check_file "$AUTHELIA_CONFIG_EXAMPLE"; then
        echo -e "${RED}Error: Example configuration file '$AUTHELIA_CONFIG_EXAMPLE' doesn't exist${NC}"
        return 1
    fi
    if ! check_file "$ENV_FILE"; then
        echo -e "${RED}Error: Environment file '$ENV_FILE' doesn't exist. Cannot retrieve domain settings.${NC}"
        return 1
    fi

    # Get the tailnet domain and hostname from .env
    local TAILNET_DOMAIN=$(grep -oP "^TAILSCALE_TAILNET_DOMAIN=\K.*" "$ENV_FILE" | tr -d '"' | tr -d "'")
    local TAILSCALE_HOSTNAME=$(grep -oP "^TAILSCALE_HOSTNAME=\K.*" "$ENV_FILE" | tr -d '"' | tr -d "'")
    local FULL_HOSTNAME="${TAILSCALE_HOSTNAME}.${TAILNET_DOMAIN}"
    local WILDCARD_DOMAIN="*.${TAILNET_DOMAIN}"

    if [ -z "$TAILNET_DOMAIN" ] || [ -z "$TAILSCALE_HOSTNAME" ]; then
        echo -e "${RED}Error: Could not read TAILSCALE_TAILNET_DOMAIN or TAILSCALE_HOSTNAME from $ENV_FILE${NC}"
        return 1
    fi

    # If config file doesn't exist, create it from example
    if [ ! -f "$AUTHELIA_CONFIG" ]; then
        echo -e "${YELLOW}Authelia configuration file '$AUTHELIA_CONFIG' doesn't exist, creating from example...${NC}"
        cp "$AUTHELIA_CONFIG_EXAMPLE" "$AUTHELIA_CONFIG"
        echo -e "${GREEN}Created new Authelia configuration file.${NC}"
        # Proceed to update the newly created file
    else
        echo -e "${BLUE}This will update your Authelia configuration '$AUTHELIA_CONFIG' based on '$AUTHELIA_CONFIG_EXAMPLE'${NC}"
        echo -e "${BLUE}Your Tailscale domain settings from '$ENV_FILE' will be applied.${NC}"
        echo -e "${BLUE}Backup will be created as: $AUTHELIA_CONFIG_BACKUP${NC}"
        echo -e "${YELLOW}Continue? [y/N]:${NC}"
        read -r answer
        if [[ ! "$answer" =~ ^[Yy]$ ]]; then
            echo -e "${RED}Authelia config update cancelled.${NC}"
            return 0
        fi
        # Create backup of current config only if it exists
        create_backup "$AUTHELIA_CONFIG" "$AUTHELIA_CONFIG_BACKUP"
    fi

    # Check for yq
    if check_yq; then
        echo -e "${BLUE}Using 'yq' to update Authelia configuration...${NC}"
        # Create a temporary file from the example
        local TEMP_CONFIG="${AUTHELIA_CONFIG}.tmp"
        cp "$AUTHELIA_CONFIG_EXAMPLE" "$TEMP_CONFIG"

        # Preserve specific existing values if the original config exists
        if [ -f "$AUTHELIA_CONFIG_BACKUP" ]; then # Use backup as source of truth for existing values
            echo -e "${BLUE}Attempting to preserve existing secrets and notifier settings...${NC}"
            local existing_jwt_secret=$(yq e '.identity_validation.reset_password.jwt_secret // ""' "$AUTHELIA_CONFIG_BACKUP")
            local existing_session_secret=$(yq e '.session.secret // ""' "$AUTHELIA_CONFIG_BACKUP")
            local existing_storage_key=$(yq e '.storage.encryption_key // ""' "$AUTHELIA_CONFIG_BACKUP")
            local existing_redis_pass=$(yq e '.session.redis.password // ""' "$AUTHELIA_CONFIG_BACKUP")
            local existing_notifier=$(yq e '.notifier // ""' "$AUTHELIA_CONFIG_BACKUP")

            # Update secrets in temp file if they existed in the backup
            if [[ -n "$existing_jwt_secret" && "$existing_jwt_secret" != '""' && "$existing_jwt_secret" != "null" ]]; then
                yq e -i '.identity_validation.reset_password.jwt_secret = strenv(existing_jwt_secret)' --env existing_jwt_secret="$existing_jwt_secret" "$TEMP_CONFIG"
            fi
            if [[ -n "$existing_session_secret" && "$existing_session_secret" != '""' && "$existing_session_secret" != "null" ]]; then
                yq e -i '.session.secret = strenv(existing_session_secret)' --env existing_session_secret="$existing_session_secret" "$TEMP_CONFIG"
            fi
            if [[ -n "$existing_storage_key" && "$existing_storage_key" != '""' && "$existing_storage_key" != "null" ]]; then
                yq e -i '.storage.encryption_key = strenv(existing_storage_key)' --env existing_storage_key="$existing_storage_key" "$TEMP_CONFIG"
            fi
            if [[ -n "$existing_redis_pass" && "$existing_redis_pass" != '""' && "$existing_redis_pass" != "null" ]]; then
                yq e -i '.session.redis.password = strenv(existing_redis_pass)' --env existing_redis_pass="$existing_redis_pass" "$TEMP_CONFIG"
            fi
            if [[ -n "$existing_notifier" && "$existing_notifier" != '""' && "$existing_notifier" != "null" ]]; then
                yq e -i '.notifier = strenv(existing_notifier)' --env existing_notifier="$existing_notifier" "$TEMP_CONFIG"
            fi
        fi

        # Update domain settings from .env
        echo -e "${BLUE}Applying Tailscale domain settings...${NC}"

        if command -v yq &> /dev/null; then
            # Replace placeholder full hostname
            yq e -i "walk(if type == \"string\" and . == \"tailscale-nas.your-tailnet.ts.net\" then \"${FULL_HOSTNAME}\" else . end)" "$TEMP_CONFIG"

            # Replace placeholder tailnet domain
            yq e -i "walk(if type == \"string\" and . == \"your-tailnet.ts.net\" then \"${TAILNET_DOMAIN}\" else . end)" "$TEMP_CONFIG"

            # Move temp file to final location
            mv "$TEMP_CONFIG" "$AUTHELIA_CONFIG"

            echo -e "${GREEN}Authelia configuration updated successfully!${NC}"
        else
            echo -e "${YELLOW}Warning: 'yq' is not installed. Using sed to update configuration.${NC}"
            echo -e "${YELLOW}This is less reliable and may not preserve all settings.${NC}"

            # Create a new file from the example
            cp "$AUTHELIA_CONFIG_EXAMPLE" "$AUTHELIA_CONFIG.new"

            # Replace placeholders manually
            sed -i "s|tailscale-nas\.your-tailnet\.ts\.net|${FULL_HOSTNAME}|g" "$AUTHELIA_CONFIG.new"
            sed -i "s|your-tailnet\.ts\.net|${TAILNET_DOMAIN}|g" "$AUTHELIA_CONFIG.new"

            # Move the new file to the final location
            mv "$AUTHELIA_CONFIG.new" "$AUTHELIA_CONFIG"

            echo -e "${YELLOW}Authelia configuration updated with sed.${NC}"
    fi
    
    echo -e "${GREEN}Authelia configuration update completed.${NC}"
    echo -e "${BLUE}Remember to restart Authelia for changes to take effect:${NC}"
    echo -e "${CYAN}  docker compose restart authelia${NC}"
}


##################################################
# PART 3: Update service configurations
##################################################

update_service_configs() {
    print_header "Service Configuration Update Tool"
    
    # Check if Docker is running
    if ! docker ps &> /dev/null; then
        echo -e "${RED}Error: Docker is not running or you don't have permission to use it.${NC}"
        echo -e "${YELLOW}Please start Docker and ensure your user has permission to run Docker commands.${NC}"
        return 1
    fi
    
    if ! check_file "$ENV_FILE"; then
        echo -e "${RED}Error: Environment file '$ENV_FILE' doesn't exist. Cannot update services.${NC}"
        return 1
    fi
    
    # Get CONFIG_ROOT from .env if defined
    CONFIG_ROOT=$(grep -oP "^CONFIG_ROOT=\K.*" "$ENV_FILE" | tr -d '"' | tr -d "'" || echo ".")
    
    echo -e "${BLUE}This will update configurations for running containers and extract API keys.${NC}"
    echo -e "${BLUE}Services may be restarted during this process.${NC}"
    echo -e "${YELLOW}Continue? [y/N]:${NC}"
    read -r answer
    if [[ ! "$answer" =~ ^[Yy]$ ]]; then
        echo -e "${RED}Service configuration update cancelled.${NC}"
        return 0
    fi
    
    # Create a temporary .env file for updates
    local ENV_TEMP="${ENV_FILE}.${TIMESTAMP}.tmp"
    cp "$ENV_FILE" "$ENV_TEMP"
    
    # Service update tracker
    declare -A updated_services
    declare -A api_keys_updated
    
    # Function to restart a container if needed
    restart_container() {
        local service=$1
        if [ "${updated_services[$service]}" == "1" ]; then
            echo -e "${YELLOW}Restarting $service to apply changes...${NC}"
            if docker compose restart "$service"; then
                echo -e "${GREEN}$service restarted successfully!${NC}"
            else
                echo -e "${RED}Failed to restart $service. Please restart manually.${NC}"
            fi
        fi
    }
    
    # Function to extract API key from a container
    extract_api_key() {
        local service=$1
        local config_path=$2
        local api_key_pattern=$3
        local env_var_name=$4
        
        # Skip if service is not running
        if ! docker compose ps "$service" 2>/dev/null | grep -q "Up"; then
            echo -e "${YELLOW}Service $service is not running, skipping API key extraction.${NC}"
            return 0
        fi
        
        echo -e "${BLUE}Attempting to extract API key for $service...${NC}"
        
        # Try to get API key directly from the config file in the container
        local api_key=""
        if [ -n "$config_path" ] && [ -n "$api_key_pattern" ]; then
            api_key=$(docker compose exec -T "$service" grep -oP "$api_key_pattern" "$config_path" 2>/dev/null | head -1 | sed 's/.*[":=]\s*\([^"]*\).*/\1/g')
        fi
        
        # If API key was found and is not empty
        if [ -n "$api_key" ]; then
            echo -e "${GREEN}API key for $service extracted successfully!${NC}"
            
            # Check if the API key is already in the .env file
            if grep -q "^$env_var_name=" "$ENV_TEMP"; then
                local current_key=$(grep -oP "^$env_var_name=\K.*" "$ENV_TEMP" | tr -d '"' | tr -d "'")
                
                # Only update if key is different and not empty
                if [ "$current_key" != "$api_key" ] && [ -n "$current_key" ]; then
                    echo -e "${YELLOW}API key for $service in .env differs from extracted key.${NC}"
                    echo -e "${YELLOW}Would you like to update it? [y/N]:${NC}"
                    read -r update_key
                    if [[ ! "$update_key" =~ ^[Yy]$ ]]; then
                        echo -e "${BLUE}Keeping existing API key in .env file.${NC}"
                        return 0
                    fi
                fi
                
                # Update the existing entry
                sed -i "s|^$env_var_name=.*|$env_var_name=$api_key|" "$ENV_TEMP"
            else
                # Add the new API key to the .env file
                echo -e "\n# API key for $service (extracted $(date +"%Y-%m-%d"))" >> "$ENV_TEMP"
                echo "$env_var_name=$api_key" >> "$ENV_TEMP"
            fi
            
            api_keys_updated["$service"]="1"
        else
            echo -e "${YELLOW}Could not extract API key for $service.${NC}"
        fi
    }
    
    # Update *arr services with the correct URL base
    for service in "sonarr" "radarr" "lidarr" "prowlarr" "bazarr"; do
        # Skip if service is not running
        if ! docker compose ps "$service" 2>/dev/null | grep -q "Up"; then
            echo -e "${YELLOW}Service $service is not running, skipping.${NC}"
            continue
        fi
        
        echo -e "${BLUE}Updating $service configuration...${NC}"
        
        local config_file="/config/config.xml"
        
        # Check if the service's config file exists
        if ! docker compose exec "$service" test -f "$config_file"; then
            echo -e "${YELLOW}Configuration file for $service not found, skipping.${NC}"
            continue
        fi
        
        # Check current URLBase setting
        local current_urlbase=$(docker compose exec -T "$service" grep -oP '<UrlBase>\K[^<]+' "$config_file" 2>/dev/null || echo "")
        local desired_urlbase="/$service"
        
        # Update URLBase if needed
        if [ "$current_urlbase" != "$desired_urlbase" ]; then
            echo -e "${YELLOW}$service URL base needs to be updated from '$current_urlbase' to '$desired_urlbase'.${NC}"
            
            if [ -z "$current_urlbase" ]; then
                # If URLBase tag doesn't exist, add it
                docker compose exec "$service" sed -i "/<BaseUrl>/i\\  <UrlBase>$desired_urlbase</UrlBase>" "$config_file" || \
                docker compose exec "$service" sed -i "/<Port>/i\\  <UrlBase>$desired_urlbase</UrlBase>" "$config_file"
            else
                # If URLBase exists, update it
                docker compose exec "$service" sed -i "s|<UrlBase>[^<]*</UrlBase>|<UrlBase>$desired_urlbase</UrlBase>|g" "$config_file"
            fi
            
            if [ $? -eq 0 ]; then
                echo -e "${GREEN}$service URL base updated to $desired_urlbase${NC}"
                updated_services["$service"]="1"
            else
                echo -e "${RED}Failed to update $service URL base.${NC}"
            fi
        else
            echo -e "${GREEN}$service URL base already set to $desired_urlbase${NC}"
        fi
        
        # Extract API key based on service
        case "$service" in
            "sonarr")
                extract_api_key "$service" "$config_file" '<ApiKey>\K[^<]+' "SONARR_API_KEY"
                ;;
            "radarr")
                extract_api_key "$service" "$config_file" '<ApiKey>\K[^<]+' "RADARR_API_KEY"
                ;;
            "lidarr")
                extract_api_key "$service" "$config_file" '<ApiKey>\K[^<]+' "LIDARR_API_KEY"
                ;;
            "prowlarr")
                extract_api_key "$service" "$config_file" '<ApiKey>\K[^<]+' "PROWLARR_API_KEY"
                ;;
            "bazarr")
                # Bazarr uses a different config format (YAML/settings.ini)
                extract_api_key "$service" "/config/config/config.ini" 'auth_apikey\s*=\s*\K[^#\n]+' "BAZARR_API_KEY"
                ;;
        esac
    done
    
    # Update qBittorrent WebUI URL
    if docker compose ps "qbittorrent" 2>/dev/null | grep -q "Up"; then
        echo -e "${BLUE}Updating qBittorrent configuration...${NC}"
        
        local qbit_config_dir="/config/qBittorrent"
        local qbit_config_file="$qbit_config_dir/qBittorrent.conf"
        
        # Check if the config file exists
        if docker compose exec "qbittorrent" test -f "$qbit_config_file"; then
            # Check current WebUI\RootFolder setting
            local current_root=$(docker compose exec -T "qbittorrent" grep -oP '^WebUI\\RootFolder\s*=\s*\K.*' "$qbit_config_file" 2>/dev/null || echo "")
            local desired_root="/qbittorrent"
            
            # Update WebUI\RootFolder if needed
            if [ "$current_root" != "$desired_root" ]; then
                echo -e "${YELLOW}qBittorrent WebUI root needs to be updated from '$current_root' to '$desired_root'.${NC}"
                
                if docker compose exec "qbittorrent" grep -q "^WebUI\\\\RootFolder=" "$qbit_config_file"; then
                    # If setting exists, update it
                    docker compose exec "qbittorrent" sed -i "s|^WebUI\\\\RootFolder=.*|WebUI\\\\RootFolder=$desired_root|g" "$qbit_config_file"
                else
                    # If setting doesn't exist, add it to the [Preferences] section
                    if docker compose exec "qbittorrent" grep -q "^\[Preferences\]" "$qbit_config_file"; then
                        docker compose exec "qbittorrent" sed -i "/^\[Preferences\]/a WebUI\\\\RootFolder=$desired_root" "$qbit_config_file"
                    else
                        # Add [Preferences] section if it doesn't exist
                        docker compose exec "qbittorrent" bash -c "echo -e '\n[Preferences]\nWebUI\\\\RootFolder=$desired_root' >> $qbit_config_file"
                    fi
                fi
                
                if [ $? -eq 0 ]; then
                    echo -e "${GREEN}qBittorrent WebUI root updated to $desired_root${NC}"
                    updated_services["qbittorrent"]="1"
                else
                    echo -e "${RED}Failed to update qBittorrent WebUI root.${NC}"
                fi
            else
                echo -e "${GREEN}qBittorrent WebUI root already set to $desired_root${NC}"
            fi
        else
            echo -e "${YELLOW}qBittorrent configuration file not found, skipping.${NC}"
        fi
    fi
    
    # Extract Jellyfin API key
    if docker compose ps "jellyfin" 2>/dev/null | grep -q "Up"; then
        # Jellyfin API key is in a different location
        extract_api_key "jellyfin" "/config/config/system.xml" '<ApiKey>\K[^<]+' "JELLYFIN_API_KEY"
    fi
    
    # Extract Jellyseerr API key
    if docker compose ps "jellyseerr" 2>/dev/null | grep -q "Up"; then
        # Jellyseerr stores API key in settings.json
        extract_api_key "jellyseerr" "/app/config/settings.json" '"apiKey":\s*"\K[^"]+' "JELLYSEERR_API_KEY"
    fi
    
    # Check for optional services based on COMPOSE_PROFILES
    local profiles=$(grep -oP "^COMPOSE_PROFILES=\K.*" "$ENV_FILE" | tr -d '"' | tr -d "'" || echo "")
    
    # Process profiles as a comma-separated list
    IFS=',' read -ra profile_array <<< "$profiles"
    for profile in "${profile_array[@]}"; do
        profile=$(echo "$profile" | xargs) # Trim whitespace
        
        # Handle specific profiles
        case "$profile" in
            "sabnzbd")
                if docker compose ps "sabnzbd" 2>/dev/null | grep -q "Up"; then
                    # Extract SABnzbd API key from its config
                    extract_api_key "sabnzbd" "/config/sabnzbd.ini" 'api_key\s*=\s*\K[0-9a-f]+' "SABNZBD_API_KEY"
                fi
                ;;
            "immich")
                if docker compose ps "immich-server" 2>/dev/null | grep -q "Up"; then
                    # Extract Immich API key from user database
                    echo -e "${BLUE}Attempting to extract Immich API key...${NC}"
                    local immich_api_key=$(docker compose exec -T immich-server /usr/local/bin/cli api-key --show 2>/dev/null | grep -oP '[0-9a-f]{8}(-[0-9a-f]{4}){3}-[0-9a-f]{12}')
                    
                    if [ -n "$immich_api_key" ]; then
                        echo -e "${GREEN}Immich API key extracted successfully!${NC}"
                        
                        # Update or add the API key to the .env file
                        if grep -q "^IMMICH_API_KEY=" "$ENV_TEMP"; then
                            sed -i "s|^IMMICH_API_KEY=.*|IMMICH_API_KEY=$immich_api_key|" "$ENV_TEMP"
                        else
                            echo -e "\n# API key for Immich (extracted $(date +"%Y-%m-%d"))" >> "$ENV_TEMP"
                            echo "IMMICH_API_KEY=$immich_api_key" >> "$ENV_TEMP"
                        fi
                        
                        api_keys_updated["immich"]="1"
                    else
                        echo -e "${YELLOW}Could not extract Immich API key.${NC}"
                    fi
                fi
                ;;
            "homeassistant")
                if docker compose ps "homeassistant" 2>/dev/null | grep -q "Up"; then
                    echo -e "${BLUE}Home Assistant detected.${NC}"
                    echo -e "${YELLOW}Home Assistant requires a Long-Lived Access Token for API access.${NC}"
                    echo -e "${YELLOW}This must be generated manually in the Home Assistant UI:${NC}"
                    echo -e "${CYAN}  Profile > Long-Lived Access Tokens > Create Token${NC}"
                    
                    # Check if token already exists
                    if ! grep -q "^HOMEASSISTANT_ACCESS_TOKEN=" "$ENV_TEMP" || [ -z "$(grep -oP "^HOMEASSISTANT_ACCESS_TOKEN=\K.*" "$ENV_TEMP" | tr -d '"' | tr -d "'")" ]; then
                        echo -e "${YELLOW}Would you like to add a Home Assistant token now? [y/N]:${NC}"
                        read -r add_token
                        if [[ "$add_token" =~ ^[Yy]$ ]]; then
                            echo -e "${BLUE}Enter your Home Assistant Long-Lived Access Token:${NC}"
                            read -rs token
                            if [ -n "$token" ]; then
                                if grep -q "^HOMEASSISTANT_ACCESS_TOKEN=" "$ENV_TEMP"; then
                                    sed -i "s|^HOMEASSISTANT_ACCESS_TOKEN=.*|HOMEASSISTANT_ACCESS_TOKEN=$token|" "$ENV_TEMP"
                                else
                                    echo -e "\n# Home Assistant Access Token (added $(date +"%Y-%m-%d"))" >> "$ENV_TEMP"
                                    echo "HOMEASSISTANT_ACCESS_TOKEN=$token" >> "$ENV_TEMP"
                                fi
                                api_keys_updated["homeassistant"]="1"
                                echo -e "${GREEN}Home Assistant token added!${NC}"
                            else
                                echo -e "${RED}No token provided. Skipping.${NC}"
                            fi
                        fi
                    else
                        echo -e "${GREEN}Home Assistant token already exists in .env${NC}"
                    fi
                fi
                ;;
        esac
    done
    
    # Apply changes to the .env file if any API keys were updated
    if [ ${#api_keys_updated[@]} -gt 0 ]; then
        echo -e "${YELLOW}API keys were extracted for these services:${NC}"
        for service in "${!api_keys_updated[@]}"; do
            echo -e "  - ${CYAN}$service${NC}"
        done
        
        echo -e "${YELLOW}Update .env file with these API keys? [Y/n]:${NC}"
        read -r update_env
        if [[ ! "$update_env" =~ ^[Nn]$ ]]; then
            # Backup the current .env before replacing
            if [ ! -f "$ENV_BACKUP" ]; then
                create_backup "$ENV_FILE" "$ENV_BACKUP"
            fi
            
            # Replace the .env file with the updated temp file
            mv "$ENV_TEMP" "$ENV_FILE"
            echo -e "${GREEN}API keys updated in $ENV_FILE${NC}"
        else
            echo -e "${YELLOW}API key updates were not saved to .env file.${NC}"
            rm "$ENV_TEMP"
        fi
    else
        echo -e "${BLUE}No API key updates were found.${NC}"
        rm "$ENV_TEMP"
    fi
    
    # Restart services that were updated
    if [ ${#updated_services[@]} -gt 0 ]; then
        echo -e "${YELLOW}The following services had their configurations updated:${NC}"
        for service in "${!updated_services[@]}"; do
            echo -e "  - ${CYAN}$service${NC}"
        done
        
        echo -e "${YELLOW}Would you like to restart these services now? [Y/n]:${NC}"
        read -r restart_now
        if [[ ! "$restart_now" =~ ^[Nn]$ ]]; then
            for service in "${!updated_services[@]}"; do
                restart_container "$service"
            done
        else
            echo -e "${YELLOW}Remember to restart these services manually for changes to take effect:${NC}"
            for service in "${!updated_services[@]}"; do
                echo -e "  ${CYAN}docker compose restart $service${NC}"
            done
        fi
    fi
    
    echo -e "\n${GREEN}${BOLD}Service Configuration Update Complete!${NC}"
    if [ ${#updated_services[@]} -gt 0 ] || [ ${#api_keys_updated[@]} -gt 0 ]; then
        echo -e "${GREEN}Services were successfully updated.${NC}"
    else
        echo -e "${BLUE}No changes were needed.${NC}"
    fi
}


##################################################
# PART 4: Authelia Policy Management
##################################################

# Get the current policy for a service from Authelia config
# Usage: get_authelia_policy <service_name> <config_file>
# Output: policy (e.g., one_factor, bypass), not_found, or error
get_authelia_policy() {
    local service=$1
    local config_file=$2

    if ! check_file "$config_file"; then
        echo "error: config file not found"
        return 1
    fi

    # Use yq if available for more reliable YAML parsing
    if check_yq; then
        # Find the rule matching the path_regex for the service
        # Note: This assumes path_regex is unique enough for the service (e.g., ^/service.*)
        local policy=$(yq e ".access_control.rules[] | select(.path_regex == \"^/${service}.*\") | .policy" "$config_file" 2>/dev/null)
        if [ -n "$policy" ] && [ "$policy" != "null" ]; then
            echo "$policy"
        else
            # Check if a rule exists for the service path at all
            local rule_exists=$(yq e ".access_control.rules[] | select(.path_regex == \"^/${service}.*\")" "$config_file" 2>/dev/null)
             if [ -n "$rule_exists" ] && [ "$rule_exists" != "null" ]; then
                 echo "error: policy not found in rule" # Rule exists but policy couldn't be read
             else
                 echo "not_found" # No rule found for this service path
             fi
        fi
    else
        # Fall back to grep if yq isn't available (less reliable)
        # This is fragile and depends heavily on formatting
        local policy_line=$(grep -A 1 "path_regex: '^/${service}.*'" "$config_file" | grep "policy:")
        if [ -n "$policy_line" ]; then
            echo "$policy_line" | awk '{print $2}'
        else
             echo "not_found"
        fi
    fi
    return 0
}

# Set the policy for a service in Authelia config
# Usage: set_authelia_policy <service_name> <policy> <config_file>
set_authelia_policy() {
    local service=$1
    local policy=$2
    local config_file=$3
    local backup_file="${config_file}.${TIMESTAMP}.bak"

    if ! check_file "$config_file"; then
        echo -e "${RED}Error: Authelia configuration file '$config_file' not found.${NC}"
        return 1
    fi

    if [[ "$policy" != "one_factor" && "$policy" != "two_factor" && "$policy" != "deny" && "$policy" != "bypass" ]]; then
        echo -e "${RED}Error: Invalid policy '$policy'. Must be one of: one_factor, two_factor, deny, bypass.${NC}"
        return 1
    fi

    echo -e "${BLUE}Setting policy for service '$service' to '$policy' in '$config_file'...${NC}"

    # Create backup if it doesn't exist for this run
    if [ ! -f "$backup_file" ]; then
        create_backup "$config_file" "$backup_file"
    fi

    # Use yq if available
    if check_yq; then
        # Check if the rule exists first using the path_regex
        local rule_index=$(yq e ".access_control.rules | map(.path_regex == \"^/${service}.*\") | indexOf(true)" "$config_file" 2>/dev/null)

        if [ "$rule_index" == "-1" ] || [ -z "$rule_index" ]; then
            echo -e "${YELLOW}Warning: No rule found for service path '^/${service}.*' in '$config_file'.${NC}"
            echo -e "${YELLOW}Attempting to add a new rule...${NC}"

            # Get the tailnet domain from .env for the new rule
            local TAILNET_DOMAIN=$(grep -oP "^TAILSCALE_TAILNET_DOMAIN=\K.*" "$ENV_FILE" | tr -d '"' | tr -d "'")
            if [ -z "$TAILNET_DOMAIN" ]; then
                 echo -e "${RED}Error: Could not read TAILSCALE_TAILNET_DOMAIN from $ENV_FILE. Cannot add rule.${NC}"
                 return 1
            fi
            local WILDCARD_DOMAIN="*.${TAILNET_DOMAIN}"

            # Add the new rule to the access_control.rules array
            # Places it before the generic domain rule if it exists, otherwise at the end
            # This assumes a generic rule like "- domain: '*.domain.tld'" exists near the end
            yq e -i ".access_control.rules |= select(.domain != \"${WILDCARD_DOMAIN}\" or .path_regex != null) + [{\"domain\": \"${WILDCARD_DOMAIN}\", \"path_regex\": \"^/${service}.*\", \"policy\": \"${policy}\"}] + select(.domain == \"${WILDCARD_DOMAIN}\" and .path_regex == null)" "$config_file"

            if [ $? -eq 0 ]; then
                 echo -e "${GREEN}Added new rule for '$service' with policy '$policy'.${NC}"
                 return 0
            else
                 echo -e "${RED}Error: Failed to add new rule for '$service' using yq.${NC}"
                 return 1
            fi
        else
            # Rule exists, update the policy at the found index
            yq e -i "(.access_control.rules[$rule_index].policy) = \"$policy\"" "$config_file"
             if [ $? -eq 0 ]; then
                 echo -e "${GREEN}Policy for '$service' updated to '$policy'.${NC}"
                 return 0
             else
                 echo -e "${RED}Error: Failed to update policy for '$service' using yq.${NC}"
                 return 1
             fi
        fi
    else
        # Fallback to sed (much less reliable, especially for adding rules)
        echo -e "${YELLOW}Warning: 'yq' not found. Using 'sed' which is less reliable for YAML manipulation.${NC}"
        # Check if rule exists (simple grep)
        if grep -q "path_regex: '^/${service}.*'" "$config_file"; then
            # Attempt to find the line number of path_regex and update the policy line below it
            local line_num=$(grep -n "path_regex: '^/${service}.*'" "$config_file" | head -n 1 | cut -d: -f1) # Use head -n 1 just in case
            if [ -n "$line_num" ]; then
                # Assuming policy is the next line (fragile!)
                local policy_line_num=$((line_num + 1))
                # Check if the next line actually contains 'policy:'
                if sed -n "${policy_line_num}p" "$config_file" | grep -q "policy:"; then
                    sed -i "${policy_line_num}s/policy:.*/policy: $policy/" "$config_file"
                    if [ $? -eq 0 ]; then
                        echo -e "${GREEN}Policy for '$service' updated to '$policy' (using sed).${NC}"
                        return 0
                    else
                        echo -e "${RED}Error: Failed to update policy line using sed.${NC}"
                        return 1
                    fi
                else
                    echo -e "${RED}Error: Could not reliably find policy line for '$service' using sed (expected on line $policy_line_num).${NC}"
                    return 1
                fi
            else
                 echo -e "${RED}Error: Could not find line number for service '$service' using sed.${NC}"
                 return 1
            fi
        else
            echo -e "${RED}Error: Rule for service '$service' not found. Cannot add rule using sed.${NC}"
            return 1
        fi
    fi
}

# List services and their Authelia policy status
list_authelia_services() {
    print_header "Authelia Service Policy Status"

    if ! check_file "$AUTHELIA_CONFIG"; then
        echo -e "${RED}Error: Authelia configuration file '$AUTHELIA_CONFIG' not found.${NC}"
        return 1
    fi

    echo -e "${BLUE}Checking service policies in $AUTHELIA_CONFIG...${NC}"
    echo -e "${CYAN}SERVICE\t\tPOLICY${NC}"
    echo -e "${CYAN}-------\t\t------${NC}"

    local service_count=0
    local processed_services="" # Track processed services

    # Use yq to get all rules with path_regex if available
    if check_yq; then
        # Extract path_regex and policy for rules that have path_regex
        local rules_data=$(yq e '.access_control.rules[] | select(has("path_regex")) | {"path": .path_regex, "policy": .policy}' "$AUTHELIA_CONFIG")

        # Process each rule found
        while IFS= read -r line; do
            # Extract service name from path_regex (e.g., "^/sonarr.*" -> "sonarr")
            local path_regex=$(echo "$line" | grep -oP 'path: \K.*' | tr -d '"' | tr -d "'")
            local policy=$(echo "$line" | grep -oP 'policy: \K.*' | tr -d '"' | tr -d "'")

            if [[ "$path_regex" =~ ^\^/([a-zA-Z0-9_-]+)\.\* ]]; then
                local service="${BASH_REMATCH[1]}"

                # Skip duplicates if multiple rules somehow match the same pattern start
                if [[ "$processed_services" == *"$service"* ]]; then
                    continue
                fi
                processed_services="$processed_services $service"

                printf "${BOLD}%-20s${NC}" "$service"
                case "$policy" in
                    "one_factor"|"two_factor")
                        echo -e "${GREEN}${policy}${NC}"
                        ;;
                    "bypass")
                        echo -e "${YELLOW}${policy}${NC}"
                        ;;
                    "deny")
                        echo -e "${RED}${policy}${NC}"
                        ;;
                    *)
                        echo -e "${MAGENTA}Unknown ($policy)${NC}"
                        ;;
                esac
                service_count=$((service_count + 1))
            fi
        # Use yq -N to prevent splitting lines with spaces
        done <<< "$(echo "$rules_data" | yq -N e '.' -)"

    else
        # Fallback to grep (less reliable)
        echo -e "${YELLOW}Warning: yq not found, using grep (may miss services or show duplicates).${NC}"
        # Find lines with path_regex, then try to get the policy line after
        grep -n "path_regex: '^/.*" "$AUTHELIA_CONFIG" | while IFS=: read -r line_num line_content; do
            if [[ "$line_content" =~ path_regex:\ \'^\/([a-zA-Z0-9_-]+)\.\*\' ]]; then
                 local service="${BASH_REMATCH[1]}"
                 # Skip duplicates
                 if [[ "$processed_services" == *"$service"* ]]; then
                     continue
                 fi
                 processed_services="$processed_services $service"

                 # Try to get policy from the next line (very fragile)
                 local policy_line_num=$((line_num + 1))
                 local policy=$(sed -n "${policy_line_num}p" "$AUTHELIA_CONFIG" | grep "policy:" | awk '{print $2}')

                 printf "${BOLD}%-20s${NC}" "$service"
                 if [ -n "$policy" ]; then
                     case "$policy" in
                         "one_factor"|"two_factor") echo -e "${GREEN}${policy}${NC}" ;;
                         "bypass") echo -e "${YELLOW}${policy}${NC}" ;;
                         "deny") echo -e "${RED}${policy}${NC}" ;;
                         *) echo -e "${MAGENTA}Unknown ($policy)${NC}" ;;
                     esac
                 else
                     echo -e "${RED}Unknown (could not read policy)${NC}"
                 fi
                 service_count=$((service_count + 1))
             fi
         done
    fi

    if [ $service_count -eq 0 ]; then
        echo -e "${YELLOW}No services with path_regex rules found in $AUTHELIA_CONFIG.${NC}"
        echo -e "${YELLOW}Only services explicitly defined with path_regex rules are listed.${NC}"
    fi
    return 0
}


cleanup_backups() {
    print_header "Backup Files Cleanup"
    
    echo -e "${BLUE}Searching for backup files...${NC}"
    
    local env_backups=$(find . -maxdepth 1 -name ".env.*.bak" | sort)
    local compose_backups=$(find . -maxdepth 1 -name "docker-compose.*.bak" | sort)
    local authelia_backups=$(find ./authelia -maxdepth 1 -name "configuration.*.bak" 2>/dev/null | sort)
    
    local env_count=$(echo "$env_backups" | grep -c "^")
    local compose_count=$(echo "$compose_backups" | grep -c "^")
    local authelia_count=$(echo "$authelia_backups" | grep -c "^")
    
    echo -e "${CYAN}Found:${NC}"
    echo -e "  - ${CYAN}$env_count .env backup files${NC}"
    echo -e "  - ${CYAN}$compose_count docker-compose.yml backup files${NC}"
    echo -e "  - ${CYAN}$authelia_count Authelia configuration backup files${NC}"
    
    local total_count=$((env_count + compose_count + authelia_count))
    
    if [ $total_count -eq 0 ]; then
        echo -e "${GREEN}No backup files found.${NC}"
        return 0
    fi
    
    echo -e "${YELLOW}Do you want to delete all backup files? [y/N]:${NC}"
    read -r answer
    if [[ ! "$answer" =~ ^[Yy]$ ]]; then
        echo -e "${YELLOW}Would you like to keep the most recent backup of each type? [Y/n]:${NC}"
        read -r keep_recent
        
        if [[ "$keep_recent" =~ ^[Nn]$ ]]; then
            echo -e "${RED}Backup cleanup cancelled.${NC}"
            return 0
        fi
        
        if [ $env_count -gt 1 ]; then
            local keep_env=$(echo "$env_backups" | tail -1)
            local del_env=$(echo "$env_backups" | grep -v "$keep_env")
            echo -e "${BLUE}Keeping most recent .env backup: ${CYAN}$keep_env${NC}"
            echo -e "${BLUE}Deleting ${CYAN}$((env_count - 1))${BLUE} older .env backups${NC}"
            for file in $del_env; do
                rm "$file"
            done
        fi
        
        if [ $compose_count -gt 1 ]; then
            local keep_compose=$(echo "$compose_backups" | tail -1)
            local del_compose=$(echo "$compose_backups" | grep -v "$keep_compose")
            echo -e "${BLUE}Keeping most recent docker-compose.yml backup: ${CYAN}$keep_compose${NC}"
            echo -e "${BLUE}Deleting ${CYAN}$((compose_count - 1))${BLUE} older docker-compose.yml backups${NC}"
            for file in $del_compose; do
                rm "$file"
            done
        fi
        
        if [ $authelia_count -gt 1 ]; then
            local keep_authelia=$(echo "$authelia_backups" | tail -1)
            local del_authelia=$(echo "$authelia_backups" | grep -v "$keep_authelia")
            echo -e "${BLUE}Keeping most recent Authelia config backup: ${CYAN}$keep_authelia${NC}"
            echo -e "${BLUE}Deleting ${CYAN}$((authelia_count - 1))${BLUE} older Authelia config backups${NC}"
            for file in $del_authelia; do
                rm "$file"
            done
        fi
        
        echo -e "${GREEN}Cleanup completed with most recent backups retained.${NC}"
    else
        for file in $env_backups $compose_backups $authelia_backups; do
            rm "$file"
        done
        echo -e "${GREEN}All $total_count backup files have been deleted.${NC}"
    fi
}

# Interactive menu for managing Authelia policies
manage_authelia_policies() {
    print_header "Authelia Policy Management"

    if ! check_file "$AUTHELIA_CONFIG"; then
        echo -e "${RED}Error: Authelia configuration file '$AUTHELIA_CONFIG' not found.${NC}"
        return 1
    fi

    while true; do
        echo -e "\n${BLUE}Choose an option:${NC}"
        echo -e "  ${CYAN}1. ${NC}List services and their current Authelia policy"
        echo -e "  ${CYAN}2. ${NC}Set Authelia policy for a service"
        echo -e "  ${CYAN}3. ${NC}Return to main menu"
        echo

        local choice
        echo -e "${YELLOW}Enter your choice [1-3]: ${NC}"
        read -r choice

        case "$choice" in
            1)
                list_authelia_services
                ;;
            2)
                echo -e "${BLUE}Enter the service name to set the policy for (e.g., sonarr, radarr):${NC}"
                read -r service
                if [ -z "$service" ]; then
                    echo -e "${RED}No service name provided.${NC}"
                    continue
                fi

                echo -e "${BLUE}Select the desired policy for '$service':${NC}"
                echo -e "  ${CYAN}1. ${GREEN}one_factor${NC} (Requires login)"
                echo -e "  ${CYAN}2. ${GREEN}two_factor${NC} (Requires login + 2FA)"
                echo -e "  ${CYAN}3. ${YELLOW}bypass${NC}     (No login required)"
                echo -e "  ${CYAN}4. ${RED}deny${NC}       (Access denied)"
                echo -e "  ${CYAN}5. ${NC}Cancel"

                local policy_choice
                local policy=""
                while true; do
                    echo -e "${YELLOW}Enter policy choice [1-5]: ${NC}"
                    read -r policy_choice
                    case "$policy_choice" in
                        1) policy="one_factor"; break ;;
                        2) policy="two_factor"; break ;;
                        3) policy="bypass"; break ;;
                        4) policy="deny"; break ;;
                        5) policy=""; break ;; # Cancel
                        *) echo -e "${RED}Invalid choice.${NC}" ;;
                    esac
                done

                if [ -z "$policy" ]; then
                    echo -e "${YELLOW}Policy change cancelled.${NC}"
                    continue
                fi

                # Call the function to set the policy
                set_authelia_policy "$service" "$policy" "$AUTHELIA_CONFIG"

                # Check the return status of set_authelia_policy
                if [ $? -eq 0 ]; then
                    echo -e "\n${YELLOW}Remember to restart Authelia for the policy change to take effect:${NC}"
                    echo -e "  ${CYAN}docker compose restart authelia${NC}"
                else
                    echo -e "${RED}Failed to set policy. Please check errors above.${NC}"
                    # Optionally offer to restore backup here
                fi
                ;;
            3)
                return 0
                ;;
            *)
                echo -e "${RED}Invalid choice. Please try again.${NC}"
                ;;
        esac
        echo # Add a newline for better readability between menu iterations
    done
}


##################################################
# PART 5: Authelia Account Management
##################################################

# Function to generate a secure, random passphrase
generate_passphrase() {
    # Generate a random passphrase using common words and numbers
    local words=("apple" "banana" "orange" "grape" "melon" "cherry" "lemon" "peach" "plum" "kiwi"
                "red" "blue" "green" "yellow" "purple" "cyan" "magenta" "white" "black" "grey"
                "dog" "cat" "bird" "fish" "horse" "tiger" "lion" "bear" "wolf" "fox"
                "river" "ocean" "mountain" "forest" "desert" "island" "valley" "canyon" "lake" "hill")
    
    local w1=${words[$((RANDOM % ${#words[@]}))]}
    local w2=${words[$((RANDOM % ${#words[@]}))]}
    local w3=${words[$((RANDOM % ${#words[@]}))]}
    local num=$((1000 + RANDOM % 9000))  # 4-digit number
    
    echo "${w1^}${w2^}${w3^}${num}"  # Capitalize first letter of each word
}

manage_authelia_accounts() {
    print_header "Authelia Account Management"
    
    local users_file="${CONFIG_ROOT:-.}/authelia/users_database.yml"
    
    if [ ! -f "$users_file" ]; then
        echo -e "${RED}Error: users_database.yml not found at $users_file${NC}"
        echo -e "${YELLOW}Would you like to create a new users database file? [y/N]:${NC}"
        read -r answer
        if [[ "$answer" =~ ^[Yy]$ ]]; then
            cat > "$users_file" <<EOL
# Authelia User Database
# Documentation: https://www.authelia.com/configuration/security/authentication/file/

users:
  admin:
    displayname: "Admin User"
    password: "$argon2id$v=19$m=102400,t=1,p=8$PBf/L9l3s7LwN6jX/B3tVg$9+q3kL8VAbpWj9Gv9Z6uA5bA4zT1fB2fH3aD5c6b7e8"
    email: admin@example.com
    groups:
      - admins
      - users
EOL
            echo -e "${GREEN}Created new users_database.yml file${NC}"
        else
            echo -e "${RED}Account management cancelled.${NC}"
            return 1
        fi
    fi
    
    local backup_file="${users_file}.${TIMESTAMP}.bak"
    cp "$users_file" "$backup_file"
    echo -e "${BLUE}Backed up users database to ${backup_file}${NC}"
    
    while true; do
        echo -e "\n${CYAN}${BOLD}Add Authelia User${NC}"
        echo -e "${BLUE}Enter a username (or press Enter to finish):${NC}"
        read -r username
        
        if [ -z "$username" ]; then
            break
        fi
        
        if grep -q "^[[:space:]]*${username}:" "$users_file"; then
            echo -e "${YELLOW}Warning: User '${username}' already exists.${NC}"
            echo -e "${YELLOW}Would you like to update this user? [y/N]:${NC}"
            read -r answer
            if [[ ! "$answer" =~ ^[Yy]$ ]]; then
                echo -e "${RED}Skipping user '${username}'.${NC}"
                continue
            fi
        fi
        
        echo -e "${BLUE}Enter display name for ${username}:${NC}"
        read -r displayname
        if [ -z "$displayname" ]; then
            displayname="$username"
        fi
        
        echo -e "${BLUE}Enter email for ${username}:${NC}"
        read -r email
        if [ -z "$email" ]; then
            email="${username}@example.com"
        fi
        
        echo -e "${BLUE}Select group(s) for ${username} (comma-separated, default: users):${NC}"
        echo -e "${CYAN}Available groups: admins, users${NC}"
        read -r groups
        if [ -z "$groups" ]; then
            groups="users"
        fi
        
        IFS=',' read -ra group_array <<< "$groups"
        formatted_groups=""
        for group in "${group_array[@]}"; do
            formatted_groups+="      - $(echo $group | xargs)\n"
        done
        
        generated_passphrase=$(generate_passphrase)
        echo -e "${GREEN}Generated passphrase: ${BOLD}${generated_passphrase}${NC}"
        echo -e "${YELLOW}Do you want to use this passphrase? [Y/n]:${NC}"
        read -r use_generated
        
        if [[ "$use_generated" =~ ^[Nn]$ ]]; then
            echo -e "${BLUE}Enter a custom password for ${username}:${NC}"
            read -rs password
            if [ -z "$password" ]; then
                echo -e "${RED}Error: Password cannot be empty.${NC}"
                continue
            fi
        else
            password="$generated_passphrase"
        fi
        
        echo -e "${CYAN}Generating password hash...${NC}"
        
        password_hash=$(docker compose run --rm authelia authelia crypto hash generate argon2 --password "$password" 2>/dev/null)
        
        if [ -z "$password_hash" ]; then
            echo -e "${RED}Error: Failed to generate password hash. Is the authelia container available?${NC}"
            echo -e "${YELLOW}Trying direct docker run method...${NC}"
            password_hash=$(docker run --rm authelia/authelia:latest authelia crypto hash generate argon2 --password "$password" 2>/dev/null)
            
            if [ -z "$password_hash" ]; then
                echo -e "${RED}Error: Both methods failed to generate a password hash. Skipping user.${NC}"
                continue
            fi
        fi
        
        password_hash=$(echo "$password_hash" | sed 's/^Digest: //')
        
        if [[ ! "$password_hash" =~ ^\$argon2id.*$ ]]; then
            echo -e "${RED}Error: Generated hash does not have the expected format. Actual value:${NC}"
            echo -e "${YELLOW}$password_hash${NC}"
            echo -e "${RED}Skipping user creation.${NC}"
            continue
        fi
        
        echo -e "${GREEN}Password hash generated successfully.${NC}"
        
        if grep -q "^[[:space:]]*${username}:" "$users_file"; then
            sed -i "/^[[:space:]]*${username}:/,/^[[:space:]]*[a-zA-Z0-9_-]\+:/ s/^/# UPDATED: /" "$users_file"
            sed -i "0,/^# UPDATED: [[:space:]]*[a-zA-Z0-9_-]\+:/ s/^# UPDATED: //" "$users_file"
        fi
        
        cat >> "$users_file" <<EOL
  ${username}:
    displayname: "${displayname}"
    password: "${password_hash}"
    email: ${email}
    groups:
$(echo -e "$formatted_groups")
EOL
        
        echo -e "${GREEN}User '${username}' added successfully!${NC}"
        echo -e "${CYAN}Password: ${BOLD}${password}${NC}"
        echo -e "${YELLOW}Please save this password securely.${NC}"
    done
    
    echo -e "\n${GREEN}${BOLD}Account Management Complete!${NC}"
    echo -e "${BLUE}Users file updated: ${users_file}${NC}"
    echo -e "${BLUE}Backup saved as: ${backup_file}${NC}"
    
    echo -e "${YELLOW}Would you like to restart Authelia to apply changes? [y/N]:${NC}"
    read -r restart
    if [[ "$restart" =~ ^[Yy]$ ]]; then
        if docker compose restart authelia; then
            echo -e "${GREEN}Authelia restarted successfully!${NC}"
        else
            echo -e "${RED}Failed to restart Authelia. Please restart manually.${NC}"
        fi
    else
        echo -e "${YELLOW}Remember to restart Authelia to apply these changes:${NC}"
        echo -e "${CYAN}  docker compose restart authelia${NC}"
    fi
}

##################################################
# MAIN SCRIPT
##################################################

# Function to display help message
show_help() {
    print_header "Docker Compose NAS - Setup Tool"
    echo -e "${BLUE}Usage: $0 [command] [arguments]${NC}"
    echo -e ""
    echo -e "${BLUE}Available Commands:${NC}"
    echo -e "  ${CYAN}update-env${NC}                Update .env file from .env.example, preserving values."
    echo -e "  ${CYAN}update-authelia${NC}           Update Authelia configuration (configuration.yml) from example,"
    echo -e "                           applying domain settings from .env and preserving secrets."
    echo -e "  ${CYAN}update-services${NC}           Update configurations for running *arr/qBittorrent/Bazarr containers"
    echo -e "                           (sets URL base, extracts API keys to .env)."
    echo -e "  ${CYAN}manage-accounts${NC}           Interactively add/update Authelia users in users_database.yml."
    echo -e "  ${CYAN}manage-policies${NC}           Interactively manage Authelia access policies for services."
    echo -e "  ${CYAN}list-policies${NC}             List services and their current Authelia policy."
    echo -e "  ${CYAN}set-policy <svc> <pol>${NC}    Set Authelia policy for a service (e.g., 'one_factor', 'bypass')."
    echo -e "  ${CYAN}cleanup${NC}                   Interactively clean up old backup files (.bak)."
    echo -e "  ${CYAN}all${NC}                       Run 'update-env', 'update-authelia', and 'update-services'."
    echo -e "  ${CYAN}help${NC}                      Show this help message."
    echo -e ""
    echo -e "${BLUE}Examples:${NC}"
    echo -e "  $0 update-authelia"
    echo -e "  $0 set-policy sonarr one_factor"
    echo -e "  $0 set-policy radarr bypass"
    echo -e "  $0 manage-policies"
    echo -e "  $0 all"
    echo -e ""
    echo -e "${YELLOW}Note:${NC} Some commands require Docker to be running and may restart containers."
    echo -e "${YELLOW}Policy changes require an Authelia restart ('docker compose restart authelia').${NC}"
}

# Check if any arguments were provided
if [ $# -eq 0 ]; then
    show_help
    exit 0
fi

# Process command line arguments
case "$1" in
    update-env)
        update_env_file
        ;;
    update-authelia)
        update_authelia_config
        ;;
    update-services)
        update_service_configs
        ;;
    manage-accounts)
        manage_authelia_accounts # Interactive
        ;;
    manage-policies)
        manage_authelia_policies # Interactive
        ;;
    list-policies)
        list_authelia_services
        ;;
    set-policy)
        if [ -z "$2" ] || [ -z "$3" ]; then
            echo -e "${RED}Error: Service name and policy are required.${NC}" >&2
            echo -e "Usage: $0 set-policy <service_name> <policy>${NC}" >&2
            echo -e "Valid policies: one_factor, two_factor, bypass, deny" >&2
            exit 1
        fi
        if ! check_file "$AUTHELIA_CONFIG"; then exit 1; fi
        # Backup is handled within set_authelia_policy if needed
        set_authelia_policy "$2" "$3" "$AUTHELIA_CONFIG"
        if [ $? -eq 0 ]; then
            echo -e "\n${YELLOW}Remember to restart Authelia for the policy change to take effect:${NC}"
            echo -e "  ${CYAN}docker compose restart authelia${NC}"
        fi
        ;;
    cleanup)
        cleanup_backups
        ;;
    all)
        print_header "Running All Updates"
        update_env_file
        update_authelia_config
        update_service_configs
        echo -e "\n${GREEN}${BOLD}All core updates completed!${NC}"
        echo -e "${BLUE}Review output for any required actions (e.g., setting new .env variables).${NC}"
        echo -e "${BLUE}Consider running 'manage-accounts' if needed.${NC}"
        echo -e "${YELLOW}Remember to restart relevant services or the full stack if necessary.${NC}"
        ;;
    help|-h|--help)
        show_help
        ;;
    *)
        echo -e "${RED}Unknown command: $1${NC}" >&2
        echo -e "${BLUE}Run '$0 help' for usage information.${NC}" >&2
        exit 1
        ;;
esac

exit 0