#!/usr/bin/env bash
# directory-size-exporter.sh — Prometheus exporter for directory sizes
#
# Monitors directory disk usage that node_exporter can't see.
# Node exporter only reports mounted filesystem totals — this script
# tracks individual directories like /var/log, /home, /opt, or any
# path you care about.
#
# Author:  Phil Connor
# Contact: contact@mylinux.work
# License: MIT
# Version: 1.0.1

set -euo pipefail

EXPORTER_NAME="directory_size"
DEFAULT_PORT=9101
OUTPUT_MODE="stdout"
OUTPUT_FILE=""
PORT="${DIRECTORY_SIZE_PORT:-$DEFAULT_PORT}"
TIMEOUT="${DIRECTORY_SIZE_TIMEOUT:-300}"
VERBOSE=false
QUIET=false
DRY_RUN=false
TARGET_DIRECTORIES=()

# ── Metrics Collection ──────────────────────────────────────────────

log_verbose() {
    [[ "$VERBOSE" == true ]] && echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" >&2 || true
}

log_info() {
    [[ "$QUIET" == false ]] && echo "[$(date '+%Y-%m-%d %H:%M:%S')] $*" >&2 || true
}

collect_metrics() {
    local start_time
    start_time=$(date +%s%N)

    local success=1
    local size_lines="" pct_lines=""

    for directory in "${TARGET_DIRECTORIES[@]}"; do
        log_verbose "Running du for: $directory"

        local du_output
        du_output=$(timeout "$TIMEOUT" du --block-size=1 --summarize "$directory" 2>/dev/null) || {
            log_info "WARNING: du failed for $directory"
            success=0
            continue
        }

        local size_bytes
        size_bytes=$(echo "$du_output" | awk '{print $1}')
        size_lines+="node_directory_size_bytes{directory=\"${directory}\"} ${size_bytes}"$'\n'

        local pct
        pct=$(df --output=pcent "$directory" 2>/dev/null | tail -n 1 | tr -d ' %')
        if [[ "$pct" =~ ^[0-9]+$ ]]; then
            pct_lines+="node_directory_filesystem_usage_percent{directory=\"${directory}\"} ${pct}"$'\n'
        fi
    done

    echo "# HELP node_directory_size_bytes Disk space used by directory"
    echo "# TYPE node_directory_size_bytes gauge"
    printf "%s" "$size_lines"

    echo ""
    echo "# HELP node_directory_filesystem_usage_percent Filesystem usage percentage for the directory mount point"
    echo "# TYPE node_directory_filesystem_usage_percent gauge"
    printf "%s" "$pct_lines"

    # ── Script runtime ──
    local end_time runtime
    end_time=$(date +%s%N)
    runtime=$(awk "BEGIN {printf \"%.3f\", ($end_time - $start_time) / 1000000000}")

    echo ""
    echo "# HELP ${EXPORTER_NAME}_duration_seconds Script execution time"
    echo "# TYPE ${EXPORTER_NAME}_duration_seconds gauge"
    echo "${EXPORTER_NAME}_duration_seconds ${runtime}"

    echo ""
    echo "# HELP ${EXPORTER_NAME}_last_run_timestamp Last successful run"
    echo "# TYPE ${EXPORTER_NAME}_last_run_timestamp gauge"
    echo "${EXPORTER_NAME}_last_run_timestamp $(date +%s)"

    echo ""
    echo "# HELP ${EXPORTER_NAME}_success Whether the exporter ran successfully"
    echo "# TYPE ${EXPORTER_NAME}_success gauge"
    echo "${EXPORTER_NAME}_success ${success}"
}

# ── HTTP Request Handler ────────────────────────────────────────────

handle_request() {
    read -r method path version

    while IFS= read -r header; do
        [[ "$header" == $'\r' || -z "$header" ]] && break
    done

    if [[ "$path" == "/metrics" ]]; then
        local metrics length
        metrics=$(collect_metrics)
        length=${#metrics}

        printf "HTTP/1.1 200 OK\r\n"
        printf "Content-Type: text/plain; version=0.0.4; charset=utf-8\r\n"
        printf "Content-Length: %d\r\n" "$length"
        printf "Connection: close\r\n"
        printf "\r\n"
        printf "%s" "$metrics"
    else
        local body="404 Not Found"
        printf "HTTP/1.1 404 Not Found\r\n"
        printf "Content-Type: text/plain\r\n"
        printf "Content-Length: %d\r\n" "${#body}"
        printf "Connection: close\r\n"
        printf "\r\n"
        printf "%s" "$body"
    fi
}

# ── Help ─────────────────────────────────────────────────────────────

show_help() {
    cat <<EOF
Usage: $0 [OPTIONS] <directory> [directory2 ...]

Monitor directory sizes for Prometheus. Node exporter only reports
mounted filesystem totals — this script tracks individual directories.

Output modes:
  (default)        Print metrics to stdout
  --textfile       Write to node_exporter textfile collector
  -o FILE          Write to a specific file
  --http           Run as HTTP server (default port: ${DEFAULT_PORT})

Options:
  --port PORT      HTTP listen port (default: ${DEFAULT_PORT})
  --timeout SECS   du command timeout (default: 300)
  --dry-run        Show what would be written without writing
  --verbose, -v    Enable verbose debug output
  --quiet, -q      Suppress non-error output
  -h, --help       Show this help message

Environment variables:
  DIRECTORY_SIZE_PORT      HTTP listen port (default: ${DEFAULT_PORT})
  DIRECTORY_SIZE_TIMEOUT   du command timeout in seconds (default: 300)

Examples:
  $0 /var/log /home /opt
  $0 --textfile /var/log /var/lib/mysql
  $0 --http --port 9101 /var/log /home
  $0 -o /tmp/dir_sizes.prom /var/log
EOF
}

# ── Argument Parsing ────────────────────────────────────────────────

while [[ $# -gt 0 ]]; do
    case "$1" in
        --textfile)
            OUTPUT_MODE="textfile"
            shift
            ;;
        -o)
            OUTPUT_MODE="file"
            OUTPUT_FILE="$2"
            shift 2
            ;;
        --http)
            OUTPUT_MODE="http"
            shift
            ;;
        --port)
            PORT="$2"
            shift 2
            ;;
        --timeout)
            TIMEOUT="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        --verbose|-v)
            VERBOSE=true
            shift
            ;;
        --quiet|-q)
            QUIET=true
            shift
            ;;
        --handle-request)
            OUTPUT_MODE="handle-request"
            shift
            ;;
        -h|--help)
            show_help
            exit 0
            ;;
        -*)
            echo "Unknown option: $1" >&2
            exit 1
            ;;
        *)
            TARGET_DIRECTORIES+=("$1")
            shift
            ;;
    esac
done

# Validate directories
if [[ ${#TARGET_DIRECTORIES[@]} -eq 0 ]]; then
    echo "Error: at least one directory argument is required" >&2
    echo "Run with --help for usage" >&2
    exit 1
fi

for dir in "${TARGET_DIRECTORIES[@]}"; do
    if [[ ! -d "$dir" ]]; then
        echo "Error: directory does not exist: $dir" >&2
        exit 1
    fi
    if [[ ! -r "$dir" ]]; then
        echo "Error: directory is not readable: $dir" >&2
        exit 1
    fi
done

# ── Output ──────────────────────────────────────────────────────────

if [[ "$DRY_RUN" == true ]]; then
    log_info "DRY RUN — metrics that would be written:"
    collect_metrics
    exit 0
fi

case "$OUTPUT_MODE" in
    handle-request)
        handle_request
        exit 0
        ;;
    stdout)
        collect_metrics
        ;;
    textfile)
        output_dir="/var/lib/node_exporter"
        OUTPUT_FILE="${output_dir}/${EXPORTER_NAME}.prom"
        mkdir -p "$output_dir"
        temp_file=$(mktemp "${output_dir}/.${EXPORTER_NAME}.XXXXXX")
        collect_metrics > "$temp_file"
        chmod 644 "$temp_file"
        mv -f "$temp_file" "$OUTPUT_FILE"
        ;;
    file)
        temp_file=$(mktemp "${OUTPUT_FILE}.XXXXXX")
        collect_metrics > "$temp_file"
        chmod 644 "$temp_file"
        mv -f "$temp_file" "$OUTPUT_FILE"
        ;;
    http)
        if ! command -v socat &>/dev/null; then
            echo "ERROR: socat is required for --http mode" >&2
            echo "Install it: apt install socat  or  dnf install socat" >&2
            exit 1
        fi
        echo "${EXPORTER_NAME} listening on port ${PORT}..."
        echo "Monitoring directories: ${TARGET_DIRECTORIES[*]}"
        socat TCP-LISTEN:"$PORT",reuseaddr,fork EXEC:"$0 --handle-request ${TARGET_DIRECTORIES[*]}"
        ;;
esac