#!/usr/bin/env bash

set -euo pipefail

# init_kerberos.sh
# Purpose:
# - Configure Kerberos realm settings for CloudNative-PG pods
# - Optionally run a Kerberos authentication test inside a Kubernetes pod
# - Optionally provision an in-cluster port forwarder to a Samba AD DC
#
# Usage:
#   ./init_kerberos.sh initialize   # apply krb5.conf to CNPG pods (best-effort)
#   ./init_kerberos.sh test         # run test pod + kinit (with optional AD port forward)
#   ./init_kerberos.sh status       # show detected config + resources
#   ./init_kerberos.sh cleanup      # remove AD forwarder resources

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

ACTION=${1:-initialize}

# Load env without leaking our positional args to the env script (some env.sh may `exec "$@"`).
__PROLE_SAVED_ARGS=("$@")
if [[ -n "${PROLE_HOME:-}" && -f "$PROLE_HOME/env.sh" ]]; then
  set --
  # shellcheck disable=SC1090
  source "$PROLE_HOME/env.sh"
elif [[ -f "$HOME/.prole/env.sh" ]]; then
  set --
  # shellcheck disable=SC1090
  source "$HOME/.prole/env.sh"
fi
set -- "${__PROLE_SAVED_ARGS[@]}"
unset __PROLE_SAVED_ARGS

# Load config values from prole.cfg before applying defaults/CLI overrides.
# shellcheck disable=SC1090
source "$SCRIPT_DIR/prole_cfg.sh"

NAMESPACE=${NAMESPACE:-default}
CNPG_CLUSTER_NAME=${CNPG_CLUSTER_NAME:-prole-db}

KRB5_REALM=${KRB5_REALM:-${REALM:-}}
KRB5_KDC=${KRB5_KDC:-}
KRB5_ADMIN=${KRB5_ADMIN:-}
KRB5_USER=${KRB5_USER:-${KRB5_USERNAME:-}}
KRB5_PASSWORD=${KRB5_PASSWORD:-}

# AD DC forwarding (socat proxy) options
KRB5_AD_PORT_FORWARD=${KRB5_AD_PORT_FORWARD:-}
KRB5_AD_PROXY_NAME=${KRB5_AD_PROXY_NAME:-prole-kerberos-ad-forwarder}
KRB5_AD_SERVICE_NAME=${KRB5_AD_SERVICE_NAME:-prole-kerberos-ad-dc}
KRB5_AD_PROXY_IMAGE=${KRB5_AD_PROXY_IMAGE:-alpine/socat}
KRB5_AD_PROXY_HOST_NETWORK=${KRB5_AD_PROXY_HOST_NETWORK:-1}
KRB5_AD_TCP_PORTS=${KRB5_AD_TCP_PORTS:-"88 389 445 464 636"}
KRB5_AD_UDP_PORTS=${KRB5_AD_UDP_PORTS:-"88 464"}

log() { printf '%s\n' "$*"; }
err() { printf '%s\n' "$*" >&2; }

ensure_tools() {
  for t in kubectl; do
    command -v "$t" >/dev/null || { err "Missing required tool: $t"; exit 1; }
  done
}

ensure_namespace() {
  if ! kubectl get namespace "$NAMESPACE" >/dev/null 2>&1; then
    log "Creating namespace '$NAMESPACE' ..."
    kubectl create namespace "$NAMESPACE" >/dev/null 2>&1 || true
  fi
}

primary_kdc() {
  if [[ -z "${KRB5_KDC}" ]]; then
    echo ""
    return 0
  fi
  printf '%s' "${KRB5_KDC}" | awk -F',' '{print $1}' | xargs
}

is_private_ip() {
  local ip="$1"
  if [[ "$ip" =~ ^10\.|^192\.168\.|^172\.(1[6-9]|2[0-9]|3[0-1])\.|^169\.254\. ]]; then
    return 0
  fi
  return 1
}

default_port_forward_if_local() {
  if [[ -n "${KRB5_AD_PORT_FORWARD}" ]]; then
    return
  fi
  local kdc_host
  kdc_host=$(primary_kdc)
  if [[ -n "$kdc_host" ]] && is_private_ip "$kdc_host"; then
    KRB5_AD_PORT_FORWARD=1
  else
    KRB5_AD_PORT_FORWARD=0
  fi
}

ensure_krb5_conf_configmap() {
  if kubectl -n "$NAMESPACE" get configmap prole-krb5-conf >/dev/null 2>&1; then
    return
  fi
  log "ConfigMap prole-krb5-conf not found. Running init_openbao.sh update..."
  if [[ -x "$SCRIPT_DIR/init_openbao.sh" ]]; then
    KRB5_REALM="$KRB5_REALM" KRB5_KDC="$KRB5_KDC" KRB5_ADMIN="$KRB5_ADMIN" \
      "$SCRIPT_DIR/init_openbao.sh" update || true
  else
    err "init_openbao.sh not found; cannot create krb5.conf ConfigMap."
    exit 1
  fi
}

get_krb5_conf() {
  kubectl -n "$NAMESPACE" get configmap prole-krb5-conf -o jsonpath='{.data.krb5\.conf}' 2>/dev/null || true
}

apply_krb5_conf_to_cnpg_pods() {
  local conf
  conf=$(get_krb5_conf)
  if [[ -z "$conf" ]]; then
    err "Missing krb5.conf data in ConfigMap prole-krb5-conf."
    return 1
  fi

  local pods
  pods=$(kubectl -n "$NAMESPACE" get pods -l "cnpg.io/cluster=$CNPG_CLUSTER_NAME" -o jsonpath='{.items[*].metadata.name}' 2>/dev/null || true)
  if [[ -z "$pods" ]]; then
    err "No CNPG pods found for cluster '$CNPG_CLUSTER_NAME' in namespace '$NAMESPACE'."
    return 1
  fi

  for pod in $pods; do
    log "Updating /etc/krb5.conf in pod $pod ..."
    if printf '%s' "$conf" | kubectl -n "$NAMESPACE" exec -i "$pod" -c postgres -- sh -c 'cat > /etc/krb5.conf'; then
      log "[OK] Updated /etc/krb5.conf in $pod"
    else
      err "[WARN] Unable to write /etc/krb5.conf in $pod (permission?)."
      err "       Consider mounting ConfigMap prole-krb5-conf into the CNPG pods or updating the image."
    fi
  done
}

apply_ad_forwarder() {
  local kdc_host
  kdc_host=$(primary_kdc)
  if [[ -z "$kdc_host" ]]; then
    err "KRB5_KDC is required to set up the AD port forwarder."
    exit 1
  fi

  local host_net_block=""
  local dns_policy_block=""
  if [[ "$KRB5_AD_PROXY_HOST_NETWORK" == "1" ]]; then
    host_net_block="      hostNetwork: true"
    dns_policy_block="      dnsPolicy: ClusterFirstWithHostNet"
  fi

  local tcp_cmd=""
  local udp_cmd=""
  local port
  for port in $KRB5_AD_TCP_PORTS; do
    tcp_cmd+="socat -d -d TCP-LISTEN:${port},fork,reuseaddr TCP:${kdc_host}:${port} & "
  done
  for port in $KRB5_AD_UDP_PORTS; do
    udp_cmd+="socat -d -d UDP-LISTEN:${port},fork,reuseaddr UDP:${kdc_host}:${port} & "
  done

  local forward_cmd="${tcp_cmd}${udp_cmd}wait"

  log "Applying AD DC forwarder '${KRB5_AD_PROXY_NAME}' in namespace '${NAMESPACE}' (target ${kdc_host}) ..."
  cat <<EOF | kubectl apply -n "$NAMESPACE" -f -
apiVersion: apps/v1
kind: Deployment
metadata:
  name: ${KRB5_AD_PROXY_NAME}
  namespace: ${NAMESPACE}
spec:
  replicas: 1
  selector:
    matchLabels:
      app: ${KRB5_AD_PROXY_NAME}
  template:
    metadata:
      labels:
        app: ${KRB5_AD_PROXY_NAME}
    spec:
${host_net_block}
${dns_policy_block}
      containers:
        - name: ad-forwarder
          image: ${KRB5_AD_PROXY_IMAGE}
          command: ["/bin/sh", "-c"]
          args:
            - >-
              ${forward_cmd}
          ports:
$(for port in $KRB5_AD_TCP_PORTS; do printf "            - containerPort: %s\n              protocol: TCP\n" "$port"; done)
$(for port in $KRB5_AD_UDP_PORTS; do printf "            - containerPort: %s\n              protocol: UDP\n" "$port"; done)
---
apiVersion: v1
kind: Service
metadata:
  name: ${KRB5_AD_SERVICE_NAME}
  namespace: ${NAMESPACE}
spec:
  selector:
    app: ${KRB5_AD_PROXY_NAME}
  ports:
$(
  for port in $KRB5_AD_TCP_PORTS; do
    name="tcp-${port}"
    printf "    - name: %s\n      port: %s\n      targetPort: %s\n      protocol: TCP\n" "$name" "$port" "$port"
  done
  for port in $KRB5_AD_UDP_PORTS; do
    name="udp-${port}"
    printf "    - name: %s\n      port: %s\n      targetPort: %s\n      protocol: UDP\n" "$name" "$port" "$port"
  done
)
EOF

  kubectl -n "$NAMESPACE" rollout status deploy/${KRB5_AD_PROXY_NAME} --timeout=120s || true
}

cleanup_ad_forwarder() {
  log "Removing AD forwarder resources (if present)..."
  kubectl -n "$NAMESPACE" delete service "$KRB5_AD_SERVICE_NAME" --ignore-not-found
  kubectl -n "$NAMESPACE" delete deployment "$KRB5_AD_PROXY_NAME" --ignore-not-found
}

initialize() {
  ensure_tools
  ensure_namespace

  if [[ -z "${KRB5_REALM}" || -z "${KRB5_KDC}" ]]; then
    err "ERROR: KRB5_REALM and KRB5_KDC must be set for Kerberos initialization."
    exit 1
  fi

  ensure_krb5_conf_configmap
  apply_krb5_conf_to_cnpg_pods
  log "Kerberos initialization complete."
}

run_test() {
  ensure_tools
  ensure_namespace

  if [[ -z "${KRB5_REALM}" || -z "${KRB5_KDC}" || -z "${KRB5_USER}" || -z "${KRB5_PASSWORD}" ]]; then
    err "ERROR: Missing Kerberos configuration. Ensure KRB5_REALM, KRB5_KDC, KRB5_USER, KRB5_PASSWORD are set."
    exit 1
  fi

  default_port_forward_if_local

  local effective_kdc="$KRB5_KDC"
  if [[ "${KRB5_AD_PORT_FORWARD}" == "1" ]]; then
    apply_ad_forwarder
    effective_kdc="${KRB5_AD_SERVICE_NAME}"
  fi

  log "Using KDC endpoint for test: ${effective_kdc}"

  # Ensure krb5.conf configmap is updated for the test endpoint
  if [[ -x "$SCRIPT_DIR/init_openbao.sh" ]]; then
    KRB5_REALM="$KRB5_REALM" KRB5_KDC="$effective_kdc" KRB5_ADMIN="$KRB5_ADMIN" \
      "$SCRIPT_DIR/init_openbao.sh" update || true
  fi

  if [[ -x "$SCRIPT_DIR/init_kerberos_test.sh" ]]; then
    KRB5_REALM="$KRB5_REALM" KRB5_KDC="$effective_kdc" KRB5_ADMIN="$KRB5_ADMIN" \
      KRB5_USER="$KRB5_USER" KRB5_PASSWORD="$KRB5_PASSWORD" \
      "$SCRIPT_DIR/init_kerberos_test.sh" test
  else
    err "init_kerberos_test.sh not found."
    exit 1
  fi
}

status() {
  ensure_tools
  log "--- init_kerberos status ---"
  log "Namespace: $NAMESPACE"
  log "CNPG cluster: $CNPG_CLUSTER_NAME"
  log "KRB5_REALM: ${KRB5_REALM:-<unset>}"
  log "KRB5_KDC: ${KRB5_KDC:-<unset>}"
  log "KRB5_USER: ${KRB5_USER:-<unset>}"
  log "AD forwarder enabled: ${KRB5_AD_PORT_FORWARD:-<auto>}"
  log "AD forwarder deployment: $KRB5_AD_PROXY_NAME"
  log "AD forwarder service: $KRB5_AD_SERVICE_NAME"
  if kubectl -n "$NAMESPACE" get configmap prole-krb5-conf >/dev/null 2>&1; then
    log "[OK] ConfigMap prole-krb5-conf present"
  else
    log "[MISSING] ConfigMap prole-krb5-conf"
  fi
  if kubectl -n "$NAMESPACE" get deploy "$KRB5_AD_PROXY_NAME" >/dev/null 2>&1; then
    log "[OK] AD forwarder deployment present"
  else
    log "[INFO] AD forwarder deployment not present"
  fi
  if kubectl -n "$NAMESPACE" get svc "$KRB5_AD_SERVICE_NAME" >/dev/null 2>&1; then
    log "[OK] AD forwarder service present"
  else
    log "[INFO] AD forwarder service not present"
  fi
}

case "$ACTION" in
  initialize|configure|update|reload)
    initialize
    ;;
  test)
    run_test
    ;;
  status)
    status
    ;;
  cleanup)
    ensure_tools
    ensure_namespace
    cleanup_ad_forwarder
    ;;
  *)
    err "Usage: $0 {initialize|configure|update|reload|test|status|cleanup}" >&2
    exit 2
    ;;
esac