#!/usr/bin/bash
#
# Version 1.0 03.09.2026
#
set -euo pipefail

# ===== CONFIG =====
PGHOST="lin8"
PGPORT="5432"
PGDATABASE="mydb"
PGUSER="postgres"
PGPASSWORD="changeme"
TABLE="compression_test"
ROWS=10000
REPEAT_STR="ABCDEFGHIJKLMNOPQRSTUVWXYZabcdefghijklmnopqrstuvwxyz0123456789"
REPEAT_CNT=100

export PGPASSWORD

PSQL_ARGS=("-c" "SELECT * FROM $TABLE;")
#PSQL_ARGS=("-c" "\COPY $TABLE TO '/dev/null'")
#PSQL_ARGS=("-c" "BEGIN;" "-c" "TRUNCATE TABLE $TABLE;" "-c" "\COPY $TABLE FROM '/tmp/data.dmp';" "-c" "COMMIT;")

# ===== Helper (unchanged – used only for setup) =====
run_psql() { psql -h "$PGHOST" -p "$PGPORT" -d "$PGDATABASE" -U "$PGUSER" -c "$1" -t -A; }

# ===== Ensure test data (unchanged) =====
if [[ $(run_psql "SELECT EXISTS (SELECT 1 FROM information_schema.tables WHERE table_name='$TABLE');") != "t" ]]; then
    echo "Creating table $TABLE..."
    run_psql "CREATE TABLE $TABLE (id SERIAL PRIMARY KEY, payload TEXT);
              INSERT INTO $TABLE (payload) SELECT repeat('$REPEAT_STR', $REPEAT_CNT) FROM generate_series(1,$ROWS);" >/dev/null
    run_psql "\COPY $TABLE TO '/tmp/data.dmp'" >/dev/null
else
    echo "Table $TABLE already exists."
fi

# ===== Measure =====
measure() {
    local label="$1"
    local comp="$2"
    echo "--- $label ---"

    local time_file=$(mktemp)
    local ss_file=$(mktemp)

    # Resolve PGHOST to an IP address (numeric)
    local PGHOST_IP=$(getent hosts "$PGHOST" | awk '{print $1}' | head -1)
    if [[ -z "$PGHOST_IP" ]]; then
        # fallback: use host command (if getent not available)
        PGHOST_IP=$(host "$PGHOST" 2>/dev/null | head -1 | awk '{print $NF}')
    fi
    if [[ -z "$PGHOST_IP" ]]; then
        # final fallback: use the hostname as given (might not match ss -n)
        PGHOST_IP="$PGHOST"
    fi

    # Run psql with the query, capture time, and write ss output to file
    {
        time -p PGPASSWORD=$PGPASSWORD PGCOMPRESSION=$comp \
            psql -h "$PGHOST" -p "$PGPORT" -d "$PGDATABASE" -U "$PGUSER" \
                "${PSQL_ARGS[@]}" \
                -c "\! sleep 1" \
                -c "\! ss -t -i -p -n > $ss_file"
    } 2> "$time_file" > /dev/null

    # Now extract counters from the ss file using the remote endpoint
    local psql_block=$(grep -A10 "$PGHOST_IP:$PGPORT" "$ss_file" | grep -A1 "psql")
    local bytes_received=$(echo "$psql_block" | grep -oE 'bytes_received:[0-9]+' | head -1 | cut -d: -f2)
    local bytes_sent=$(echo "$psql_block" | grep -oE 'bytes_sent:[0-9]+' | head -1 | cut -d: -f2)
    local runtime=$(grep '^real' "$time_file" | awk '{print $2}')

    # Print results
    echo "  psql network data received: $(awk "BEGIN {printf \"%.2f\", ${bytes_received:-0}/1024}") KB"
    echo "  psql network data sent:     $(awk "BEGIN {printf \"%.2f\", ${bytes_sent:-0}/1024}") KB"
    echo "  runtime:                    ${runtime:-N/A} seconds"
    echo

    # Clean up
    rm -f "$time_file" "$ss_file"
}

# ===== Run measurements =====
measure "without compression" "off"
measure "with zstd compression" "zstd"