#!/bin/bash

# Runs a MongoDB backup using the *local* mongodump binary.
# It port-forwards to the MongoDB primary pod in the cluster, then runs mongodump
# against localhost and writes output to a local directory.
# Pass --tls-dir (or --tls-ca-file / --tls-cert-key) to enable TLS/SSL.

set -euo pipefail

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

log_info() { echo -e "${BLUE}[INFO]${NC} $1"; }
log_success() { echo -e "${GREEN}[SUCCESS]${NC} $1"; }
log_warning() { echo -e "${YELLOW}[WARNING]${NC} $1"; }
log_error() { echo -e "${RED}[ERROR]${NC} $1" >&2; }

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

NAMESPACE="${NAMESPACE:-hubble-system}"
MONGO_USER="${MONGO_USER:-root}"
AUTH_DB="${AUTH_DB:-admin}"
REMOTE_PORT="${REMOTE_PORT:-27017}"

LOCAL_HOST="${LOCAL_HOST:-127.0.0.1}"
LOCAL_PORT="${LOCAL_PORT:-27017}"

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

ALLOW_INSECURE_TLS="${ALLOW_INSECURE_TLS:-true}"

POD_NUMBER="${POD_NUMBER:-}"
LOCAL_DIR="${LOCAL_DIR:-./spectrodump}"
DATABASES=(hubbledb hubble_timeseriesdb hubble_archivedb)
NUM_PARALLEL_COLLECTIONS="${NUM_PARALLEL_COLLECTIONS:-1}"
KUBECONFIG_ARG=""
MONGO_PASSWORD=""

usage() {
    cat << EOF
Backup MongoDB using local mongodump (TLS optional)

This script port-forwards to the MongoDB primary pod (mongo-0/1/2) and runs your local
mongodump. TLS is off by default; pass --tls-dir to enable TLS/SSL. When enabled, the
TLS directory should contain (e.g. from extract-mongo-tls-certs.sh):
  - ca.crt
  - tls-combined.pem

Usage: $0 -p|--password PASS [options]

Required:
  -p, --password PASS        MongoDB password

Options:
  -k, --kubeconfig FILE      Path to kubeconfig (default: use KUBECONFIG env or default)
  -l, --local-dir DIR        Local output directory for mongodump (default: ./spectrodump)
      --tls-dir DIR          Enable TLS; directory with ca.crt and tls-combined.pem
      --tls-ca-file FILE     CA file path (implies TLS; default: <tls-dir>/ca.crt)
      --tls-cert-key FILE    Client cert+key PEM (implies TLS; default: <tls-dir>/tls-combined.pem)
  -N, --pod-number NUM       Mongo pod number (0/1/2). If omitted, auto-discovers primary.
      --local-port PORT      Local port to bind for port-forward (default: 27017)
      --remote-port PORT     Remote mongo port in pod (default: 27017)
      --num-parallel N       mongodump --numParallelCollections (default: 1)
      --strict-tls           Do NOT pass --tlsInsecure (default: insecure)
  -h, --help                 Show this help

Env overrides:
  NAMESPACE, MONGO_USER, AUTH_DB, LOCAL_HOST, LOCAL_PORT, REMOTE_PORT,
  ALLOW_INSECURE_TLS, NUM_PARALLEL_COLLECTIONS

Examples:
  $0 --kubeconfig ./edge.kubeconfig -p 'secret' -l ./bk1
  $0 --kubeconfig ./edge.kubeconfig -p 'secret' -l ./bk1 --tls-dir ./bk1/mongo-tls
EOF
    exit 1
}

require_cmd() {
    local cmd="$1"
    if ! command -v "$cmd" &>/dev/null; then
        log_error "Required command not found: $cmd"
        exit 1
    fi
}

cleanup_pf() {
    if [[ -n "${PF_PID:-}" ]]; then
        kill "$PF_PID" &>/dev/null || true
        wait "$PF_PID" &>/dev/null || true
        unset PF_PID
    fi
}

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
            ;;
        -l|--local-dir)
            LOCAL_DIR="$2"
            shift 2
            ;;
        --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
            ;;
        -N|--pod-number)
            POD_NUMBER="$2"
            shift 2
            ;;
        --local-port)
            LOCAL_PORT="$2"
            shift 2
            ;;
        --remote-port)
            REMOTE_PORT="$2"
            shift 2
            ;;
        --strict-tls)
            ALLOW_INSECURE_TLS="false"
            shift 1
            ;;
        --num-parallel)
            NUM_PARALLEL_COLLECTIONS="$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

require_cmd kubectl
require_cmd mongodump

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

if [[ "$USE_TLS" == "true" ]]; then
    if [[ -z "$TLS_DIR" && ( -z "$TLS_CA_FILE" || -z "$TLS_CERT_KEY_FILE" ) ]]; then
        log_error "TLS enabled but --tls-dir is required when --tls-ca-file or --tls-cert-key is not both set"
        exit 1
    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
    if [[ ! -f "$TLS_CA_FILE" ]]; then
        log_error "TLS CA file not found: $TLS_CA_FILE"
        exit 1
    fi
    if [[ ! -f "$TLS_CERT_KEY_FILE" ]]; then
        log_error "TLS cert+key PEM not found: $TLS_CERT_KEY_FILE"
        exit 1
    fi
fi

if [[ -z "$POD_NUMBER" ]]; then
    if [[ ! -x "${SCRIPT_DIR}/get-mongo-primary-pod.sh" ]]; then
        log_error "Primary discovery script not found or not executable: ${SCRIPT_DIR}/get-mongo-primary-pod.sh"
        exit 1
    fi
    log_info "Discovering MongoDB primary pod number..."
    POD_NUMBER="$("${SCRIPT_DIR}/get-mongo-primary-pod.sh" -p "$MONGO_PASSWORD"${KUBECONFIG_ARG:+ -k "$KUBECONFIG_ARG"})"
fi

POD="mongo-${POD_NUMBER}"

trap cleanup_pf EXIT INT TERM

log_info "Starting port-forward to $NAMESPACE/$POD ${LOCAL_HOST}:${LOCAL_PORT} -> ${REMOTE_PORT} ..."
kubectl -n "$NAMESPACE" port-forward "pod/${POD}" "${LOCAL_PORT}:${REMOTE_PORT}" >/dev/null 2>&1 &
PF_PID=$!

# Simple readiness check (avoid hanging forever on a bad forward)
for _ in {1..20}; do
    if kill -0 "$PF_PID" 2>/dev/null; then
        if nc -z "$LOCAL_HOST" "$LOCAL_PORT" >/dev/null 2>&1; then
            break
        fi
    else
        log_error "port-forward process exited unexpectedly"
        exit 1
    fi
    sleep 0.25
done

if ! nc -z "$LOCAL_HOST" "$LOCAL_PORT" >/dev/null 2>&1; then
    log_warning "Port-forward may not be ready yet; proceeding anyway."
fi

mkdir -p "$LOCAL_DIR"

if [[ "$USE_TLS" == "true" ]]; then
    log_info "Running local mongodump to $LOCAL_DIR (TLS enabled) ..."
else
    log_info "Running local mongodump to $LOCAL_DIR (TLS disabled) ..."
fi

dump_args=(
  --host "$LOCAL_HOST"
  --port "$LOCAL_PORT"
  --username "$MONGO_USER"
  --password "$MONGO_PASSWORD"
  --authenticationDatabase "$AUTH_DB"
)

if [[ "$USE_TLS" == "true" ]]; then
    # Legacy SSL flags (supported) plus optional --tlsInsecure for localhost port-forward.
    dump_args+=(--ssl --sslCAFile "$TLS_CA_FILE" --sslPEMKeyFile "$TLS_CERT_KEY_FILE")
    if [[ "${ALLOW_INSECURE_TLS}" == "true" ]]; then
        dump_args+=(--tlsInsecure)
    fi
fi

mongodump \
  "${dump_args[@]}" \
  --numParallelCollections "$NUM_PARALLEL_COLLECTIONS" \
  --gzip \
  --out="$LOCAL_DIR"

log_success "mongodump completed successfully: $LOCAL_DIR"

missing_required=0
for dbname in "${DATABASES[@]}"; do
    log_info "Checking that dump contains database directory: $dbname"
    if [[ ! -d "${LOCAL_DIR}/${dbname}" ]]; then
        if [[ "$dbname" == "hubble_archivedb" ]]; then
            log_warning "Optional database dump directory missing: ${LOCAL_DIR}/${dbname}"
            continue
        fi
        log_error "Dump verification failed: expected directory not found: ${LOCAL_DIR}/${dbname}"
        missing_required=1
        continue
    fi
    log_success "Dump verified: directory '${dbname}' found in ${LOCAL_DIR}"
done

if [[ "$missing_required" -ne 0 ]]; then
    exit 1
fi

log_success "Backup verified successfully (hubble_archivedb may be absent on some clusters)"

