#!/bin/bash
#############################################################
#### Database Backup Exporter for Prometheus              ####
#### Monitor MySQL and PostgreSQL backup freshness,       ####
#### size, and status via node_exporter textfile collector ####
####                                                     ####
#### Author: Phil Connor                                 ####
#### Contact: contact@mylinux.work                       ####
#### License: MIT                                        ####
#### Version: 1.0                                        ####
####                                                     ####
#### Usage: ./database-backup-exporter.sh [OPTIONS]      ####
#############################################################

set -euo pipefail

# -----------------------------
# Defaults
# -----------------------------
BACKUP_DIR="/opt/backups"
MAX_AGE=86400
PROM_FILE="/var/lib/node_exporter/database_backups.prom"
INTERVAL=300
RUN_ONCE=false

# -----------------------------
# Color codes
# -----------------------------
RED='\033[0;31m'
YELLOW='\033[1;33m'
GREEN='\033[0;32m'
NC='\033[0m'

# -----------------------------
# Logging
# -----------------------------
log_info() {
    echo -e "${GREEN}[INFO]${NC}  $(date '+%Y-%m-%d %H:%M:%S') $*"
}

log_warn() {
    echo -e "${YELLOW}[WARN]${NC}  $(date '+%Y-%m-%d %H:%M:%S') $*" >&2
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $(date '+%Y-%m-%d %H:%M:%S') $*" >&2
}

# -----------------------------
# Usage
# -----------------------------
usage() {
    cat <<EOF
Database Backup Exporter for Prometheus

Scans a backup directory for MySQL and PostgreSQL dump files and writes
metrics to a node_exporter textfile collector .prom file.

Usage:
  $(basename "$0") [OPTIONS]

Options:
  --backup-dir PATH    Directory containing backup files (default: /opt/backups)
  --max-age SECONDS    Maximum acceptable backup age in seconds (default: 86400 / 24h)
  --prom-file PATH     Path to write .prom metrics file (default: /var/lib/node_exporter/database_backups.prom)
  --interval SECONDS   Collection interval when running as daemon (default: 300)
  --once               Run once and exit instead of looping
  --help               Show this help message

Supported file extensions:
  .sql  .sql.gz  .dump  .pgdump

Filename pattern:
  <dbname>_YYYYMMDD[HHMMSS].<ext>
  Examples: myapp_20260309.sql.gz  orders_20260308120000.pgdump

EOF
    exit 0
}

# -----------------------------
# Parse arguments
# -----------------------------
parse_args() {
    while [[ $# -gt 0 ]]; do
        case "$1" in
            --backup-dir)
                BACKUP_DIR="$2"
                shift 2
                ;;
            --max-age)
                MAX_AGE="$2"
                shift 2
                ;;
            --prom-file)
                PROM_FILE="$2"
                shift 2
                ;;
            --interval)
                INTERVAL="$2"
                shift 2
                ;;
            --once)
                RUN_ONCE=true
                shift
                ;;
            --help)
                usage
                ;;
            *)
                log_error "Unknown option: $1"
                usage
                ;;
        esac
    done
}

# -----------------------------
# Detect backup type from ext
# -----------------------------
detect_type() {
    local filename="$1"
    case "$filename" in
        *.pgdump|*.dump)
            echo "postgres"
            ;;
        *.sql|*.sql.gz)
            echo "mysql"
            ;;
        *)
            echo "unknown"
            ;;
    esac
}

# -----------------------------
# Extract database name
# -----------------------------
extract_dbname() {
    local filename
    filename="$(basename "$1")"
    # Strip all known extensions
    filename="${filename%.gz}"
    filename="${filename%.sql}"
    filename="${filename%.dump}"
    filename="${filename%.pgdump}"
    # Expect pattern: dbname_YYYYMMDD... — grab everything before the date segment
    echo "$filename" | sed -E 's/_[0-9]{8,14}$//'
}

# -----------------------------
# Collect and write metrics
# -----------------------------
collect_metrics() {
    local backup_dir="$1"
    local max_age="$2"
    local now
    now="$(date +%s)"

    if [[ ! -d "$backup_dir" ]]; then
        log_error "Backup directory does not exist: $backup_dir"
        return 1
    fi

    # Associative arrays keyed by "dbname|type"
    declare -A latest_ts
    declare -A latest_size
    declare -A file_count

    # Scan for backup files
    local found=0
    while IFS= read -r -d '' file; do
        local base
        base="$(basename "$file")"
        local btype
        btype="$(detect_type "$base")"
        [[ "$btype" == "unknown" ]] && continue

        local dbname
        dbname="$(extract_dbname "$file")"
        [[ -z "$dbname" ]] && continue

        local key="${dbname}|${btype}"
        local mtime
        mtime="$(stat -c '%Y' "$file" 2>/dev/null)" || continue
        local fsize
        fsize="$(stat -c '%s' "$file" 2>/dev/null)" || continue

        # Track count
        file_count[$key]=$(( ${file_count[$key]:-0} + 1 ))

        # Track most recent
        if [[ -z "${latest_ts[$key]:-}" ]] || (( mtime > latest_ts[$key] )); then
            latest_ts[$key]="$mtime"
            latest_size[$key]="$fsize"
        fi

        found=$((found + 1))
    done < <(find "$backup_dir" -type f \( -name '*.sql' -o -name '*.sql.gz' -o -name '*.dump' -o -name '*.pgdump' \) -print0 2>/dev/null)

    log_info "Found $found backup file(s) in $backup_dir"

    # Build output
    local output=""

    output+="# HELP db_backup_last_timestamp Unix timestamp of most recent backup.\n"
    output+="# TYPE db_backup_last_timestamp gauge\n"
    for key in "${!latest_ts[@]}"; do
        local dbname="${key%%|*}"
        local btype="${key##*|}"
        output+="db_backup_last_timestamp{database=\"${dbname}\",type=\"${btype}\"} ${latest_ts[$key]}\n"
    done

    output+="# HELP db_backup_age_seconds Seconds since most recent backup.\n"
    output+="# TYPE db_backup_age_seconds gauge\n"
    for key in "${!latest_ts[@]}"; do
        local dbname="${key%%|*}"
        local btype="${key##*|}"
        local age=$(( now - latest_ts[$key] ))
        output+="db_backup_age_seconds{database=\"${dbname}\",type=\"${btype}\"} ${age}\n"
    done

    output+="# HELP db_backup_size_bytes Size of most recent backup file in bytes.\n"
    output+="# TYPE db_backup_size_bytes gauge\n"
    for key in "${!latest_size[@]}"; do
        local dbname="${key%%|*}"
        local btype="${key##*|}"
        output+="db_backup_size_bytes{database=\"${dbname}\",type=\"${btype}\"} ${latest_size[$key]}\n"
    done

    output+="# HELP db_backup_count Number of backup files found.\n"
    output+="# TYPE db_backup_count gauge\n"
    for key in "${!file_count[@]}"; do
        local dbname="${key%%|*}"
        local btype="${key##*|}"
        output+="db_backup_count{database=\"${dbname}\",type=\"${btype}\"} ${file_count[$key]}\n"
    done

    output+="# HELP db_backup_fresh 1 if backup is within max_age, 0 if stale.\n"
    output+="# TYPE db_backup_fresh gauge\n"
    for key in "${!latest_ts[@]}"; do
        local dbname="${key%%|*}"
        local btype="${key##*|}"
        local age=$(( now - latest_ts[$key] ))
        local fresh=1
        if (( age > max_age )); then
            fresh=0
            log_warn "Stale backup: database=${dbname} type=${btype} age=${age}s exceeds max_age=${max_age}s"
        fi
        output+="db_backup_fresh{database=\"${dbname}\",type=\"${btype}\"} ${fresh}\n"
    done

    output+="# HELP db_backup_exporter_last_run Timestamp of last exporter run.\n"
    output+="# TYPE db_backup_exporter_last_run gauge\n"
    output+="db_backup_exporter_last_run ${now}\n"

    echo "$output"
}

# -----------------------------
# Write metrics atomically
# -----------------------------
write_metrics() {
    local content="$1"
    local prom_file="$2"

    local prom_dir
    prom_dir="$(dirname "$prom_file")"

    if [[ ! -d "$prom_dir" ]]; then
        log_error "Prom directory does not exist: $prom_dir"
        return 1
    fi

    local tmp_file
    tmp_file="$(mktemp "${prom_dir}/.database_backups.prom.XXXXXX")"

    echo -e "$content" > "$tmp_file"
    mv "$tmp_file" "$prom_file"

    log_info "Metrics written to $prom_file"
}

# -----------------------------
# Main
# -----------------------------
main() {
    parse_args "$@"

    log_info "Database Backup Exporter starting"
    log_info "Backup directory: $BACKUP_DIR"
    log_info "Max backup age:   ${MAX_AGE}s"
    log_info "Prom file:        $PROM_FILE"

    while true; do
        local metrics
        metrics="$(collect_metrics "$BACKUP_DIR" "$MAX_AGE")" || true

        if [[ -n "$metrics" ]]; then
            write_metrics "$metrics" "$PROM_FILE"
        fi

        if [[ "$RUN_ONCE" == true ]]; then
            log_info "Single run complete, exiting"
            break
        fi

        log_info "Sleeping ${INTERVAL}s until next collection"
        sleep "$INTERVAL"
    done
}

main "$@"