#!/usr/bin/env bash
#
# RabbitMQ Prometheus Metrics Exporter
#
# Prometheus textfile collector exporter for RabbitMQ.
# Uses the RabbitMQ Management HTTP API to collect cluster health,
# queue depth, message rates, connection/channel counts, memory/disk
# watermarks, Erlang processes, and cluster partition status.
#
# Usage:
#   RABBITMQ_USER="admin" RABBITMQ_PASS="password" ./rabbitmq-exporter.sh
#   RABBITMQ_USER="admin" RABBITMQ_PASS="password" ./rabbitmq-exporter.sh --textfile
#   RABBITMQ_USER="admin" RABBITMQ_PASS="password" ./rabbitmq-exporter.sh --install
#
# Parameters:
#   --textfile   Write to textfile collector directory
#   --install    Create cron job for automatic collection
#   --help       Show usage
#
# Environment:
#   RABBITMQ_URL     RabbitMQ Management API URL (default: http://localhost:15672)
#   RABBITMQ_USER    Username (required)
#   RABBITMQ_PASS    Password (required)
#   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:
#     - rabbitmq_up
#     - rabbitmq_exporter_info{version}
#     - rabbitmq_connections_total
#     - rabbitmq_channels_total
#     - rabbitmq_consumers_total
#     - rabbitmq_queues_total
#
#   Message Rates:
#     - rabbitmq_messages_published_total
#     - rabbitmq_messages_delivered_total
#     - rabbitmq_messages_acknowledged_total
#
#   Queues:
#     - rabbitmq_queue_messages{vhost,queue}
#     - rabbitmq_queue_messages_ready{vhost,queue}
#     - rabbitmq_queue_consumers{vhost,queue}
#
#   Nodes:
#     - rabbitmq_memory_used_bytes
#     - rabbitmq_memory_limit_bytes
#     - rabbitmq_memory_alarm
#     - rabbitmq_disk_free_bytes
#     - rabbitmq_disk_free_limit_bytes
#     - rabbitmq_disk_alarm
#     - rabbitmq_erlang_processes_used
#     - rabbitmq_cluster_partitions
#
#   Exporter:
#     - rabbitmq_exporter_duration_seconds
#     - rabbitmq_exporter_last_run_timestamp

set -euo pipefail

# --- Configuration ---
readonly VERSION="1.0"
readonly SCRIPT_NAME="$(basename "$0")"
RABBITMQ_URL="${RABBITMQ_URL:-http://localhost:15672}"
RABBITMQ_USER="${RABBITMQ_USER:-}"
RABBITMQ_PASS="${RABBITMQ_PASS:-}"
TEXTFILE_DIR="${TEXTFILE_DIR:-/var/lib/node_exporter/textfile_collector}"
CURL_TIMEOUT="${CURL_TIMEOUT:-10}"
TEXTFILE_MODE=false
OUTPUT=""
START_TIME=""

# --- Functions ---

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

RabbitMQ 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:
  RABBITMQ_URL     RabbitMQ Management API URL (default: http://localhost:15672)
  RABBITMQ_USER    Username (required)
  RABBITMQ_PASS    Password (required)
  TEXTFILE_DIR     Output directory (default: /var/lib/node_exporter/textfile_collector)
  CURL_TIMEOUT     Request timeout in seconds (default: 10)

Examples:
  RABBITMQ_USER="admin" RABBITMQ_PASS="password" $SCRIPT_NAME
  RABBITMQ_USER="admin" RABBITMQ_PASS="password" $SCRIPT_NAME --textfile
  RABBITMQ_URL="http://rabbitmq.example.com:15672" RABBITMQ_USER="admin" RABBITMQ_PASS="password" $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() {
    # Strip trailing slash
    RABBITMQ_URL="${RABBITMQ_URL%/}"

    if [[ -z "$RABBITMQ_USER" ]]; then
        echo "ERROR: RABBITMQ_USER environment variable is required" >&2
        exit 1
    fi

    if [[ -z "$RABBITMQ_PASS" ]]; then
        echo "ERROR: RABBITMQ_PASS environment variable is required" >&2
        exit 1
    fi
}

api_get() {
    local endpoint="$1"
    local curl_args=(-sf --max-time "$CURL_TIMEOUT")

    curl_args+=(-u "${RABBITMQ_USER}:${RABBITMQ_PASS}")

    curl "${curl_args[@]}" "${RABBITMQ_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_overview() {
    local overview_json
    overview_json=$(api_get "/api/overview")

    if [[ -z "$overview_json" ]]; then
        add_metric "rabbitmq_up" "gauge" "RabbitMQ reachability (1=up, 0=down)" "0"
        return 1
    fi

    add_metric "rabbitmq_up" "gauge" "RabbitMQ reachability (1=up, 0=down)" "1"

    # Object totals
    local connections channels consumers queues
    connections=$(echo "$overview_json" | jq '.object_totals.connections // 0' 2>/dev/null)
    channels=$(echo "$overview_json" | jq '.object_totals.channels // 0' 2>/dev/null)
    consumers=$(echo "$overview_json" | jq '.object_totals.consumers // 0' 2>/dev/null)
    queues=$(echo "$overview_json" | jq '.object_totals.queues // 0' 2>/dev/null)

    add_metric "rabbitmq_connections_total" "gauge" "Total number of open connections" "${connections:-0}"
    add_metric "rabbitmq_channels_total" "gauge" "Total number of open channels" "${channels:-0}"
    add_metric "rabbitmq_consumers_total" "gauge" "Total number of consumers" "${consumers:-0}"
    add_metric "rabbitmq_queues_total" "gauge" "Total number of queues" "${queues:-0}"

    # Message rates
    local published delivered acknowledged
    published=$(echo "$overview_json" | jq '.message_stats.publish // 0' 2>/dev/null)
    delivered=$(echo "$overview_json" | jq '.message_stats.deliver_get // 0' 2>/dev/null)
    acknowledged=$(echo "$overview_json" | jq '.message_stats.ack // 0' 2>/dev/null)

    add_metric "rabbitmq_messages_published_total" "counter" "Total messages published" "${published:-0}"
    add_metric "rabbitmq_messages_delivered_total" "counter" "Total messages delivered to consumers" "${delivered:-0}"
    add_metric "rabbitmq_messages_acknowledged_total" "counter" "Total messages acknowledged by consumers" "${acknowledged:-0}"

    return 0
}

collect_queues() {
    local queues_json
    queues_json=$(api_get "/api/queues")

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

    local queue_count
    queue_count=$(echo "$queues_json" | jq 'length' 2>/dev/null)

    if [[ "${queue_count:-0}" -eq 0 ]]; then
        return
    fi

    # Add HELP/TYPE headers once, then individual values per queue
    OUTPUT+="# HELP rabbitmq_queue_messages Total messages in queue
# TYPE rabbitmq_queue_messages gauge
"

    local i vhost queue messages
    for ((i = 0; i < queue_count; i++)); do
        vhost=$(echo "$queues_json" | jq -r ".[$i].vhost // \"\"" 2>/dev/null)
        queue=$(echo "$queues_json" | jq -r ".[$i].name // \"\"" 2>/dev/null)
        messages=$(echo "$queues_json" | jq ".[$i].messages // 0" 2>/dev/null)

        if [[ -n "$queue" ]]; then
            add_metric_value "rabbitmq_queue_messages" "${messages:-0}" "vhost=\"${vhost}\",queue=\"${queue}\""
        fi
    done

    OUTPUT+="# HELP rabbitmq_queue_messages_ready Messages ready to be delivered
# TYPE rabbitmq_queue_messages_ready gauge
"

    local messages_ready
    for ((i = 0; i < queue_count; i++)); do
        vhost=$(echo "$queues_json" | jq -r ".[$i].vhost // \"\"" 2>/dev/null)
        queue=$(echo "$queues_json" | jq -r ".[$i].name // \"\"" 2>/dev/null)
        messages_ready=$(echo "$queues_json" | jq ".[$i].messages_ready // 0" 2>/dev/null)

        if [[ -n "$queue" ]]; then
            add_metric_value "rabbitmq_queue_messages_ready" "${messages_ready:-0}" "vhost=\"${vhost}\",queue=\"${queue}\""
        fi
    done

    OUTPUT+="# HELP rabbitmq_queue_consumers Number of consumers on queue
# TYPE rabbitmq_queue_consumers gauge
"

    local queue_consumers
    for ((i = 0; i < queue_count; i++)); do
        vhost=$(echo "$queues_json" | jq -r ".[$i].vhost // \"\"" 2>/dev/null)
        queue=$(echo "$queues_json" | jq -r ".[$i].name // \"\"" 2>/dev/null)
        queue_consumers=$(echo "$queues_json" | jq ".[$i].consumers // 0" 2>/dev/null)

        if [[ -n "$queue" ]]; then
            add_metric_value "rabbitmq_queue_consumers" "${queue_consumers:-0}" "vhost=\"${vhost}\",queue=\"${queue}\""
        fi
    done
}

collect_nodes() {
    local nodes_json
    nodes_json=$(api_get "/api/nodes")

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

    local node_count
    node_count=$(echo "$nodes_json" | jq 'length' 2>/dev/null)

    if [[ "${node_count:-0}" -eq 0 ]]; then
        return
    fi

    # Aggregate metrics across all nodes
    local total_mem_used=0
    local total_mem_limit=0
    local total_disk_free=0
    local total_disk_limit=0
    local total_erlang_procs=0
    local total_partitions=0
    local mem_alarm=0
    local disk_alarm=0

    local i
    for ((i = 0; i < node_count; i++)); do
        local mem_used mem_limit disk_free disk_limit erlang_procs partitions node_mem_alarm node_disk_alarm

        mem_used=$(echo "$nodes_json" | jq ".[$i].mem_used // 0" 2>/dev/null)
        mem_limit=$(echo "$nodes_json" | jq ".[$i].mem_limit // 0" 2>/dev/null)
        disk_free=$(echo "$nodes_json" | jq ".[$i].disk_free // 0" 2>/dev/null)
        disk_limit=$(echo "$nodes_json" | jq ".[$i].disk_free_limit // 0" 2>/dev/null)
        erlang_procs=$(echo "$nodes_json" | jq ".[$i].proc_used // 0" 2>/dev/null)
        partitions=$(echo "$nodes_json" | jq ".[$i].partitions | length // 0" 2>/dev/null)
        node_mem_alarm=$(echo "$nodes_json" | jq ".[$i].mem_alarm // false" 2>/dev/null)
        node_disk_alarm=$(echo "$nodes_json" | jq ".[$i].disk_free_alarm // false" 2>/dev/null)

        total_mem_used=$((total_mem_used + ${mem_used:-0}))
        total_mem_limit=$((total_mem_limit + ${mem_limit:-0}))
        total_disk_free=$((total_disk_free + ${disk_free:-0}))
        total_disk_limit=$((total_disk_limit + ${disk_limit:-0}))
        total_erlang_procs=$((total_erlang_procs + ${erlang_procs:-0}))
        total_partitions=$((total_partitions + ${partitions:-0}))

        if [[ "$node_mem_alarm" == "true" ]]; then
            mem_alarm=1
        fi

        if [[ "$node_disk_alarm" == "true" ]]; then
            disk_alarm=1
        fi
    done

    # Memory metrics
    add_metric "rabbitmq_memory_used_bytes" "gauge" "Memory used by RabbitMQ in bytes" "$total_mem_used"
    add_metric "rabbitmq_memory_limit_bytes" "gauge" "Memory high watermark limit in bytes" "$total_mem_limit"
    add_metric "rabbitmq_memory_alarm" "gauge" "Whether the memory alarm is active (1=active, 0=inactive)" "$mem_alarm"

    # Disk metrics
    add_metric "rabbitmq_disk_free_bytes" "gauge" "Free disk space in bytes" "$total_disk_free"
    add_metric "rabbitmq_disk_free_limit_bytes" "gauge" "Disk free space low watermark limit in bytes" "$total_disk_limit"
    add_metric "rabbitmq_disk_alarm" "gauge" "Whether the disk alarm is active (1=active, 0=inactive)" "$disk_alarm"

    # Erlang processes
    add_metric "rabbitmq_erlang_processes_used" "gauge" "Number of Erlang processes used" "$total_erlang_procs"

    # Cluster partitions
    add_metric "rabbitmq_cluster_partitions" "gauge" "Number of network partitions detected" "$total_partitions"
}

write_output() {
    if [[ "$TEXTFILE_MODE" == true ]]; then
        local output_file="${TEXTFILE_DIR}/rabbitmq.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/rabbitmq-exporter <<EOF
# RabbitMQ Prometheus Exporter — runs every 2 minutes
RABBITMQ_URL=${RABBITMQ_URL}
RABBITMQ_USER=${RABBITMQ_USER}
RABBITMQ_PASS=${RABBITMQ_PASS}
TEXTFILE_DIR=${TEXTFILE_DIR}
*/2 * * * * root ${script_path} --textfile 2>/dev/null
EOF

    chmod 644 /etc/cron.d/rabbitmq-exporter
    echo "Installed cron job: /etc/cron.d/rabbitmq-exporter"
    echo "Metrics will be written to: ${TEXTFILE_DIR}/rabbitmq.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

    START_TIME=$(date +%s%N)

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

    # Collect metrics
    if collect_overview; then
        collect_queues
        collect_nodes
    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 "rabbitmq_exporter_duration_seconds" "gauge" "Time to generate all metrics" "$duration"
    add_metric "rabbitmq_exporter_last_run_timestamp" "gauge" "Unix timestamp of last successful run" "$(date +%s)"

    write_output
}

main "$@"