#!/bin/bash
################################################################################
# Script Name: postgresql-exporter.sh
# Version: 1.0
# Description: Prometheus textfile exporter for PostgreSQL — connections, TPS,
#              cache hit ratio, database sizes, dead tuples, locks, WAL
#              replication, and checkpoint stats via psql
#
# Author: Phil Connor
# Contact: contact@mylinux.work
# Website: https://mylinux.work
# License: MIT
#
# Prerequisites:
#   - psql CLI client
#   - Credentials via ~/.pgpass or PGPASSWORD env
#   - Network access to PostgreSQL server
#   - netcat (nc) for HTTP mode
#
# Usage:
#   ./postgresql-exporter.sh
#   ./postgresql-exporter.sh --http -p 9208
#   ./postgresql-exporter.sh --textfile
#
# Configuration:
#   Default HTTP port: 9208
#   Textfile directory: /var/lib/node_exporter
#
################################################################################

TEXTFILE_DIR="/var/lib/node_exporter"
OUTPUT_FILE=""
HTTP_MODE=false
HTTP_PORT=9208
PGHOST="${PGHOST:-127.0.0.1}"
PGPORT="${PGPORT:-5432}"
PGUSER="${PGUSER:-}"
PGDATABASE="${PGDATABASE:-postgres}"

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

Export PostgreSQL statistics as Prometheus metrics.

MODES:
    --textfile      Write to node_exporter textfile collector
    --http          Run HTTP server on port $HTTP_PORT

OPTIONS:
    -p, --port      HTTP port (default: 9208)
    -o, --output    Output file path

EXAMPLES:
    $0 --textfile
    $0 --http --port 9208
    PGUSER=monitoring PGPASSWORD=secret $0 --textfile

EOF
    exit 0
}

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

build_psql_cmd() {
    local cmd="psql -h $PGHOST -p $PGPORT"
    [ -n "$PGUSER" ] && cmd="$cmd -U $PGUSER"
    cmd="$cmd -d $PGDATABASE -t -A -F'|'"
    echo "$cmd"
}

check_postgresql() {
    if ! command -v psql >/dev/null 2>&1; then
        echo "ERROR: psql client not found" >&2
        return 1
    fi
    local psql_cmd
    psql_cmd=$(build_psql_cmd)
    if ! $psql_cmd -c "SELECT 1" >/dev/null 2>&1; then
        echo "ERROR: Cannot connect to PostgreSQL server" >&2
        return 1
    fi
    return 0
}

psql_query() {
    local query="$1"
    local psql_cmd
    psql_cmd=$(build_psql_cmd)
    $psql_cmd -c "$query" 2>/dev/null
}

generate_metrics() {
    local script_start
    script_start=$(date +%s)

    if ! check_postgresql; then
        cat <<EOF
# HELP postgresql_overview_up PostgreSQL exporter status
# TYPE postgresql_overview_up gauge
postgresql_overview_up 0
EOF
        return
    fi

    local version uptime max_conn active_conn
    version=$(psql_query "SELECT version()" | head -1 | awk '{print $2}')
    uptime=$(psql_query "SELECT EXTRACT(EPOCH FROM now() - pg_postmaster_start_time())::int" | head -1)
    max_conn=$(psql_query "SHOW max_connections" | head -1)
    active_conn=$(psql_query "SELECT count(*) FROM pg_stat_activity WHERE state = 'active'" | head -1)

    local txn_committed txn_rolledback
    txn_committed=$(psql_query "SELECT sum(xact_commit) FROM pg_stat_database" | head -1)
    txn_rolledback=$(psql_query "SELECT sum(xact_rollback) FROM pg_stat_database" | head -1)

    local cache_hit_ratio
    cache_hit_ratio=$(psql_query "SELECT CASE WHEN blks_hit + blks_read = 0 THEN 1.0 ELSE blks_hit::float / (blks_hit + blks_read) END FROM pg_stat_database WHERE datname = current_database()" | head -1)

    local dead_tuples
    dead_tuples=$(psql_query "SELECT sum(n_dead_tup) FROM pg_stat_user_tables" | head -1)

    local longest_query
    longest_query=$(psql_query "SELECT COALESCE(EXTRACT(EPOCH FROM max(now() - query_start))::int, 0) FROM pg_stat_activity WHERE state = 'active' AND query NOT LIKE 'autovacuum%'" | head -1)

    local wal_size
    wal_size=$(psql_query "SELECT pg_wal_lsn_diff(pg_current_wal_lsn(), '0/0')::bigint" 2>/dev/null | head -1)

    local temp_bytes
    temp_bytes=$(psql_query "SELECT sum(temp_bytes) FROM pg_stat_database" | head -1)

    local checkpoints
    checkpoints=$(psql_query "SELECT checkpoints_timed + checkpoints_req FROM pg_stat_bgwriter" | head -1)

    local checkpoint_write_time
    checkpoint_write_time=$(psql_query "SELECT checkpoint_write_time / 1000.0 FROM pg_stat_bgwriter" | head -1)

    cat <<EOF
# HELP postgresql_overview_up PostgreSQL exporter status
# TYPE postgresql_overview_up gauge
postgresql_overview_up 1

# HELP postgresql_overview_version_info PostgreSQL server version
# TYPE postgresql_overview_version_info gauge
postgresql_overview_version_info{version="${version:-unknown}"} 1

# HELP postgresql_overview_uptime_seconds PostgreSQL server uptime
# TYPE postgresql_overview_uptime_seconds gauge
postgresql_overview_uptime_seconds ${uptime:-0}

# HELP postgresql_overview_active_connections Current active connections
# TYPE postgresql_overview_active_connections gauge
postgresql_overview_active_connections ${active_conn:-0}

# HELP postgresql_overview_max_connections Maximum allowed connections
# TYPE postgresql_overview_max_connections gauge
postgresql_overview_max_connections ${max_conn:-0}

# HELP postgresql_overview_transactions_committed_total Total committed transactions
# TYPE postgresql_overview_transactions_committed_total counter
postgresql_overview_transactions_committed_total ${txn_committed:-0}

# HELP postgresql_overview_transactions_rolledback_total Total rolled back transactions
# TYPE postgresql_overview_transactions_rolledback_total counter
postgresql_overview_transactions_rolledback_total ${txn_rolledback:-0}

# HELP postgresql_overview_cache_hit_ratio Cache hit ratio
# TYPE postgresql_overview_cache_hit_ratio gauge
postgresql_overview_cache_hit_ratio ${cache_hit_ratio:-0}

# HELP postgresql_overview_dead_tuples Dead tuples needing vacuum
# TYPE postgresql_overview_dead_tuples gauge
postgresql_overview_dead_tuples ${dead_tuples:-0}

# HELP postgresql_overview_longest_query_duration_seconds Duration of longest running query
# TYPE postgresql_overview_longest_query_duration_seconds gauge
postgresql_overview_longest_query_duration_seconds ${longest_query:-0}

# HELP postgresql_overview_wal_size_bytes WAL size in bytes
# TYPE postgresql_overview_wal_size_bytes gauge
postgresql_overview_wal_size_bytes ${wal_size:-0}

# HELP postgresql_overview_temp_bytes Temp file usage in bytes
# TYPE postgresql_overview_temp_bytes gauge
postgresql_overview_temp_bytes ${temp_bytes:-0}

# HELP postgresql_overview_checkpoints_total Total checkpoints
# TYPE postgresql_overview_checkpoints_total counter
postgresql_overview_checkpoints_total ${checkpoints:-0}

# HELP postgresql_overview_checkpoint_write_time_seconds Checkpoint write time
# TYPE postgresql_overview_checkpoint_write_time_seconds gauge
postgresql_overview_checkpoint_write_time_seconds ${checkpoint_write_time:-0}

EOF

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

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

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

EOF
}

run_http_server() {
    echo "Starting PostgreSQL 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"
                echo "<html><head><title>PostgreSQL Exporter</title></head><body><h1>PostgreSQL Prometheus Exporter</h1><p><a href=\"/metrics\">Metrics</a></p></body></html>"
            fi
        } | nc -l -p "$HTTP_PORT" -q 1 2>/dev/null
    done
}

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}/.postgresql_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 "$@"