#!/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)

# Load environment and config via prole_cfg.sh
# shellcheck disable=SC1090
source "$SCRIPT_DIR/prole_cfg.sh"

ACTION=${1:-initialize}

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"}
CNPG_WAIT_TIMEOUT=${CNPG_WAIT_TIMEOUT:-300}

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
}

resolve_krb5_defaults() {
  # If 'dog' KDC is available in 'authority' namespace, use it as default
  if kubectl get pod -n "authority" -l app=dog >/dev/null 2>&1; then
    KRB5_KDC="dog.authority.svc.cluster.local"
    KRB5_ADMIN="dog.authority.svc.cluster.local"
    if [[ -z "${KRB5_REALM}" ]]; then
        KRB5_REALM="PROLE.ORG"
    fi
  fi

  if [[ -z "${KRB5_REALM}" && -n "${REALM:-}" ]]; then
    KRB5_REALM="$REALM"
  fi
  if [[ -z "${KRB5_KDC}" && -n "${DOMAIN:-}" ]]; then
    KRB5_KDC="kdc.${DOMAIN}"
  fi
  if [[ -z "${KRB5_ADMIN}" ]]; then
    if [[ -n "${KRB5_KDC}" ]]; then
      KRB5_ADMIN=$(printf '%s' "${KRB5_KDC}" | awk -F',' '{print $1}' | xargs)
    elif [[ -n "${DOMAIN:-}" ]]; then
      KRB5_ADMIN="kdc.${DOMAIN}"
    fi
  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() {
  log "Creating/updating ConfigMap prole-krb5-conf ..."

  local kdc_val admin_val realm_val domain_val
  realm_val="$KRB5_REALM"
  domain_val="${DOMAIN:-}"
  if [[ -z "$domain_val" && -n "$realm_val" ]]; then
    domain_val=$(printf '%s' "$realm_val" | tr '[:upper:]' '[:lower:]')
  fi
  kdc_val="$KRB5_KDC"
  admin_val="${KRB5_ADMIN:-}"

  if [[ -z "$realm_val" || -z "$kdc_val" ]]; then
    err "ERROR: KRB5_REALM and KRB5_KDC must be set to build krb5.conf."
    return 1
  fi

  if [[ -z "$admin_val" ]]; then
    admin_val=$(printf '%s' "$kdc_val" | awk -F',' '{print $1}' | xargs)
  fi

  local TMP
  TMP=$(mktemp)
  cat >"$TMP" <<EOF
[libdefaults]
 default_realm = $realm_val
 dns_lookup_realm = true
 dns_lookup_kdc = true

[realms]
 $realm_val = {
  kdc = $kdc_val
  admin_server = $admin_val
 }

[domain_realm]
 .$domain_val = $realm_val
 $domain_val = $realm_val
EOF
  kubectl -n "$NAMESPACE" create configmap prole-krb5-conf --from-file=krb5.conf="$TMP" --dry-run=client -o yaml | kubectl apply -f -
  rm -f "$TMP"
}

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

wait_for_cnpg_pods() {
  local timeout=${1:-$CNPG_WAIT_TIMEOUT}
  local start_time
  start_time=$(date +%s)

  while true; do
    if kubectl -n "$NAMESPACE" get cluster "$CNPG_CLUSTER_NAME" >/dev/null 2>&1; then
      if kubectl -n "$NAMESPACE" get pods -l "cnpg.io/cluster=$CNPG_CLUSTER_NAME" --no-headers 2>/dev/null | grep -q .; then
        if kubectl -n "$NAMESPACE" wait --for=condition=Ready pod -l "cnpg.io/cluster=$CNPG_CLUSTER_NAME" --timeout=60s >/dev/null 2>&1; then
          return 0
        fi
      fi
    fi

    if (( $(date +%s) - start_time > timeout )); then
      err "ERROR: Timed out waiting for CNPG pods for cluster '$CNPG_CLUSTER_NAME' in namespace '$NAMESPACE'."
      return 1
    fi
    log "Waiting for CNPG pods for cluster '$CNPG_CLUSTER_NAME'..."
    sleep 3
  done
}

patch_cnpg_cluster_for_auth() {
  log "Patching CNPG Cluster $CNPG_CLUSTER_NAME to enable GSSAPI (best-effort) ..."
  log "Current Context: $(kubectl config current-context 2>/dev/null || echo 'unknown')"
  log "Target Namespace: $NAMESPACE"

  if ! kubectl -n "$NAMESPACE" get cluster "$CNPG_CLUSTER_NAME" >/dev/null 2>&1; then
    err "ERROR: Cluster '$CNPG_CLUSTER_NAME' not found in namespace '$NAMESPACE'."
    return 1
  fi

  # Best-effort patch; schema may vary with CNPG version.
  # We remove krb_srvname as it is unrecognized in some PG 17 builds.
  # We revert certificates to default to avoid operator TLS issues.
  
  local kerberos_enabled="${KERBEROS_ENABLED:-${ENABLED:-false}}"
  local pg_hba
  if [[ "$kerberos_enabled" == "true" || "$kerberos_enabled" == "True" || "$kerberos_enabled" == "1" ]]; then
     pg_hba="[
          \"local all postgres trust\",
          \"host all postgres all scram-sha-256\",
          \"host all all all gss include_realm=1 krb_realm=$KRB5_REALM\",
          \"host all all all scram-sha-256\"
        ]"
  else
     pg_hba="[
          \"local all postgres trust\",
          \"host all postgres all scram-sha-256\",
          \"host all all all scram-sha-256\"
        ]"
  fi

  if ! kubectl -n "$NAMESPACE" patch cluster "$CNPG_CLUSTER_NAME" --type merge -p "{
    \"spec\": {
      \"postgresql\": {
        \"parameters\": {\"krb_srvname\": null},
        \"pg_hba\": $pg_hba
      },
      \"certificates\": {
        \"serverCASecret\": null,
        \"clientCASecret\": null
      }
    }
  }"; then
    err "Note: patch failed, but cluster exists. It may need adjustment for your CNPG version."
  fi
}

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
  resolve_krb5_defaults

  # Support both names for the toggle from prole.cfg
  local enabled="${KERBEROS_ENABLED:-${ENABLED:-false}}"
  if [[ "$enabled" == "false" || "$enabled" == "0" || "$enabled" == "False" ]]; then
    log "Kerberos is disabled (KERBEROS_ENABLED=$enabled). Skipping initialization."
    # Cleanup any existing Kerberos config if it was previously enabled
    kubectl -n "$NAMESPACE" delete configmap prole-krb5-conf --ignore-not-found
    return 0
  fi

  if [[ -z "${KRB5_REALM}" || -z "${KRB5_KDC}" ]]; then
    err "ERROR: KRB5_REALM and KRB5_KDC must be set for Kerberos initialization."
    err "       Set KRB5_REALM/KRB5_KDC or REALM/DOMAIN in env or config."
    exit 1
  fi

  ensure_krb5_conf_configmap
  patch_cnpg_cluster_for_auth || true
  wait_for_cnpg_pods "$CNPG_WAIT_TIMEOUT"
  apply_krb5_conf_to_cnpg_pods
  log "Kerberos initialization complete."
}

run_test() {
  ensure_tools
  ensure_namespace
  resolve_krb5_defaults

  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}"

  KRB5_KDC="$effective_kdc"
  ensure_krb5_conf_configmap

  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