#!/usr/bin/env bash

#########################################################################################
#### squid-exporter.sh — Prometheus metrics exporter for Squid proxy                 ####
#### Exports cache performance, hit rates, request counts, client/server stats,      ####
#### memory usage, and file descriptor utilization as Prometheus metrics              ####
#### Requires: bash 4+, squidclient or curl, jq optional                             ####
####                                                                                 ####
#### Author: Phil Connor                                                             ####
#### Contact: contact@mylinux.work                                                   ####
#### License: MIT                                                                    ####
#### Version 1.00                                                                    ####
####                                                                                 ####
#### Usage:                                                                          ####
####   ./squid-exporter.sh --http --port 9619                                        ####
####                                                                                 ####
#### See --help for all options.                                                     ####
#########################################################################################

set -uo pipefail

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

SQUID_HOST="${SQUID_HOST:-localhost}"
SQUID_PORT="${SQUID_PORT:-3128}"
SQUID_MGR_PORT="${SQUID_MGR_PORT:-3128}"
TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9619
EXPORTER_VERSION="1.00"

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

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

Export Squid proxy statistics as Prometheus metrics.

MODES:
    --textfile      Write to node_exporter textfile collector
    --http          Run HTTP server on port $HTTP_PORT
    -o, --output    Write to a specific file

OPTIONS:
    -p, --port        HTTP port (default: $HTTP_PORT)
    --squid-host      Squid host (default: $SQUID_HOST)
    --squid-port      Squid port (default: $SQUID_PORT)
    --mgr-port        Squid manager port (default: $SQUID_MGR_PORT)

ENVIRONMENT VARIABLES:
    SQUID_HOST        Squid host (default: localhost)
    SQUID_PORT        Squid port (default: 3128)
    SQUID_MGR_PORT    Squid manager port (default: 3128)

EXAMPLES:
    $0                                    # Output to stdout
    $0 --textfile                         # Write to textfile collector
    $0 --http --port 9619                 # Run HTTP server
    $0 -o /tmp/squid.prom                 # Write to custom file

NOTE:
    Uses squidclient if available, otherwise falls back to curl.
    For curl mode, ensure cachemgr_passwd is configured in squid.conf.

EOF
    exit 0
}

parse_args() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            -h|--help) show_usage ;;
            --textfile) OUTPUT_FILE="$TEXTFILE_DIR/squid.prom"; shift ;;
            --http) HTTP_MODE=true; shift ;;
            -p|--port) HTTP_PORT="$2"; shift 2 ;;
            -o|--output) OUTPUT_FILE="$2"; shift 2 ;;
            --squid-host) SQUID_HOST="$2"; shift 2 ;;
            --squid-port) SQUID_PORT="$2"; shift 2 ;;
            --mgr-port) SQUID_MGR_PORT="$2"; shift 2 ;;
            *) echo "Unknown option: $1" >&2; exit 1 ;;
        esac
    done
}

# ============================================================================
# SQUID DATA COLLECTION
# ============================================================================

# Fetch a squid manager page
# Args: $1 - manager page (e.g., "info", "counters", "5min")
# Returns: raw text output on stdout
fetch_mgr_page() {
    local page="$1"

    if command -v squidclient >/dev/null 2>&1; then
        squidclient -h "$SQUID_HOST" -p "$SQUID_PORT" "mgr:${page}" 2>/dev/null
    elif command -v curl >/dev/null 2>&1; then
        curl -sf --max-time 10 "http://${SQUID_HOST}:${SQUID_MGR_PORT}/squid-internal-mgr/${page}" 2>/dev/null
    else
        echo "ERROR: Neither squidclient nor curl found" >&2
        return 1
    fi
}

# Check if Squid is responding
# Returns: 0 if OK, 1 if error
check_squid() {
    local output
    output=$(fetch_mgr_page "info" 2>/dev/null)
    if [[ -z "$output" ]]; then
        return 1
    fi
    return 0
}

# Extract a numeric value from squid output
# Args: $1 - text data, $2 - field pattern (grep), $3 - awk field number (default 2)
# Returns: numeric value or 0
extract_value() {
    local data="$1"
    local pattern="$2"
    local field="${3:-0}"
    local result

    if [[ "$field" == "0" ]]; then
        # Take the last numeric token on the matched line
        result=$(echo "$data" | grep -i "$pattern" | head -1 | grep -oE '[0-9]+\.?[0-9]*' | tail -1)
    else
        result=$(echo "$data" | grep -i "$pattern" | head -1 | awk "{print \$$field}")
    fi

    # Strip any non-numeric characters except dot
    result=$(echo "$result" | tr -cd '0-9.')
    echo "${result:-0}"
}

# Extract a specific counter from mgr:counters output
# Format: "counter_name = value"
# Args: $1 - counters data, $2 - counter name
# Returns: numeric value or 0
extract_counter() {
    local data="$1"
    local name="$2"
    local result

    result=$(echo "$data" | grep -E "^${name} = " | head -1 | awk -F' = ' '{print $2}')
    result=$(echo "$result" | tr -cd '0-9.')
    echo "${result:-0}"
}

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

generate_metrics() {
    local script_start
    script_start=$(date +%s%N 2>/dev/null || date +%s)

    if ! check_squid; then
        cat <<EOF
# HELP squid_up Whether Squid is responding (1=up, 0=down)
# TYPE squid_up gauge
squid_up 0
EOF
        return
    fi

    # Fetch all data upfront
    local info_data counters_data fivemin_data
    info_data=$(fetch_mgr_page "info")
    counters_data=$(fetch_mgr_page "counters")
    fivemin_data=$(fetch_mgr_page "5min")

    cat <<EOF
# HELP squid_up Whether Squid is responding (1=up, 0=down)
# TYPE squid_up gauge
squid_up 1
EOF

    echo ""

    # ========================================================================
    # SERVICE TIME (from 5min averages)
    # ========================================================================

    local http_svc_time dns_svc_time icp_svc_time
    http_svc_time=$(extract_value "$fivemin_data" "HTTP Requests.*Service Time" 0)
    dns_svc_time=$(extract_value "$fivemin_data" "DNS Lookups.*Service Time" 0)
    icp_svc_time=$(extract_value "$fivemin_data" "ICP Queries.*Service Time" 0)

    # Convert from milliseconds to seconds if value seems to be in ms
    http_svc_time=$(awk "BEGIN {v=$http_svc_time; if (v > 1) printf \"%.6f\", v/1000; else printf \"%.6f\", v}")
    dns_svc_time=$(awk "BEGIN {v=$dns_svc_time; if (v > 1) printf \"%.6f\", v/1000; else printf \"%.6f\", v}")
    icp_svc_time=$(awk "BEGIN {v=$icp_svc_time; if (v > 1) printf \"%.6f\", v/1000; else printf \"%.6f\", v}")

    cat <<EOF
# HELP squid_service_time_seconds Median service time over 5 minutes
# TYPE squid_service_time_seconds gauge
squid_service_time_seconds{type="http"} $http_svc_time
squid_service_time_seconds{type="dns"} $dns_svc_time
squid_service_time_seconds{type="icp"} $icp_svc_time
EOF

    echo ""

    # ========================================================================
    # CLIENT HTTP METRICS (from counters)
    # ========================================================================

    local client_requests client_hits client_errors client_bytes_in client_bytes_out
    client_requests=$(extract_counter "$counters_data" "client_http.requests")
    client_hits=$(extract_counter "$counters_data" "client_http.hits")
    client_errors=$(extract_counter "$counters_data" "client_http.errors")
    client_bytes_in=$(extract_counter "$counters_data" "client_http.kbytes_in")
    client_bytes_out=$(extract_counter "$counters_data" "client_http.kbytes_out")

    # Convert KB to bytes
    client_bytes_in=$(awk "BEGIN {printf \"%.0f\", $client_bytes_in * 1024}")
    client_bytes_out=$(awk "BEGIN {printf \"%.0f\", $client_bytes_out * 1024}")

    cat <<EOF
# HELP squid_client_http_requests_total Total client HTTP requests
# TYPE squid_client_http_requests_total counter
squid_client_http_requests_total $client_requests

# HELP squid_client_http_hits_total Total client HTTP cache hits
# TYPE squid_client_http_hits_total counter
squid_client_http_hits_total $client_hits

# HELP squid_client_http_errors_total Total client HTTP errors
# TYPE squid_client_http_errors_total counter
squid_client_http_errors_total $client_errors

# HELP squid_client_http_bytes_in_total Total client bytes received
# TYPE squid_client_http_bytes_in_total counter
squid_client_http_bytes_in_total $client_bytes_in

# HELP squid_client_http_bytes_out_total Total client bytes sent
# TYPE squid_client_http_bytes_out_total counter
squid_client_http_bytes_out_total $client_bytes_out
EOF

    echo ""

    # ========================================================================
    # SERVER HTTP METRICS (from counters)
    # ========================================================================

    local server_requests server_errors server_bytes_in server_bytes_out
    server_requests=$(extract_counter "$counters_data" "server.all.requests")
    server_errors=$(extract_counter "$counters_data" "server.all.errors")
    server_bytes_in=$(extract_counter "$counters_data" "server.all.kbytes_in")
    server_bytes_out=$(extract_counter "$counters_data" "server.all.kbytes_out")

    server_bytes_in=$(awk "BEGIN {printf \"%.0f\", $server_bytes_in * 1024}")
    server_bytes_out=$(awk "BEGIN {printf \"%.0f\", $server_bytes_out * 1024}")

    cat <<EOF
# HELP squid_server_http_requests_total Total server (origin) HTTP requests
# TYPE squid_server_http_requests_total counter
squid_server_http_requests_total $server_requests

# HELP squid_server_http_errors_total Total server HTTP errors
# TYPE squid_server_http_errors_total counter
squid_server_http_errors_total $server_errors

# HELP squid_server_http_bytes_in_total Total server bytes received
# TYPE squid_server_http_bytes_in_total counter
squid_server_http_bytes_in_total $server_bytes_in

# HELP squid_server_http_bytes_out_total Total server bytes sent
# TYPE squid_server_http_bytes_out_total counter
squid_server_http_bytes_out_total $server_bytes_out
EOF

    echo ""

    # ========================================================================
    # CACHE HIT RATIOS (from 5min and info)
    # ========================================================================

    local hit_ratio_5min hit_ratio_60min byte_hit_ratio_5min byte_hit_ratio_60min

    # From mgr:info — look for "Request Hit Ratios" lines
    hit_ratio_5min=$(echo "$info_data" | grep -i "Request Hit Ratios" | head -1 | grep -oE '[0-9]+\.[0-9]+' | head -1)
    hit_ratio_60min=$(echo "$info_data" | grep -i "Request Hit Ratios" | head -1 | grep -oE '[0-9]+\.[0-9]+' | sed -n '2p')
    byte_hit_ratio_5min=$(echo "$info_data" | grep -i "Byte Hit Ratios" | head -1 | grep -oE '[0-9]+\.[0-9]+' | head -1)
    byte_hit_ratio_60min=$(echo "$info_data" | grep -i "Byte Hit Ratios" | head -1 | grep -oE '[0-9]+\.[0-9]+' | sed -n '2p')

    hit_ratio_5min="${hit_ratio_5min:-0}"
    hit_ratio_60min="${hit_ratio_60min:-0}"
    byte_hit_ratio_5min="${byte_hit_ratio_5min:-0}"
    byte_hit_ratio_60min="${byte_hit_ratio_60min:-0}"

    cat <<EOF
# HELP squid_cache_hit_ratio Cache request hit ratio percentage
# TYPE squid_cache_hit_ratio gauge
squid_cache_hit_ratio{period="5min"} $hit_ratio_5min
squid_cache_hit_ratio{period="60min"} $hit_ratio_60min

# HELP squid_cache_byte_hit_ratio Cache byte hit ratio percentage
# TYPE squid_cache_byte_hit_ratio gauge
squid_cache_byte_hit_ratio{period="5min"} $byte_hit_ratio_5min
squid_cache_byte_hit_ratio{period="60min"} $byte_hit_ratio_60min
EOF

    echo ""

    # ========================================================================
    # DNS METRICS (from counters)
    # ========================================================================

    local dns_lookups dns_replies
    dns_lookups=$(extract_counter "$counters_data" "dns.lookups")
    dns_replies=$(extract_counter "$counters_data" "dns.replies")

    # Fallback: try server.all if dns.lookups not found
    if [[ "$dns_lookups" == "0" ]]; then
        dns_lookups=$(extract_counter "$counters_data" "dns.svctime.count")
    fi

    cat <<EOF
# HELP squid_dns_lookups_total Total DNS lookups
# TYPE squid_dns_lookups_total counter
squid_dns_lookups_total $dns_lookups

# HELP squid_dns_replies_total Total DNS replies
# TYPE squid_dns_replies_total counter
squid_dns_replies_total $dns_replies
EOF

    echo ""

    # ========================================================================
    # MEMORY METRICS (from mgr:info)
    # ========================================================================

    local mem_total mem_free mem_cache mem_objects

    # Total memory allocated (look for "Total accounted")
    mem_total=$(echo "$info_data" | grep -i "Total accounted" | head -1 | grep -oE '[0-9]+' | head -1)
    # Free memory
    mem_free=$(echo "$info_data" | grep -i "Free Memory" | head -1 | grep -oE '[0-9]+' | head -1)
    # Memory used by cache/store
    mem_cache=$(echo "$info_data" | grep -iE "Storage Mem|StoreEntries" | head -1 | grep -oE '[0-9]+' | head -1)
    # Memory in use for objects
    mem_objects=$(echo "$info_data" | grep -i "mem_pool_allocated" | head -1 | grep -oE '[0-9]+' | head -1)

    # Convert KB values to bytes where needed
    mem_total=$(awk "BEGIN {printf \"%.0f\", ${mem_total:-0} * 1024}")
    mem_free=$(awk "BEGIN {printf \"%.0f\", ${mem_free:-0} * 1024}")
    mem_cache=$(awk "BEGIN {printf \"%.0f\", ${mem_cache:-0} * 1024}")
    mem_objects=$(awk "BEGIN {printf \"%.0f\", ${mem_objects:-0} * 1024}")

    cat <<EOF
# HELP squid_memory_usage_bytes Memory usage breakdown in bytes
# TYPE squid_memory_usage_bytes gauge
squid_memory_usage_bytes{type="total"} $mem_total
squid_memory_usage_bytes{type="free"} $mem_free
squid_memory_usage_bytes{type="cache"} $mem_cache
squid_memory_usage_bytes{type="objects"} $mem_objects
EOF

    echo ""

    # ========================================================================
    # FILE DESCRIPTOR METRICS (from mgr:info)
    # ========================================================================

    local fd_used fd_max

    fd_used=$(echo "$info_data" | grep -i "Number of file desc currently in use" | head -1 | grep -oE '[0-9]+' | tail -1)
    fd_max=$(echo "$info_data" | grep -i "Maximum number of file desc" | head -1 | grep -oE '[0-9]+' | tail -1)

    # Fallback patterns
    if [[ -z "$fd_used" || "$fd_used" == "0" ]]; then
        fd_used=$(echo "$info_data" | grep -i "file descriptor" | grep -i "used\|in use" | head -1 | grep -oE '[0-9]+' | head -1)
    fi
    if [[ -z "$fd_max" || "$fd_max" == "0" ]]; then
        fd_max=$(echo "$info_data" | grep -i "file descriptor" | grep -i "max\|available" | head -1 | grep -oE '[0-9]+' | head -1)
    fi

    fd_used="${fd_used:-0}"
    fd_max="${fd_max:-0}"

    cat <<EOF
# HELP squid_fd_used File descriptors currently in use
# TYPE squid_fd_used gauge
squid_fd_used $fd_used

# HELP squid_fd_max Maximum file descriptors available
# TYPE squid_fd_max gauge
squid_fd_max $fd_max
EOF

    echo ""

    # ========================================================================
    # CPU USAGE (from mgr:info)
    # ========================================================================

    local cpu_user cpu_system

    cpu_user=$(echo "$info_data" | grep -i "CPU Usage.*user" | head -1 | grep -oE '[0-9]+\.?[0-9]*' | head -1)
    cpu_system=$(echo "$info_data" | grep -i "CPU Usage.*system\|system.*CPU" | head -1 | grep -oE '[0-9]+\.?[0-9]*' | head -1)

    # Fallback: try separate lines
    if [[ -z "$cpu_user" ]]; then
        cpu_user=$(echo "$info_data" | grep -i "user time" | head -1 | grep -oE '[0-9]+\.?[0-9]*' | head -1)
    fi
    if [[ -z "$cpu_system" ]]; then
        cpu_system=$(echo "$info_data" | grep -i "system time" | head -1 | grep -oE '[0-9]+\.?[0-9]*' | head -1)
    fi

    cpu_user="${cpu_user:-0}"
    cpu_system="${cpu_system:-0}"

    cat <<EOF
# HELP squid_cpu_usage_seconds CPU time consumed in seconds
# TYPE squid_cpu_usage_seconds counter
squid_cpu_usage_seconds{type="user"} $cpu_user
squid_cpu_usage_seconds{type="system"} $cpu_system
EOF

    echo ""

    # ========================================================================
    # ACTIVE CONNECTIONS (from mgr:info)
    # ========================================================================

    local conn_client conn_server conn_icp

    conn_client=$(echo "$info_data" | grep -iE "Number of clients accessing cache" | head -1 | grep -oE '[0-9]+' | tail -1)
    conn_server=$(echo "$info_data" | grep -iE "Number of HTTP requests being processed" | head -1 | grep -oE '[0-9]+' | tail -1)
    conn_icp=$(echo "$info_data" | grep -iE "Number of ICP" | head -1 | grep -oE '[0-9]+' | tail -1)

    conn_client="${conn_client:-0}"
    conn_server="${conn_server:-0}"
    conn_icp="${conn_icp:-0}"

    cat <<EOF
# HELP squid_connections_active Currently active connections
# TYPE squid_connections_active gauge
squid_connections_active{type="client"} $conn_client
squid_connections_active{type="server"} $conn_server
squid_connections_active{type="icp"} $conn_icp
EOF

    echo ""

    # ========================================================================
    # SWAP / DISK CACHE SIZE (from mgr:info)
    # ========================================================================

    local swap_size

    swap_size=$(echo "$info_data" | grep -i "Storage Swap size" | head -1 | grep -oE '[0-9]+' | head -1)
    swap_size=$(awk "BEGIN {printf \"%.0f\", ${swap_size:-0} * 1024}")

    cat <<EOF
# HELP squid_swap_size_bytes Swap (disk cache) size in bytes
# TYPE squid_swap_size_bytes gauge
squid_swap_size_bytes $swap_size
EOF

    echo ""

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

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

    # Calculate duration — handle nanosecond timestamps if available
    if [[ ${#script_start} -gt 10 ]]; then
        local end_ns
        end_ns=$(date +%s%N 2>/dev/null || date +%s)
        script_duration=$(awk "BEGIN {printf \"%.3f\", ($end_ns - $script_start) / 1000000000}")
    else
        script_duration=$(( $(date +%s) - script_start ))
    fi

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

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

    echo ""
}

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

run_http_server() {
    echo "Starting Squid exporter on port $HTTP_PORT..." >&2

    if ! command -v nc >/dev/null 2>&1; then
        echo "ERROR: netcat (nc) required for HTTP mode" >&2
        exit 1
    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>Squid Exporter v${EXPORTER_VERSION}</title></head>
<body>
<h1>Squid Prometheus Exporter v${EXPORTER_VERSION}</h1>
<p><a href="/metrics">Metrics</a></p>
<p>Exports Squid proxy cache performance and operational metrics.</p>
</body>
</html>
EOF
            fi
        } | nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
    done
}

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

main() {
    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}/.squid_metrics.XXXXXX")

        if ! generate_metrics > "$temp_file" 2>/dev/null; then
            rm -f "$temp_file"
            echo "ERROR: Failed to generate metrics" >&2
            exit 1
        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"
            echo "ERROR: Metrics file too small ($file_lines lines), keeping previous" >&2
            exit 1
        fi

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

        echo "Metrics written to $OUTPUT_FILE ($file_lines lines)" >&2
    else
        generate_metrics
    fi
}

main "$@"