#!/usr/bin/env bash

#########################################################################################
#### pihole-exporter.sh — Prometheus metrics exporter for Pi-hole DNS               ####
#### Exports DNS query stats, blocking rates, top domains, client counts,           ####
#### and gravity database health as Prometheus metrics                               ####
#### Requires: bash 4+, curl, jq                                                    ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./pihole-exporter.sh --http --port 9617                                       ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -uo pipefail

# ============================================================================
# COLOR VARIABLES
# ============================================================================

RED=""
GREEN=""
YELLOW=""
BLUE=""
RESET=""

setup_colors() {
    if [[ -t 2 ]] && [[ "${TERM:-}" != "dumb" ]]; then
        RED='\033[0;31m'
        GREEN='\033[0;32m'
        YELLOW='\033[0;33m'
        BLUE='\033[0;34m'
        RESET='\033[0m'
    fi
}

# ============================================================================
# LOGGING HELPERS
# ============================================================================

log()     { echo -e "${GREEN}[INFO]${RESET} $*" >&2; }
warn()    { echo -e "${YELLOW}[WARN]${RESET} $*" >&2; }
err()     { echo -e "${RED}[ERROR]${RESET} $*" >&2; }
verbose() { [[ "${VERBOSE:-false}" == "true" ]] && echo -e "${BLUE}[DEBUG]${RESET} $*" >&2; }
die()     { err "$@"; exit 1; }

# ============================================================================
# CONFIGURATION VARIABLES
# ============================================================================

PIHOLE_URL="${PIHOLE_URL:-http://localhost/admin/api.php}"
PIHOLE_API_KEY="${PIHOLE_API_KEY:-}"
TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9617
VERBOSE=false

# ============================================================================
# HELPER FUNCTIONS
# ============================================================================

show_usage() {
    cat <<EOF
Usage: $0 [OPTIONS]

Export Pi-hole DNS statistics as Prometheus metrics.

MODES:
    --textfile      Write to node_exporter textfile collector
    --http          Run HTTP server on port $HTTP_PORT
    (default)       Print metrics to stdout

OPTIONS:
    -p, --port        HTTP port (default: 9617)
    -o, --output      Output file path
    -u, --url         Pi-hole API URL (default: $PIHOLE_URL)
    -k, --key         Pi-hole API key (or set PIHOLE_API_KEY)
    -v, --verbose     Enable verbose logging
    -h, --help        Show this help

ENVIRONMENT VARIABLES:
    PIHOLE_URL        Pi-hole Admin API URL
    PIHOLE_API_KEY    Pi-hole API authentication key

EXAMPLES:
    $0                                        # Print to stdout
    $0 --textfile                             # Write to textfile collector
    $0 --http --port 9617                     # Run HTTP server
    $0 -o /tmp/pihole.prom                    # Write to custom file
    PIHOLE_API_KEY=abc123 $0 --http           # With API key

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --textfile) OUTPUT_FILE="$TEXTFILE_DIR/pihole.prom"; shift ;;
            --http) HTTP_MODE=true; shift ;;
            -p|--port) HTTP_PORT="$2"; shift 2 ;;
            -o|--output) OUTPUT_FILE="$2"; shift 2 ;;
            -u|--url) PIHOLE_URL="$2"; shift 2 ;;
            -k|--key) PIHOLE_API_KEY="$2"; shift 2 ;;
            -v|--verbose) VERBOSE=true; shift ;;
            *) echo "Unknown option: $1" >&2; exit 1 ;;
        esac
    done
}

# Build API URL with optional authentication
# Args: $1 - query parameter (e.g., "summary", "topItems")
# Returns: full URL string
build_url() {
    local query="$1"
    local url="${PIHOLE_URL}?${query}"
    if [[ -n "$PIHOLE_API_KEY" ]]; then
        url="${url}&auth=${PIHOLE_API_KEY}"
    fi
    echo "$url"
}

# Fetch JSON from Pi-hole API
# Args: $1 - query parameter
# Returns: JSON response on stdout, empty string on error
api_fetch() {
    local query="$1"
    local url
    url=$(build_url "$query")
    verbose "Fetching: ${PIHOLE_URL}?${query}"
    curl -sf --max-time 10 "$url" 2>/dev/null
}

# Check prerequisites
# Returns: 0 if OK, 1 if error
check_deps() {
    if ! command -v curl >/dev/null 2>&1; then
        err "curl not found (required for API calls)"
        return 1
    fi

    if ! command -v jq >/dev/null 2>&1; then
        err "jq not found (required for JSON parsing)"
        return 1
    fi

    return 0
}

# Escape special characters in Prometheus label values
# Args: $1 - string to escape
# Returns: escaped string safe for Prometheus labels
prom_escape() {
    local val="$1"
    val="${val//\\/\\\\}"
    val="${val//\"/\\\"}"
    val="${val//$'\n'/}"
    echo "$val"
}

# ============================================================================
# METRIC GENERATION
# ============================================================================

# Generate all Prometheus metrics
# Returns: Prometheus text format metrics on stdout
generate_metrics() {
    local script_start
    script_start=$(date +%s%N 2>/dev/null || date +%s)

    # Check dependencies
    if ! check_deps; then
        cat <<EOF
# HELP pihole_up Pi-hole API reachability (1=up, 0=down)
# TYPE pihole_up gauge
pihole_up 0
EOF
        return
    fi

    # Fetch summary data
    local summary_json
    summary_json=$(api_fetch "summary")

    if [[ -z "$summary_json" ]]; then
        cat <<EOF
# HELP pihole_up Pi-hole API reachability (1=up, 0=down)
# TYPE pihole_up gauge
pihole_up 0
EOF
        return
    fi

    # Validate we got actual JSON with expected fields
    local status
    status=$(echo "$summary_json" | jq -r '.status // empty' 2>/dev/null)

    if [[ -z "$status" ]]; then
        cat <<EOF
# HELP pihole_up Pi-hole API reachability (1=up, 0=down)
# TYPE pihole_up gauge
pihole_up 0
EOF
        return
    fi

    cat <<EOF
# HELP pihole_up Pi-hole API reachability (1=up, 0=down)
# TYPE pihole_up gauge
pihole_up 1
EOF

    echo ""

    # ========================================================================
    # STATUS
    # ========================================================================

    local status_val=0
    if [[ "$status" == "enabled" ]]; then
        status_val=1
    fi

    cat <<EOF
# HELP pihole_status Pi-hole blocking status (1=enabled, 0=disabled)
# TYPE pihole_status gauge
pihole_status $status_val
EOF

    echo ""

    # ========================================================================
    # SUMMARY METRICS
    # ========================================================================

    local domains_blocked dns_queries_today ads_blocked_today ads_percentage_today
    local unique_domains queries_forwarded queries_cached
    local clients_ever_seen unique_clients

    domains_blocked=$(echo "$summary_json" | jq -r '.domains_being_blocked // 0' 2>/dev/null | tr -d ',')
    dns_queries_today=$(echo "$summary_json" | jq -r '.dns_queries_today // 0' 2>/dev/null | tr -d ',')
    ads_blocked_today=$(echo "$summary_json" | jq -r '.ads_blocked_today // 0' 2>/dev/null | tr -d ',')
    ads_percentage_today=$(echo "$summary_json" | jq -r '.ads_percentage_today // 0' 2>/dev/null | tr -d ',')
    unique_domains=$(echo "$summary_json" | jq -r '.unique_domains // 0' 2>/dev/null | tr -d ',')
    queries_forwarded=$(echo "$summary_json" | jq -r '.queries_forwarded // 0' 2>/dev/null | tr -d ',')
    queries_cached=$(echo "$summary_json" | jq -r '.queries_cached // 0' 2>/dev/null | tr -d ',')
    clients_ever_seen=$(echo "$summary_json" | jq -r '.clients_ever_seen // 0' 2>/dev/null | tr -d ',')
    unique_clients=$(echo "$summary_json" | jq -r '.unique_clients // 0' 2>/dev/null | tr -d ',')

    cat <<EOF
# HELP pihole_domains_blocked Number of domains on the blocklist
# TYPE pihole_domains_blocked gauge
pihole_domains_blocked $domains_blocked

# HELP pihole_dns_queries_today Total DNS queries today
# TYPE pihole_dns_queries_today gauge
pihole_dns_queries_today $dns_queries_today

# HELP pihole_ads_blocked_today Ads blocked today
# TYPE pihole_ads_blocked_today gauge
pihole_ads_blocked_today $ads_blocked_today

# HELP pihole_ads_percentage_today Percentage of queries blocked today
# TYPE pihole_ads_percentage_today gauge
pihole_ads_percentage_today $ads_percentage_today

# HELP pihole_unique_domains Unique domains queried
# TYPE pihole_unique_domains gauge
pihole_unique_domains $unique_domains

# HELP pihole_queries_forwarded Queries forwarded to upstream DNS
# TYPE pihole_queries_forwarded gauge
pihole_queries_forwarded $queries_forwarded

# HELP pihole_queries_cached Queries answered from cache
# TYPE pihole_queries_cached gauge
pihole_queries_cached $queries_cached

# HELP pihole_clients_ever_seen Total clients ever seen
# TYPE pihole_clients_ever_seen gauge
pihole_clients_ever_seen $clients_ever_seen

# HELP pihole_unique_clients Unique active clients
# TYPE pihole_unique_clients gauge
pihole_unique_clients $unique_clients
EOF

    echo ""

    # ========================================================================
    # REPLY TYPE METRICS
    # ========================================================================

    local reply_nodata reply_nxdomain reply_cname reply_ip

    reply_nodata=$(echo "$summary_json" | jq -r '.reply_NODATA // 0' 2>/dev/null | tr -d ',')
    reply_nxdomain=$(echo "$summary_json" | jq -r '.reply_NXDOMAIN // 0' 2>/dev/null | tr -d ',')
    reply_cname=$(echo "$summary_json" | jq -r '.reply_CNAME // 0' 2>/dev/null | tr -d ',')
    reply_ip=$(echo "$summary_json" | jq -r '.reply_IP // 0' 2>/dev/null | tr -d ',')

    cat <<EOF
# HELP pihole_reply_type DNS reply types
# TYPE pihole_reply_type gauge
pihole_reply_type{type="NODATA"} $reply_nodata
pihole_reply_type{type="NXDOMAIN"} $reply_nxdomain
pihole_reply_type{type="CNAME"} $reply_cname
pihole_reply_type{type="IP"} $reply_ip
EOF

    echo ""

    # ========================================================================
    # GRAVITY LAST UPDATE
    # ========================================================================

    local gravity_ts
    gravity_ts=$(echo "$summary_json" | jq -r '.gravity_last_updated.absolute // 0' 2>/dev/null)

    cat <<EOF
# HELP pihole_gravity_last_update Unix timestamp of last gravity update
# TYPE pihole_gravity_last_update gauge
pihole_gravity_last_update $gravity_ts
EOF

    echo ""

    # ========================================================================
    # QUERY TYPE METRICS
    # ========================================================================

    local querytypes_json
    querytypes_json=$(api_fetch "getQueryTypes")

    cat <<EOF
# HELP pihole_query_type Percentage of queries by DNS record type
# TYPE pihole_query_type gauge
EOF

    if [[ -n "$querytypes_json" ]]; then
        echo "$querytypes_json" | jq -r '
            .querytypes // {} | to_entries[] |
            "\(.key) \(.value)"
        ' 2>/dev/null | while read -r qtype qval; do
            [[ -z "$qtype" ]] && continue
            local clean_type="${qtype#(}"
            clean_type="${clean_type%)}"
            echo "pihole_query_type{type=\"$(prom_escape "$clean_type")\"} $qval"
        done
    fi

    echo ""

    # ========================================================================
    # TOP DOMAINS (PERMITTED)
    # ========================================================================

    local topitems_json
    topitems_json=$(api_fetch "topItems")

    cat <<EOF
# HELP pihole_top_domain_queries Top permitted domains by query count
# TYPE pihole_top_domain_queries gauge
EOF

    if [[ -n "$topitems_json" ]]; then
        echo "$topitems_json" | jq -r '
            .top_queries // {} | to_entries | sort_by(-.value) | .[0:5] | .[] |
            "\(.key) \(.value)"
        ' 2>/dev/null | while read -r domain count; do
            [[ -z "$domain" ]] && continue
            echo "pihole_top_domain_queries{domain=\"$(prom_escape "$domain")\"} $count"
        done
    fi

    echo ""

    # ========================================================================
    # TOP BLOCKED DOMAINS
    # ========================================================================

    cat <<EOF
# HELP pihole_top_blocked_queries Top blocked domains by query count
# TYPE pihole_top_blocked_queries gauge
EOF

    if [[ -n "$topitems_json" ]]; then
        echo "$topitems_json" | jq -r '
            .top_ads // {} | to_entries | sort_by(-.value) | .[0:5] | .[] |
            "\(.key) \(.value)"
        ' 2>/dev/null | while read -r domain count; do
            [[ -z "$domain" ]] && continue
            echo "pihole_top_blocked_queries{domain=\"$(prom_escape "$domain")\"} $count"
        done
    fi

    echo ""

    # ========================================================================
    # TOP CLIENTS
    # ========================================================================

    local topclients_json
    topclients_json=$(api_fetch "topClients")

    cat <<EOF
# HELP pihole_top_client_queries Top clients by query count
# TYPE pihole_top_client_queries gauge
EOF

    if [[ -n "$topclients_json" ]]; then
        echo "$topclients_json" | jq -r '
            .top_sources // {} | to_entries | sort_by(-.value) | .[0:5] | .[] |
            "\(.key) \(.value)"
        ' 2>/dev/null | while read -r client count; do
            [[ -z "$client" ]] && continue
            echo "pihole_top_client_queries{client=\"$(prom_escape "$client")\"} $count"
        done
    fi

    echo ""

    # ========================================================================
    # EXPORTER RUNTIME
    # ========================================================================

    local script_end script_duration
    script_end=$(date +%s)

    if [[ "$script_start" =~ [0-9]{10,} ]]; then
        local end_ns
        end_ns=$(date +%s%N 2>/dev/null || echo "${script_end}000000000")
        script_duration=$(( (end_ns - script_start) / 1000000000 ))
    else
        script_duration=$(( script_end - script_start ))
    fi

    cat <<EOF
# HELP pihole_exporter_duration_seconds Time to generate all metrics
# TYPE pihole_exporter_duration_seconds gauge
pihole_exporter_duration_seconds $script_duration

# HELP pihole_exporter_last_run_timestamp Unix timestamp of last successful run
# TYPE pihole_exporter_last_run_timestamp gauge
pihole_exporter_last_run_timestamp $script_end
EOF

    echo ""
}

# ============================================================================
# HTTP SERVER MODE
# ============================================================================

# Run simple HTTP server using netcat
# Serves metrics on /metrics endpoint
run_http_server() {
    log "Starting Pi-hole exporter on port $HTTP_PORT..."

    if ! command -v nc >/dev/null 2>&1; then
        die "netcat (nc) required for HTTP mode"
    fi

    while true; do
        {
            read -r request
            if [[ "$request" =~ ^GET\ /metrics ]]; then
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/plain; version=0.0.4\r\n\r"
                generate_metrics
            else
                echo -e "HTTP/1.1 200 OK\r\nContent-Type: text/html\r\n\r"
                cat <<EOF
<!DOCTYPE html>
<html>
<head><title>Pi-hole Exporter v1.00</title></head>
<body>
<h1>Pi-hole Prometheus Exporter v1.00</h1>
<p><a href="/metrics">Metrics</a></p>
<p>DNS query stats, blocking rates, top domains, client activity.</p>
</body>
</html>
EOF
            fi
        } | nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
    done
}

# ============================================================================
# MAIN EXECUTION
# ============================================================================

main() {
    setup_colors
    parse_args "$@"

    if [[ "$HTTP_MODE" == "true" ]]; then
        run_http_server
    elif [[ -n "$OUTPUT_FILE" ]]; then
        local output_dir
        output_dir="$(dirname "$OUTPUT_FILE")"
        mkdir -p "$output_dir"

        local temp_file
        temp_file=$(mktemp "${output_dir}/.pihole_metrics.XXXXXX")

        if ! generate_metrics > "$temp_file" 2>/dev/null; then
            rm -f "$temp_file"
            die "Failed to generate metrics"
        fi

        local file_lines
        file_lines=$(wc -l < "$temp_file" 2>/dev/null || echo 0)

        if [[ "$file_lines" -lt 10 ]]; then
            rm -f "$temp_file"
            die "Metrics file too small ($file_lines lines), keeping previous"
        fi

        chmod 644 "$temp_file"
        mv -f "$temp_file" "$OUTPUT_FILE"

        log "Metrics written to $OUTPUT_FILE ($file_lines lines)"
    else
        generate_metrics
    fi
}

main "$@"