#!/bin/bash

################################################
#### Wickr Enterprise Metrics Collector     ####
#### for Prometheus node_exporter textfile  ####
####                                        ####
#### Author: Phil Connor                    ####
#### Contact: contact@mylinux.work          ####
#### Version: 1.00-030726                   ####
################################################

set -o pipefail

SCRIPT_NAME=$(basename "$0")
readonly SCRIPT_NAME

# Default configuration
readonly DEFAULT_NODE_DIR="/var/lib/node_exporter"
readonly DEFAULT_COLLECTION_INTERVAL=300
readonly DEFAULT_CURL_TIMEOUT=30

# Configuration variables (can be overridden by environment)
WICKR_ADMIN_URL=${WICKR_ADMIN_URL:-}
WICKR_API_TOKEN=${WICKR_API_TOKEN:-}
NODE_DIR=${NODE_DIR:-$DEFAULT_NODE_DIR}
COLLECTION_INTERVAL=${COLLECTION_INTERVAL:-$DEFAULT_COLLECTION_INTERVAL}
DEBUG=${DEBUG:-}

# Runtime flags
RUN_MODE="once"

# Error tracking
ERRORS_TOTAL=0

handle_error() {
    local exit_code=$1
    local line_number=$2
    echo "Error: $SCRIPT_NAME failed at line $line_number with exit code $exit_code" >&2
    exit "$exit_code"
}

trap 'handle_error $? $LINENO' ERR

debug_echo() {
    if [[ -n "$DEBUG" ]]; then
        echo "[DEBUG] $*" >&2
    fi
}

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

Wickr Enterprise metrics collector for Prometheus node_exporter textfile directory.

Collects user, bot, device, security group, and network statistics from the
Wickr Enterprise Admin API and writes them as Prometheus metrics.

OPTIONS:
    --once          Run collection once and exit (default)
    --daemon        Run continuously at COLLECTION_INTERVAL
    --help, -h      Show this help message

ENVIRONMENT VARIABLES:
    WICKR_ADMIN_URL         Wickr Enterprise admin URL (required, e.g. https://wickr.example.com)
    WICKR_API_TOKEN         Admin API token (required — generate in Admin Console → API Access)
    NODE_DIR                Node exporter textfile directory (default: $DEFAULT_NODE_DIR)
    COLLECTION_INTERVAL     Seconds between collections in daemon mode (default: $DEFAULT_COLLECTION_INTERVAL)
    DEBUG                   Enable debug output

EXAMPLES:
    WICKR_ADMIN_URL=https://wickr.example.com WICKR_API_TOKEN=abc123 $SCRIPT_NAME --once
    WICKR_ADMIN_URL=https://wickr.example.com WICKR_API_TOKEN=abc123 $SCRIPT_NAME --daemon

OUTPUT:
    Writes metrics to \$NODE_DIR/textfile_collector/wickr_metrics.prom

EOF
    exit 0
}

# Parse arguments
while [[ $# -gt 0 ]]; do
    case "$1" in
        --once)     RUN_MODE="once"; shift ;;
        --daemon)   RUN_MODE="daemon"; shift ;;
        --help|-h)  show_help ;;
        *) echo "Unknown option: $1" >&2; show_help ;;
    esac
done

# Validate required variables
validate_config() {
    if [[ -z "$WICKR_ADMIN_URL" ]]; then
        echo "Error: WICKR_ADMIN_URL is required" >&2
        exit 1
    fi

    if [[ -z "$WICKR_API_TOKEN" ]]; then
        echo "Error: WICKR_API_TOKEN is required" >&2
        exit 1
    fi

    # Strip trailing slash
    WICKR_ADMIN_URL="${WICKR_ADMIN_URL%/}"

    # Check textfile collector directory
    local textfile_dir="${NODE_DIR}/textfile_collector"
    if [[ ! -d "$textfile_dir" ]]; then
        echo "Error: Textfile collector directory not found: $textfile_dir" >&2
        echo "Create it: sudo mkdir -p $textfile_dir" >&2
        exit 1
    fi
}

# API helper — makes authenticated requests to the Wickr Admin API
wickr_api() {
    local endpoint="$1"
    local url="${WICKR_ADMIN_URL}/api/v1${endpoint}"

    debug_echo "GET $url"

    local response
    response=$(curl -sf --max-time "$DEFAULT_CURL_TIMEOUT" \
        -H "Authorization: Bearer ${WICKR_API_TOKEN}" \
        -H "Accept: application/json" \
        "$url" 2>/dev/null)

    local exit_code=$?
    if [[ $exit_code -ne 0 ]]; then
        debug_echo "API call failed: $url (exit code $exit_code)"
        ERRORS_TOTAL=$((ERRORS_TOTAL + 1))
        return 1
    fi

    echo "$response"
}

# Count items from a paginated API response
# Uses jq to count array items, handles pagination via response headers
wickr_api_count() {
    local endpoint="$1"
    local url="${WICKR_ADMIN_URL}/api/v1${endpoint}"

    debug_echo "HEAD $url (count)"

    # Try to get total from response headers or count array items
    local response
    response=$(curl -sf --max-time "$DEFAULT_CURL_TIMEOUT" \
        -H "Authorization: Bearer ${WICKR_API_TOKEN}" \
        -H "Accept: application/json" \
        "$url" 2>/dev/null)

    if [[ $? -ne 0 ]]; then
        debug_echo "API count failed: $url"
        ERRORS_TOTAL=$((ERRORS_TOTAL + 1))
        echo "0"
        return 1
    fi

    # Try to count array items
    if command -v jq &>/dev/null; then
        local count
        count=$(echo "$response" | jq 'if type == "array" then length elif .data? and (.data | type == "array") then .data | length elif .total? then .total elif .count? then .count else 0 end' 2>/dev/null)
        echo "${count:-0}"
    else
        echo "0"
    fi
}

# Collect user metrics
collect_users() {
    debug_echo "Collecting user metrics..."

    local users_response
    users_response=$(wickr_api "/users") || return 1

    if command -v jq &>/dev/null && [[ -n "$users_response" ]]; then
        local total active suspended pending

        total=$(echo "$users_response" | jq 'if type == "array" then length elif .data? then (.data | length) elif .total? then .total else 0 end' 2>/dev/null)
        active=$(echo "$users_response" | jq '[if type == "array" then .[] elif .data? then .data[] else empty end | select(.status? == "active" or .status? == "Active")] | length' 2>/dev/null)
        suspended=$(echo "$users_response" | jq '[if type == "array" then .[] elif .data? then .data[] else empty end | select(.status? == "suspended" or .status? == "Suspended")] | length' 2>/dev/null)
        pending=$(echo "$users_response" | jq '[if type == "array" then .[] elif .data? then .data[] else empty end | select(.status? == "pending" or .status? == "Pending")] | length' 2>/dev/null)

        cat <<EOF
# HELP wickr_users_total Total number of Wickr users.
# TYPE wickr_users_total gauge
wickr_users_total ${total:-0}

# HELP wickr_users_active Number of active Wickr users.
# TYPE wickr_users_active gauge
wickr_users_active ${active:-0}

# HELP wickr_users_suspended Number of suspended Wickr users.
# TYPE wickr_users_suspended gauge
wickr_users_suspended ${suspended:-0}

# HELP wickr_users_pending Number of pending Wickr users.
# TYPE wickr_users_pending gauge
wickr_users_pending ${pending:-0}

EOF
    fi
}

# Collect bot metrics
collect_bots() {
    debug_echo "Collecting bot metrics..."

    local bots_response
    bots_response=$(wickr_api "/bots") || return 1

    if command -v jq &>/dev/null && [[ -n "$bots_response" ]]; then
        local total active

        total=$(echo "$bots_response" | jq 'if type == "array" then length elif .data? then (.data | length) elif .total? then .total else 0 end' 2>/dev/null)
        active=$(echo "$bots_response" | jq '[if type == "array" then .[] elif .data? then .data[] else empty end | select(.status? == "active" or .status? == "Active" or .status? == "online" or .status? == "Online")] | length' 2>/dev/null)

        cat <<EOF
# HELP wickr_bots_total Total number of Wickr bots.
# TYPE wickr_bots_total gauge
wickr_bots_total ${total:-0}

# HELP wickr_bots_active Number of active Wickr bots.
# TYPE wickr_bots_active gauge
wickr_bots_active ${active:-0}

EOF
    fi
}

# Collect security group metrics
collect_security_groups() {
    debug_echo "Collecting security group metrics..."

    local sg_response
    sg_response=$(wickr_api "/securitygroups") || return 1

    if command -v jq &>/dev/null && [[ -n "$sg_response" ]]; then
        local total
        total=$(echo "$sg_response" | jq 'if type == "array" then length elif .data? then (.data | length) elif .total? then .total else 0 end' 2>/dev/null)

        cat <<EOF
# HELP wickr_security_groups_total Total number of security groups.
# TYPE wickr_security_groups_total gauge
wickr_security_groups_total ${total:-0}

EOF

        # Per-group user counts if available
        echo "$sg_response" | jq -r '
            (if type == "array" then . elif .data? then .data else [] end)[]
            | select(.name? != null)
            | "wickr_security_group_users{group=\"\(.name)\"} \(.userCount // .user_count // 0)"
        ' 2>/dev/null | while IFS= read -r line; do
            if [[ -z "$_sg_header_printed" ]]; then
                echo "# HELP wickr_security_group_users Number of users per security group."
                echo "# TYPE wickr_security_group_users gauge"
                _sg_header_printed=1
            fi
            echo "$line"
        done
        echo ""
    fi
}

# Collect network metrics
collect_networks() {
    debug_echo "Collecting network metrics..."

    local net_response
    net_response=$(wickr_api "/networks") || return 1

    if command -v jq &>/dev/null && [[ -n "$net_response" ]]; then
        local total
        total=$(echo "$net_response" | jq 'if type == "array" then length elif .data? then (.data | length) elif .total? then .total else 0 end' 2>/dev/null)

        cat <<EOF
# HELP wickr_networks_total Total number of Wickr networks.
# TYPE wickr_networks_total gauge
wickr_networks_total ${total:-0}

EOF
    fi
}

# Collect device metrics (aggregate from users)
collect_devices() {
    debug_echo "Collecting device metrics..."

    local users_response
    users_response=$(wickr_api "/users") || return 1

    if command -v jq &>/dev/null && [[ -n "$users_response" ]]; then
        local device_count
        device_count=$(echo "$users_response" | jq '
            [if type == "array" then .[] elif .data? then .data[] else empty end
            | (.devices // .device_count // 0)
            | if type == "array" then length else . end] | add // 0
        ' 2>/dev/null)

        cat <<EOF
# HELP wickr_devices_total Total number of registered devices.
# TYPE wickr_devices_total gauge
wickr_devices_total ${device_count:-0}

EOF
    fi
}

# Check API reachability
collect_health() {
    debug_echo "Checking Wickr API health..."

    local start_time end_time duration_ms
    start_time=$(date +%s%N)

    if curl -sf --max-time "$DEFAULT_CURL_TIMEOUT" \
        -H "Authorization: Bearer ${WICKR_API_TOKEN}" \
        "${WICKR_ADMIN_URL}/api/v1/networks" &>/dev/null; then
        local up=1
    else
        local up=0
    fi

    end_time=$(date +%s%N)
    duration_ms=$(( (end_time - start_time) / 1000000 ))

    cat <<EOF
# HELP wickr_api_up Whether the Wickr Admin API is reachable (1=up, 0=down).
# TYPE wickr_api_up gauge
wickr_api_up $up

# HELP wickr_api_response_time_ms Wickr Admin API response time in milliseconds.
# TYPE wickr_api_response_time_ms gauge
wickr_api_response_time_ms $duration_ms

EOF
}

# Write collection metadata
collect_metadata() {
    cat <<EOF
# HELP wickr_exporter_errors_total Total number of errors during collection.
# TYPE wickr_exporter_errors_total counter
wickr_exporter_errors_total $ERRORS_TOTAL

# HELP wickr_exporter_last_success_timestamp Unix timestamp of last successful collection.
# TYPE wickr_exporter_last_success_timestamp gauge
wickr_exporter_last_success_timestamp $(date +%s)

EOF
}

# Main collection function
collect_all() {
    ERRORS_TOTAL=0
    local output_dir="${NODE_DIR}/textfile_collector"
    local output_file="${output_dir}/wickr_metrics.prom"
    local temp_file
    temp_file=$(mktemp "${output_file}.XXXXXX")

    debug_echo "Starting collection..."

    {
        collect_health
        collect_users
        collect_bots
        collect_security_groups
        collect_networks
        collect_devices
        collect_metadata
    } > "$temp_file" 2>/dev/null

    mv "$temp_file" "$output_file"

    debug_echo "Collection complete. Wrote to $output_file (errors: $ERRORS_TOTAL)"
}

# Main
main() {
    validate_config

    case "$RUN_MODE" in
        once)
            collect_all
            ;;
        daemon)
            echo "$SCRIPT_NAME running in daemon mode (interval: ${COLLECTION_INTERVAL}s)"
            while true; do
                collect_all
                sleep "$COLLECTION_INTERVAL"
            done
            ;;
    esac
}

main