#!/bin/bash

# Configuration
ENV="stg"
PORT=5432

# Parse arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -s|--stg) ENV="stg" ;;
        -p|--prd) ENV="prd" ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done

# Function to get IPs based on environment
get_ips() {
    if [ "$ENV" == "stg" ]; then
        # Sourcing activate-stg to ensure KUBECONFIG is correct
        if [ -f "$HOME/bin/activate-stg" ]; then
            source "$HOME/bin/activate-stg" > /dev/null
        fi
    else
        # Placeholder for production kubeconfig activation
        echo "❌ Error: Production kubeconfig activation not yet defined in this script."
        exit 1
    fi

    kubectl get nodes -o jsonpath='{range .items[*]}{.status.addresses[?(@.type=="ExternalIP")].address}{"\n"}{end}'
}

echo "🔍 Fetching Kubernetes Worker IPs for $ENV..."
IPS=$(get_ips)

if [ -z "$IPS" ]; then
    echo "❌ Error: No IPs found. Are you authenticated to the cluster?"
    exit 1
fi

echo ""
echo "==============================================================="
echo "📋 IPTABLES RULES FOR $(echo $ENV | tr '[:lower:]' '[:upper:]') POSTGRESQL ACCESS"
echo "==============================================================="
echo "# Port: $PORT"
echo "# Generated on: $(date)"
echo ""

for IP in $IPS; do
    # Skip IPv6 for old iptables
    if [[ $IP =~ .*:.* ]]; then
        continue
    fi
    echo "iptables -A INPUT -p tcp -s $IP --dport $PORT -j ACCEPT"
done

echo ""
echo "==============================================================="
echo "💡 INSTRUCTIONS:"
echo "1. SSH into the 'int' server: ssh root@72.14.176.182"
echo "2. Copy/Paste the rules above into the terminal."
echo "3. Verify with: iptables -L -n | grep $PORT"
echo "==============================================================="