#!/bin/bash
#set -x 

### Global Settings ###
# Set your Home Assistant SSH credentials
HA_SSH_KEY="SSH_PRIVATE_KEY"  # Replace with your SSH Private key for HomeAssistant
HA_SSH_USER="SSH_USER"        # Replace with your HomeAssistant server username (Default is root)
HA_SSH_HOST="SSH_HOST_IP_OR_HOSTNAME"  # Replace with your HomeAssistant server IP address
HA_SSH_PORT="SSH_PORT"        # Replace with your HomeAssistant server port

## Storage
# Set the backup destination path
BACKUP_DEST="BACKUP_FILE_DESTINATION" # Replace with the location for the backup files

# Set the location of the Home Assistant backups (Set to default location /backup)
# If you have a custom install, using the likes of Supervised
# this location should normally be /usr/share/hassio/backup
BACKUP_FILE_LOCATION="/backup"

# Storage threshold (in KB) for triggering cleanup
STORAGE_THRESHOLD_KB=2097152  # 2GB = 2 * 1024 * 1024 KB
# Set clean up old backups (adjust the number of days to keep)
LOCAL_BACKUP_CLEANUP="90"  # Set the number of days that local backup files will be kept for (Default 90 days)
REMOTE_BACKUP_CLEANUP="30" # Set the number of days that backups on the Home Assistant server will be kept for (Default 30 days)

## Notification Settings
# ntfy settings
NTFY_ENABLED=false  # Set to 'true' to enable ntfy notifications
NTFY_SERVER="NTFY_SERVER_FQDN_ADDRESS"  # Replace with your ntfy server address
NTFY_TOPIC="NTFY_TOPIC"  # Replace with your ntfy topic

# Discord webhook settings
DISCORD_ENABLED=false  # Set to 'true' to enable Discord notifications
DISCORD_WEBHOOK_URL="DISCORD_WEBHOOK_URL"  # Replace with your Discord webhook URL

### Script ###
## Sanity Check##
# Function to check if the necessary variables are set and if jq and yq are installed
sanity_check() {
    # Check if jq and yq are installed
    if ! command -v jq &> /dev/null; then
        echo "Error: jq is not installed. Please install jq before running this script."
        exit 1
    fi
    
    if ! command -v yq &> /dev/null; then
        echo "Error: yq is not installed. Please install yq before running this script."
        exit 1
    fi
    
    # Check for each variable
    if [ -z "$HA_SSH_KEY" ] || [ "$HA_SSH_KEY" == "SSH_PRIVATE_KEY" ]; then
        echo "Error: SSH_PRIVATE_KEY is not set."
        exit 1
    fi

    if [ -z "$HA_SSH_USER" ] || [ "$HA_SSH_USER" == "SSH_USER" ]; then
        echo "Error: SSH_USER is not set."
        exit 1
    fi

    if [ -z "$HA_SSH_HOST" ] || [ "$HA_SSH_HOST" == "SSH_HOST_IP_OR_HOSTNAME" ]; then
        echo "Error: SSH_HOST_IP_OR_HOSTNAME is not set."
        exit 1
    fi

    if [ -z "$HA_SSH_PORT" ] || [ "$HA_SSH_PORT" == "SSH_PORT" ]; then
        echo "Error: SSH_PORT is not set."
        exit 1
    fi

    if [ -z "$BACKUP_DEST" ] || [ "$BACKUP_DEST" == "BACKUP_FILE_DESTINATION" ]; then
        echo "Error: BACKUP_FILE_DESTINATION is not set."
        exit 1
    fi

    if [ "$NTFY_ENABLED" == "true" ]; then
        if [ -z "$NTFY_SERVER" ] || [ "$NTFY_SERVER" == "NTFY_SERVER_FQDN_ADDRESS" ]; then
            echo "Error: NTFY_SERVER_FQDN_ADDRESS is not set."
            exit 1
        fi
        if [ -z "$NTFY_TOPIC" ] || [ "$NTFY_TOPIC" == "NTFY_TOPIC" ]; then
            echo "Error: NTFY_TOPIC is not set."
            exit 1
        fi
    fi

    if [ "$DISCORD_ENABLED" == "true" ]; then
        if [ -z "$DISCORD_WEBHOOK_URL" ] || [ "$DISCORD_WEBHOOK_URL" == "DISCORD_WEBHOOK_URL" ]; then
            echo "Error: DISCORD_WEBHOOK_URL is not set."
            exit 1
        fi
    fi
}

# Call the function to check requirements and variables
sanity_check

# Function to remove old backups from Home Assistant server (older than $REMOTE_BACKUP_CLEANUP days)
remove_old_backups_ssh() {
    echo "Removing backups older than $REMOTE_BACKUP_CLEANUP days on Home Assistant server..."

    # Fetch the list of backups and parse the YAML output using yq
    BACKUPS_OUTPUT=$(ssh -i "$HA_SSH_KEY" -p "$HA_SSH_PORT" "$HA_SSH_USER@$HA_SSH_HOST" "ha backups")
    if ! echo "$BACKUPS_OUTPUT" | yq eval >/dev/null 2>&1; then
        echo "Error: The output of 'ha backups' is not valid YAML."
        return 1
    fi

    # Convert the YAML to JSON and extract old backups using jq
    OLD_BACKUPS=$(echo "$BACKUPS_OUTPUT" | yq eval -o=json | jq --argjson days "$REMOTE_BACKUP_CLEANUP" -r '.backups[] | select((.date | sub("\\.[0-9]+\\+00:00$"; "Z") | fromdateiso8601) < (now - ($days * 86400))) | .slug')

    # Remove each old backup by its slug
    if [ -n "$OLD_BACKUPS" ]; then
        for slug in $OLD_BACKUPS; do
            echo "Removing backup with slug: $slug"
            ssh -i "$HA_SSH_KEY" -p "$HA_SSH_PORT" "$HA_SSH_USER@$HA_SSH_HOST" "ha backups remove $slug"
            if [ $? -eq 0 ]; then
                echo "Successfully removed backup with slug: $slug"
                send_ntfy_notification "Old backup removed: $slug"
                send_discord_notification "Old backup removed: $slug"
            else
                echo "Failed to remove backup with slug: $slug"
            fi
        done
    else
        echo "No old backups found to remove."
    fi
}


### Function for sending notifications ###
send_ntfy_notification() {
    local message=$1
    if [ "$NTFY_ENABLED" = true ]; then
        curl -X POST -d "$message" $NTFY_URL
    fi
}

send_discord_notification() {
    local message=$1
    if [ "$DISCORD_ENABLED" = true ]; then
        curl -X POST -H "Content-Type: application/json" -d '{"content": "'"$message"'"}' $DISCORD_WEBHOOK_URL
    fi
}

# ntfy URL setup
NTFY_URL="$NTFY_SERVER/$NTFY_TOPIC"

# Timestamp variable for consistent naming
TIMESTAMP=$(date '+%Y-%m-%d_%H-%M-%S')

# Start the timer
START_TIME=$(date +%s)

# Send notification that backup has started
send_ntfy_notification "Backup process started at $TIMESTAMP"
send_discord_notification "Backup process started at $TIMESTAMP"

# Check available disk space on Home Assistant server
AVAILABLE_SPACE=$(ssh -i "$HA_SSH_KEY" -p "$HA_SSH_PORT" "$HA_SSH_USER@$HA_SSH_HOST" "df --output=avail / | tail -n 1")
if [ "$AVAILABLE_SPACE" -lt "$STORAGE_THRESHOLD_KB" ]; then
    echo "Disk space is less than 2GB on Home Assistant server. Proceeding with remote backup cleanup..."
    # Remove old backups on Home Assistant server if less than 2GB available
    remove_old_backups_ssh
else
    echo "Sufficient disk space available on Home Assistant server. Skipping remote backup cleanup."
fi

# Create a backup on the Home Assistant server and extract the slug
BACKUP_OUTPUT=$(ssh -i "$HA_SSH_KEY" -p "$HA_SSH_PORT" "$HA_SSH_USER@$HA_SSH_HOST" "ha backup new --name=auto_backup_$TIMESTAMP")
if [ $? -ne 0 ]; then
    # Send notification if backup creation fails
    send_ntfy_notification "Error: Failed to create backup on Home Assistant server."
    send_discord_notification "Error: Failed to create backup on Home Assistant server."
    echo "Error: Failed to create backup on Home Assistant server."
    exit 1
fi

# Extract the backup slug and strip any carriage return characters
BACKUP_SLUG=$(echo "$BACKUP_OUTPUT" | awk '/slug:/ {print $2}' | tr -d '\r')
if [ -z "$BACKUP_SLUG" ]; then
    # Send notification if backup slug extraction fails
    send_ntfy_notification "Error: Failed to extract backup slug."
    send_discord_notification "Error: Failed to extract backup slug."
    echo "Error: Failed to extract backup slug."
    exit 1
fi

# Send notification that backup was created
send_ntfy_notification "Backup created successfully with slug: $BACKUP_SLUG"
send_discord_notification "Backup created successfully with slug: $BACKUP_SLUG"

# Create a filename for the backup
BACKUP_FILE="$BACKUP_DEST/${TIMESTAMP}_${BACKUP_SLUG}.tar"

# Copy the backup file from Home Assistant to your local machine
scp -i "$HA_SSH_KEY" -P "$HA_SSH_PORT" "$HA_SSH_USER@$HA_SSH_HOST:$BACKUP_FILE_LOCATION/${BACKUP_SLUG}.tar" "$BACKUP_FILE"
if [ $? -ne 0 ]; then
    # Send notification if backup copy fails
    send_ntfy_notification "Error: Failed to copy backup to local machine."
    send_discord_notification "Error: Failed to copy backup to local machine."
    echo "Error: Failed to copy backup to local machine."
    exit 1
fi

# Get the size of the backup file
BACKUP_FILE_SIZE=$(stat -c %s "$BACKUP_FILE")
if [ $? -ne 0 ]; then
    send_ntfy_notification "Error: Failed to get backup file size."
    send_discord_notification "Error: Failed to get backup file size."
    echo "Error: Failed to get backup file size."
    exit 1
fi

# Clean up old backups locally (older than $LOCAL_BACKUP_CLEANUP days)
find "$BACKUP_DEST" -name '*.tar' -mtime +$LOCAL_BACKUP_CLEANUP -exec rm {} \;
if [ $? -ne 0 ]; then
    # Send notification if cleanup fails
    send_ntfy_notification "Error: Failed to clean up old backups locally."
    send_discord_notification "Error: Failed to clean up old backups locally."
    echo "Error: Failed to clean up old backups locally."
    exit 1
fi

# Send notification that the cleanup was successful
send_ntfy_notification "Old backups cleaned up successfully."
send_discord_notification "Old backups cleaned up successfully."

# Success message
echo "Backup completed successfully: $BACKUP_FILE"
send_ntfy_notification "Backup completed successfully: $BACKUP_FILE"
send_discord_notification "Backup completed successfully: $BACKUP_FILE"

# Get the total time taken by the script
END_TIME=$(date +%s)
RUNTIME=$((END_TIME - START_TIME))

# Send notification with backup file size and script runtime
send_ntfy_notification "Backup file size: $BACKUP_FILE_SIZE bytes. Script runtime: $RUNTIME seconds."
send_discord_notification "Backup file size: $BACKUP_FILE_SIZE bytes. Script runtime: $RUNTIME seconds."