#!/bin/bash

# Finds which mongo pod (mongo-0, mongo-1, mongo-2) is the primary and outputs its number.
# TLS is off by default; pass --tls (or --tls-dir / cert file flags) when the cluster requires it.
# Usage: ./get-mongo-primary-pod.sh -p <password> [options]
# Example: POD_NUMBER=$(./get-mongo-primary-pod.sh -p 'Z2zOkYikAnpQ')

set -euo pipefail

RED='\033[0;31m'
GREEN='\033[0;32m'
BLUE='\033[0;34m'
NC='\033[0m'

log_info() { echo -e "${BLUE}[INFO]${NC} $1" >&2; }
log_error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }

NAMESPACE="${NAMESPACE:-hubble-system}"
MONGO_PORT="${MONGO_PORT:-27017}"
MONGO_USER="${MONGO_USER:-root}"
AUTH_DB="${AUTH_DB:-admin}"
DEFAULT_TLS_DIR="/var/mongodb/tls"

USE_TLS=false
TLS_DIR=""
TLS_CA_FILE=""
TLS_CERT_KEY_FILE=""

usage() {
    cat << EOF >&2
Usage: $0 -p|--password <mongo_password> [options]

Required:
  -p, --password PASS    MongoDB password

Options:
  -k, --kubeconfig FILE  Path to kubeconfig
      --tls              Enable TLS (default cert paths inside pod: /var/mongodb/tls)
      --tls-dir DIR      Enable TLS; cert paths inside pod (default: /var/mongodb/tls)
      --tls-ca-file FILE CA file path inside pod (implies TLS)
      --tls-cert-key FILE Client cert+key PEM inside pod (implies TLS)
  -h, --help             Show this help

Outputs the pod number (0, 1, or 2) that is the replica set primary.

Examples:
  $0 -p 'secret' -k ./edge.kubeconfig
  $0 -p 'secret' -k ./edge.kubeconfig --tls
  $0 -p 'secret' --tls --tls-dir /var/mongodb/tls
EOF
    exit 1
}

MONGO_PASSWORD=""
KUBECONFIG_ARG=""

while [[ $# -gt 0 ]]; do
    case $1 in
        -p|--password)
            MONGO_PASSWORD="${2:?Usage: $0 -p <password>}"
            shift 2
            ;;
        -k|--kubeconfig)
            KUBECONFIG_ARG="$2"
            shift 2
            ;;
        --tls)
            USE_TLS=true
            shift 1
            ;;
        --tls-dir)
            USE_TLS=true
            TLS_DIR="$2"
            shift 2
            ;;
        --tls-ca-file)
            USE_TLS=true
            TLS_CA_FILE="$2"
            shift 2
            ;;
        --tls-cert-key)
            USE_TLS=true
            TLS_CERT_KEY_FILE="$2"
            shift 2
            ;;
        -h|--help)
            usage
            ;;
        *)
            log_error "Unknown option: $1"
            usage
            ;;
    esac
done

if [[ -z "$MONGO_PASSWORD" ]]; then
    log_error "Missing required argument: -p|--password"
    usage
fi

if [[ -n "$KUBECONFIG_ARG" ]]; then
    export KUBECONFIG="$KUBECONFIG_ARG"
fi

if [[ "$USE_TLS" == "true" ]]; then
    if [[ -z "$TLS_DIR" ]]; then
        TLS_DIR="$DEFAULT_TLS_DIR"
    fi
    if [[ -z "$TLS_CA_FILE" ]]; then
        TLS_CA_FILE="${TLS_DIR}/ca.crt"
    fi
    if [[ -z "$TLS_CERT_KEY_FILE" ]]; then
        TLS_CERT_KEY_FILE="${TLS_DIR}/tls-combined.pem"
    fi
fi

PRIMARY=""
for i in 0 1 2; do
    log_info "Checking mongo-${i} for primary role..."
    if ! kubectl -n "$NAMESPACE" get "pod/mongo-${i}" &>/dev/null; then
        log_info "mongo-${i}: pod not found, skipping"
        continue
    fi

    mongosh_args=(
        --host "mongo-${i}.mongo.${NAMESPACE}.svc.cluster.local"
        --port "$MONGO_PORT"
        --username "$MONGO_USER"
        --password "$MONGO_PASSWORD"
        --authenticationDatabase "$AUTH_DB"
    )
    if [[ "$USE_TLS" == "true" ]]; then
        mongosh_args+=(--tls --tlsCAFile "$TLS_CA_FILE" --tlsCertificateKeyFile "$TLS_CERT_KEY_FILE")
    fi
    mongosh_args+=(--quiet --eval "db.runCommand('hello').isWritablePrimary")

    exec_output=""
    exec_exit=0
    exec_output=$(kubectl -n "$NAMESPACE" exec "pod/mongo-${i}" -- env OPENSSL_CONF=/dev/null mongosh \
        "${mongosh_args[@]}" 2>&1) || exec_exit=$?

    result=$(printf '%s\n' "$exec_output" | tail -1 | tr -d '[:space:]')
    if [[ "$result" == "true" ]]; then
        log_info "mongo-${i}: isWritablePrimary=true (primary found)"
        PRIMARY="$i"
        export PRIMARY
        break
    elif [[ $exec_exit -ne 0 ]]; then
        detail=$(printf '%s\n' "$exec_output" | grep -v '^Defaulted container' | head -3 | tr '\n' '; ')
        log_info "mongo-${i}: check failed (exit ${exec_exit})${detail:+: ${detail}}"
    else
        log_info "mongo-${i}: isWritablePrimary=${result:-empty}"
    fi
done

if [[ -z "$PRIMARY" ]]; then
    log_error "No primary found among mongo-0, mongo-1, mongo-2"
    exit 1
fi

echo "$PRIMARY"
