#!/usr/bin/env bash
#
# Vaultwarden Prometheus Metrics Exporter
#
# Prometheus textfile collector exporter for Vaultwarden.
# Uses the Vaultwarden admin API to collect user counts, organization
# counts, cipher totals, attachment stats, signup/invitation status,
# server version, and database size.
#
# Usage:
#   VAULTWARDEN_URL="https://vault.example.com" VAULTWARDEN_ADMIN_TOKEN="xxx" ./vaultwarden-exporter.sh
#   VAULTWARDEN_URL="https://vault.example.com" VAULTWARDEN_ADMIN_TOKEN="xxx" ./vaultwarden-exporter.sh --textfile
#   VAULTWARDEN_URL="https://vault.example.com" VAULTWARDEN_ADMIN_TOKEN="xxx" ./vaultwarden-exporter.sh --install
#
# Parameters:
#   --textfile   Write to textfile collector directory
#   --install    Create cron job for automatic collection
#   --help       Show usage
#
# Environment:
#   VAULTWARDEN_URL          Vaultwarden base URL (required)
#   VAULTWARDEN_ADMIN_TOKEN  Admin panel token (required)
#   VAULTWARDEN_DATA_DIR     Data directory for direct DB file size check (optional)
#   TEXTFILE_DIR             Textfile collector directory (default: /var/lib/node_exporter/textfile_collector)
#   CURL_TIMEOUT             API request timeout in seconds (default: 10)
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
# Version: 1.0
#
# Metrics Exported:
#   Core:
#     - vaultwarden_up
#     - vaultwarden_exporter_info{version}
#     - vaultwarden_server_version_info{version}
#
#   Users:
#     - vaultwarden_users_total
#     - vaultwarden_users_enabled
#
#   Organizations:
#     - vaultwarden_organizations_total
#
#   Vault:
#     - vaultwarden_ciphers_total
#     - vaultwarden_attachments_total
#     - vaultwarden_attachments_size_bytes
#
#   Configuration:
#     - vaultwarden_signups_allowed
#     - vaultwarden_invitations_allowed
#
#   Database:
#     - vaultwarden_database_size_bytes
#
#   Exporter:
#     - vaultwarden_exporter_duration_seconds
#     - vaultwarden_exporter_last_run_timestamp

set -euo pipefail

# --- Configuration ---
readonly VERSION="1.0"
readonly SCRIPT_NAME="$(basename "$0")"
VAULTWARDEN_URL="${VAULTWARDEN_URL:-}"
VAULTWARDEN_ADMIN_TOKEN="${VAULTWARDEN_ADMIN_TOKEN:-}"
VAULTWARDEN_DATA_DIR="${VAULTWARDEN_DATA_DIR:-}"
TEXTFILE_DIR="${TEXTFILE_DIR:-/var/lib/node_exporter/textfile_collector}"
CURL_TIMEOUT="${CURL_TIMEOUT:-10}"
TEXTFILE_MODE=false
OUTPUT=""
START_TIME=""
COOKIE_JAR=""

# --- Functions ---

usage() {
    cat <<EOF
Usage: $SCRIPT_NAME [OPTIONS]

Vaultwarden Prometheus Metrics Exporter

Options:
  --textfile    Write metrics to textfile collector directory
  --install     Create cron job for automatic collection
  --help        Show this help message

Environment Variables:
  VAULTWARDEN_URL          Vaultwarden base URL (required)
  VAULTWARDEN_ADMIN_TOKEN  Admin panel token (required)
  VAULTWARDEN_DATA_DIR     Data directory for DB file size (optional)
  TEXTFILE_DIR             Output directory (default: /var/lib/node_exporter/textfile_collector)
  CURL_TIMEOUT             Request timeout in seconds (default: 10)

Examples:
  VAULTWARDEN_URL="https://vault.example.com" VAULTWARDEN_ADMIN_TOKEN="xxx" $SCRIPT_NAME
  VAULTWARDEN_URL="https://vault.example.com" VAULTWARDEN_ADMIN_TOKEN="xxx" $SCRIPT_NAME --textfile
  VAULTWARDEN_URL="https://vault.example.com" VAULTWARDEN_ADMIN_TOKEN="xxx" $SCRIPT_NAME --install
EOF
    exit 0
}

check_dependencies() {
    local missing=()
    for cmd in curl jq; do
        if ! command -v "$cmd" &>/dev/null; then
            missing+=("$cmd")
        fi
    done
    if [[ ${#missing[@]} -gt 0 ]]; then
        echo "ERROR: Missing required commands: ${missing[*]}" >&2
        echo "Install with: apt install ${missing[*]}  OR  dnf install ${missing[*]}" >&2
        exit 1
    fi
}

validate_config() {
    if [[ -z "$VAULTWARDEN_URL" ]]; then
        echo "ERROR: VAULTWARDEN_URL environment variable is required" >&2
        exit 1
    fi
    if [[ -z "$VAULTWARDEN_ADMIN_TOKEN" ]]; then
        echo "ERROR: VAULTWARDEN_ADMIN_TOKEN environment variable is required" >&2
        exit 1
    fi
    # Strip trailing slash
    VAULTWARDEN_URL="${VAULTWARDEN_URL%/}"
}

setup_cookie_jar() {
    COOKIE_JAR=$(mktemp)
    trap 'rm -f "$COOKIE_JAR"' EXIT
}

admin_login() {
    # POST the admin token to get a session cookie
    curl -sf --max-time "$CURL_TIMEOUT" \
        -c "$COOKIE_JAR" \
        -d "token=${VAULTWARDEN_ADMIN_TOKEN}" \
        "${VAULTWARDEN_URL}/admin" \
        -o /dev/null 2>/dev/null
}

admin_get() {
    local endpoint="$1"
    curl -sf --max-time "$CURL_TIMEOUT" \
        -b "$COOKIE_JAR" \
        -H "Accept: application/json" \
        "${VAULTWARDEN_URL}${endpoint}" 2>/dev/null || echo ""
}

add_metric() {
    local name="$1"
    local type="$2"
    local help="$3"
    local value="$4"
    local labels="${5:-}"

    if [[ -n "$labels" ]]; then
        OUTPUT+="# HELP ${name} ${help}
# TYPE ${name} ${type}
${name}{${labels}} ${value}
"
    else
        OUTPUT+="# HELP ${name} ${help}
# TYPE ${name} ${type}
${name} ${value}
"
    fi
}

add_metric_value() {
    local name="$1"
    local value="$2"
    local labels="${3:-}"

    if [[ -n "$labels" ]]; then
        OUTPUT+="${name}{${labels}} ${value}
"
    else
        OUTPUT+="${name} ${value}
"
    fi
}

collect_health() {
    local alive_response
    alive_response=$(curl -sf --max-time "$CURL_TIMEOUT" \
        "${VAULTWARDEN_URL}/alive" 2>/dev/null || echo "")

    if [[ -z "$alive_response" ]]; then
        add_metric "vaultwarden_up" "gauge" "Vaultwarden reachability (1=up, 0=down)" "0"
        return 1
    fi

    add_metric "vaultwarden_up" "gauge" "Vaultwarden reachability (1=up, 0=down)" "1"
    return 0
}

collect_diagnostics() {
    local diag_json
    diag_json=$(admin_get "/admin/diagnostics")

    if [[ -z "$diag_json" ]]; then
        return
    fi

    # Server version
    local server_version
    server_version=$(echo "$diag_json" | jq -r '.version // empty' 2>/dev/null)
    if [[ -n "$server_version" ]]; then
        add_metric "vaultwarden_server_version_info" "gauge" "Vaultwarden server version" "1" "version=\"${server_version}\""
    fi

    # Database size from diagnostics (if available)
    local db_size
    db_size=$(echo "$diag_json" | jq -r '.db_size // empty' 2>/dev/null)
    if [[ -n "$db_size" ]]; then
        add_metric "vaultwarden_database_size_bytes" "gauge" "Database file size in bytes" "$db_size"
    fi

    # Signups allowed
    local signups
    signups=$(echo "$diag_json" | jq -r '.signups_allowed // empty' 2>/dev/null)
    if [[ "$signups" == "true" ]]; then
        add_metric "vaultwarden_signups_allowed" "gauge" "Whether new signups are allowed (1=yes, 0=no)" "1"
    elif [[ "$signups" == "false" ]]; then
        add_metric "vaultwarden_signups_allowed" "gauge" "Whether new signups are allowed (1=yes, 0=no)" "0"
    fi

    # Invitations allowed
    local invitations
    invitations=$(echo "$diag_json" | jq -r '.invitations_allowed // empty' 2>/dev/null)
    if [[ "$invitations" == "true" ]]; then
        add_metric "vaultwarden_invitations_allowed" "gauge" "Whether invitations are allowed (1=yes, 0=no)" "1"
    elif [[ "$invitations" == "false" ]]; then
        add_metric "vaultwarden_invitations_allowed" "gauge" "Whether invitations are allowed (1=yes, 0=no)" "0"
    fi
}

collect_users() {
    local users_json
    users_json=$(admin_get "/admin/users/overview")

    if [[ -z "$users_json" ]]; then
        return
    fi

    # Total users
    local total_users
    total_users=$(echo "$users_json" | jq 'length // 0' 2>/dev/null)
    add_metric "vaultwarden_users_total" "gauge" "Total number of registered users" "${total_users:-0}"

    # Enabled users
    local enabled_users
    enabled_users=$(echo "$users_json" | jq '[.[] | select(.Enabled == true)] | length // 0' 2>/dev/null)
    add_metric "vaultwarden_users_enabled" "gauge" "Number of enabled users" "${enabled_users:-0}"

    # Cipher totals (sum across all users)
    local total_ciphers
    total_ciphers=$(echo "$users_json" | jq '[.[].CipherCount // 0] | add // 0' 2>/dev/null)
    if [[ -n "$total_ciphers" ]]; then
        add_metric "vaultwarden_ciphers_total" "gauge" "Total number of cipher entries" "${total_ciphers:-0}"
    fi

    # Attachment totals
    local total_attachments
    total_attachments=$(echo "$users_json" | jq '[.[].AttachmentCount // 0] | add // 0' 2>/dev/null)
    if [[ -n "$total_attachments" ]]; then
        add_metric "vaultwarden_attachments_total" "gauge" "Total number of attachments" "${total_attachments:-0}"
    fi

    # Attachment size
    local total_attachment_size
    total_attachment_size=$(echo "$users_json" | jq '[.[].AttachmentSize // 0] | add // 0' 2>/dev/null)
    if [[ -n "$total_attachment_size" ]]; then
        add_metric "vaultwarden_attachments_size_bytes" "gauge" "Total attachment size in bytes" "${total_attachment_size:-0}"
    fi
}

collect_organizations() {
    local orgs_json
    orgs_json=$(admin_get "/admin/organizations/overview")

    if [[ -z "$orgs_json" ]]; then
        return
    fi

    local total_orgs
    total_orgs=$(echo "$orgs_json" | jq 'length // 0' 2>/dev/null)
    add_metric "vaultwarden_organizations_total" "gauge" "Total number of organizations" "${total_orgs:-0}"
}

collect_database_size() {
    # Direct file size check if data dir is set and using SQLite
    if [[ -n "$VAULTWARDEN_DATA_DIR" ]]; then
        local db_file="${VAULTWARDEN_DATA_DIR}/db.sqlite3"
        if [[ -f "$db_file" ]]; then
            local file_size
            file_size=$(stat -c %s "$db_file" 2>/dev/null || stat -f %z "$db_file" 2>/dev/null || echo "")
            if [[ -n "$file_size" ]]; then
                add_metric "vaultwarden_database_size_bytes" "gauge" "Database file size in bytes" "$file_size"
            fi
        fi
    fi
}

write_output() {
    if [[ "$TEXTFILE_MODE" == true ]]; then
        local output_file="${TEXTFILE_DIR}/vaultwarden.prom"
        local temp_file="${output_file}.$$"

        mkdir -p "$TEXTFILE_DIR"
        echo "$OUTPUT" > "$temp_file"
        mv "$temp_file" "$output_file"
    else
        echo "$OUTPUT"
    fi
}

install_cron() {
    if [[ $EUID -ne 0 ]]; then
        echo "ERROR: --install requires root" >&2
        exit 1
    fi

    local script_path
    script_path=$(readlink -f "$0")

    cat > /etc/cron.d/vaultwarden-exporter <<EOF
# Vaultwarden Prometheus Exporter — runs every 2 minutes
VAULTWARDEN_URL=${VAULTWARDEN_URL}
VAULTWARDEN_ADMIN_TOKEN=${VAULTWARDEN_ADMIN_TOKEN}
TEXTFILE_DIR=${TEXTFILE_DIR}
*/2 * * * * root ${script_path} --textfile 2>/dev/null
EOF

    chmod 644 /etc/cron.d/vaultwarden-exporter
    echo "Installed cron job: /etc/cron.d/vaultwarden-exporter"
    echo "Metrics will be written to: ${TEXTFILE_DIR}/vaultwarden.prom"
}

# --- Main ---

main() {
    # Parse arguments
    for arg in "$@"; do
        case "$arg" in
            --textfile) TEXTFILE_MODE=true ;;
            --install)
                check_dependencies
                validate_config
                install_cron
                exit 0
                ;;
            --help|-h) usage ;;
            *) echo "Unknown option: $arg" >&2; usage ;;
        esac
    done

    check_dependencies
    validate_config
    setup_cookie_jar

    START_TIME=$(date +%s%N)

    # Exporter info
    add_metric "vaultwarden_exporter_info" "gauge" "Exporter version information" "1" "version=\"${VERSION}\""

    # Collect metrics
    if collect_health; then
        admin_login
        collect_diagnostics
        collect_users
        collect_organizations
        collect_database_size
    fi

    # Exporter performance
    local end_time duration
    end_time=$(date +%s%N)
    duration=$(echo "scale=2; ($end_time - $START_TIME) / 1000000000" | bc 2>/dev/null || echo "0")
    add_metric "vaultwarden_exporter_duration_seconds" "gauge" "Time to generate all metrics" "$duration"
    add_metric "vaultwarden_exporter_last_run_timestamp" "gauge" "Unix timestamp of last successful run" "$(date +%s)"

    write_output
}

main "$@"
