#!/bin/bash

# Runs a MongoDB restore using the *local* mongorestore binary.
# It port-forwards to the MongoDB primary pod in the cluster, then runs mongorestore
# against localhost and reads input from a local dump 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; }

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)
KUBECONFIG_ARG=""
MONGO_PASSWORD=""

usage() {
    local exit_code="${1:-1}"
    cat << EOF
Restore MongoDB using local mongorestore (TLS optional)

This script port-forwards to the MongoDB primary pod (mongo-0/1/2) and runs your local
mongorestore. 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 -N|--pod-number NUM [options]

Required:
  -p, --password PASS        MongoDB password
  -N, --pod-number NUM       Mongo pod number (0/1/2), e.g. 0 for mongo-0

Options:
  -k, --kubeconfig FILE      Path to kubeconfig (default: use KUBECONFIG env or default)
  -l, --local-dir DIR        Local dump directory to restore from (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)
      --local-port PORT      Local port to bind for port-forward (default: 27017)
      --remote-port PORT     Remote mongo port in pod (default: 27017)
      --strict-tls           Do NOT pass --tlsInsecure (default: insecure)
      --no-drop              Do NOT pass --drop (default: drop enabled)
  -h, --help                 Show this help

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

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

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
}

DROP_FLAG="--drop"

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
            ;;
        --no-drop)
            DROP_FLAG=""
            shift 1
            ;;
        -h|--help)
            usage 0
            ;;
        *)
            log_error "Unknown option: $1"
            usage
            ;;
    esac
done

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

require_cmd kubectl
require_cmd mongorestore
require_cmd nc

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 [[ ! -d "$LOCAL_DIR" ]]; then
    log_error "Local dump directory not found: $LOCAL_DIR"
    exit 1
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=$!

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

# Legacy SSL flags (supported) plus optional --tlsInsecure for localhost port-forward.
restore_base_args=(
  --host "$LOCAL_HOST"
  --port "$LOCAL_PORT"
  --username "$MONGO_USER"
  --password "$MONGO_PASSWORD"
  --authenticationDatabase "$AUTH_DB"
)
if [[ "$USE_TLS" == "true" ]]; then
    restore_base_args+=(--ssl --sslCAFile "$TLS_CA_FILE" --sslPEMKeyFile "$TLS_CERT_KEY_FILE")
    if [[ "${ALLOW_INSECURE_TLS}" == "true" ]]; then
        restore_base_args+=(--tlsInsecure)
    fi
fi

if [[ "$USE_TLS" == "true" ]]; then
    log_info "Running local mongorestore from $LOCAL_DIR (TLS enabled) ..."
else
    log_info "Running local mongorestore from $LOCAL_DIR (TLS disabled) ..."
fi
missing_required=0
for db in "${DATABASES[@]}"; do
    if [[ ! -d "${LOCAL_DIR}/${db}" ]]; then
        if [[ "$db" == "hubble_archivedb" ]]; then
            log_warning "Optional database dump directory missing: ${LOCAL_DIR}/${db}; skipping restore"
            continue
        fi
        log_error "Required database dump directory missing: ${LOCAL_DIR}/${db}"
        missing_required=1
        continue
    fi

    log_info "Restoring database: $db"
    mongorestore \
      "${restore_base_args[@]}" \
      --db "$db" \
      --dir "${LOCAL_DIR}/${db}" \
      ${DROP_FLAG:+$DROP_FLAG} \
      --gzip
    log_success "Restored $db"
done

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

log_success "mongorestore completed successfully"

