#!/bin/bash
set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
cd "$SCRIPT_DIR"
[ -f "${SCRIPT_DIR}/env-nats" ]         && source "${SCRIPT_DIR}/env-nats"
[ -f "${SCRIPT_DIR}/_credentials.env" ] && source "${SCRIPT_DIR}/_credentials.env"

NATS_NAMESPACE="${NATS_NAMESPACE:-core-service}"
NATS_NAME="${NATS_NAME:-nats}"
NATS_IMAGE="${NATS_IMAGE:-nats:2.10-alpine}"
NATS_PORT="${NATS_PORT:-4222}"
NATS_MONITOR_PORT="${NATS_MONITOR_PORT:-8222}"
NATS_DATA_DIR="${NATS_DATA_DIR:-/data}"
NATS_STORAGE_CLASS="${NATS_STORAGE_CLASS:-longhorn-retain}"
NATS_STORAGE_SIZE="${NATS_STORAGE_SIZE:-1Gi}"
CMD_TSK="${CMD_TSK:-${1:-install}}"

_resolve_kubeconfig() {
  if [ -n "${KUBECONFIG:-}" ] && [ -f "$KUBECONFIG" ]; then
    return
  fi
  for candidate in \
    /var/lib/k0s/pki/admin.conf \
    /etc/kubernetes/admin.conf \
    /root/.kube/config; do
    if [ -f "$candidate" ]; then
      export KUBECONFIG="$candidate"
      return
    fi
  done
  if command -v k0s >/dev/null 2>&1; then
    local tmp
    tmp="$(mktemp)"
    k0s kubeconfig admin > "$tmp" 2>/dev/null && export KUBECONFIG="$tmp" && return
    rm -f "$tmp"
  fi
  echo "ERROR: kubeconfig not found — set KUBECONFIG explicitly" >&2; exit 1
}
_resolve_kubeconfig

_kubectl() {
  if command -v kubectl >/dev/null 2>&1; then
    command kubectl "$@"
  elif command -v k0s >/dev/null 2>&1; then
    k0s kubectl "$@"
  else
    echo "ERROR: no kubectl or k0s binary found" >&2; exit 127
  fi
}

_require_env() {
  local var="$1"
  if [ -z "${!var:-}" ]; then
    echo "ERROR: required variable $var is not set" >&2; exit 1
  fi
}

_namespace_exists() {
  _kubectl get namespace "$NATS_NAMESPACE" >/dev/null 2>&1
}

_pod_name() {
  _kubectl get pod \
    --namespace "$NATS_NAMESPACE" \
    --selector "app.kubernetes.io/name=nats" \
    -o jsonpath='{.items[0].metadata.name}' 2>/dev/null
}

_wait_for_pod() {
  local timeout="${1:-180}"
  _kubectl wait pod \
    --namespace "$NATS_NAMESPACE" \
    --selector "app.kubernetes.io/name=nats" \
    --for=condition=Ready \
    --timeout="${timeout}s"
}

_statefulset_exists() {
  _kubectl get statefulset "$NATS_NAME" --namespace "$NATS_NAMESPACE" >/dev/null 2>&1
}

_do_install() {
  _require_env NATS_USER
  _require_env NATS_PASS

  local already_running=false
  if _namespace_exists && _statefulset_exists; then
    already_running=true
  fi

  if ! _namespace_exists; then
    _kubectl create namespace "$NATS_NAMESPACE"
  fi

  local nkey_args=()
  if [ -n "${NATS_NKEY_PUB:-}" ]; then
    nkey_args+=(--from-literal=NATS_NKEY_PUB="$NATS_NKEY_PUB")
  fi
  _kubectl create secret generic nats-credentials \
    --namespace "$NATS_NAMESPACE" \
    --from-literal=NATS_USER="$NATS_USER" \
    --from-literal=NATS_PASS="$NATS_PASS" \
    "${nkey_args[@]}" \
    --dry-run=client -o yaml | _kubectl apply -f -

  for manifest in storageclass.yaml namespace.yaml configmap.yaml pvc.yaml statefulset.yaml service.yaml; do
    local path="$SCRIPT_DIR/templates/${manifest}"
    [ -f "$path" ] && _kubectl apply -f "$path"
  done

  if $already_running; then
    echo "NATS already running — rolling restart to apply config changes..."
    _kubectl rollout restart statefulset/"$NATS_NAME" --namespace "$NATS_NAMESPACE"
  fi

  echo "Waiting for NATS pod to be ready..."
  _wait_for_pod 180

  echo "NATS installed in namespace $NATS_NAMESPACE"
}

_do_update() {
  for manifest in configmap.yaml statefulset.yaml; do
    local path="$SCRIPT_DIR/templates/${manifest}"
    [ -f "$path" ] && _kubectl apply -f "$path"
  done
  _kubectl rollout restart statefulset/"$NATS_NAME" --namespace "$NATS_NAMESPACE"
  echo "Waiting for NATS to restart..."
  _wait_for_pod 120
  echo "NATS updated"
}

_do_delete() {
  for manifest in statefulset.yaml service.yaml configmap.yaml; do
    local path="$SCRIPT_DIR/manifests/${manifest}"
    [ -f "$path" ] && _kubectl delete -f "$path" --ignore-not-found
  done
  echo "NATS workload removed from namespace $NATS_NAMESPACE (pvc and namespace preserved)"
}

_do_backup() {
  _require_env BACKUP_DIR
  local pod timestamp backup_file
  pod=$(_pod_name)
  timestamp=$(date +%Y%m%d_%H%M%S)
  backup_file="${BACKUP_DIR}/nats_${timestamp}.tar.gz"
  mkdir -p "$BACKUP_DIR"
  _kubectl exec --namespace "$NATS_NAMESPACE" "$pod" -- \
    tar czf - -C "$NATS_DATA_DIR" . > "$backup_file"
  echo "Backup written to $backup_file"
}

_do_health() {
  local pod
  pod=$(_pod_name 2>/dev/null || true)
  if [ -z "$pod" ]; then
    echo "UNHEALTHY: no nats pod found in namespace $NATS_NAMESPACE" >&2
    exit 1
  fi
  local status_code
  status_code=$(_kubectl exec --namespace "$NATS_NAMESPACE" "$pod" -- \
    sh -c "wget -qO /dev/null -S http://localhost:${NATS_MONITOR_PORT}/healthz 2>&1 | grep 'HTTP/' | awk '{print \$2}'" \
    2>/dev/null || echo "0")
  if [ "$status_code" = "200" ]; then
    echo "HEALTHY: nats /healthz returned 200"
  else
    echo "UNHEALTHY: nats /healthz returned $status_code" >&2
    exit 1
  fi
}

case "$CMD_TSK" in
  install) _do_install ;;
  update)  _do_update  ;;
  delete)  _do_delete  ;;
  backup)  _do_backup  ;;
  health)  _do_health  ;;
  *)
    echo "ERROR: unknown CMD_TSK='$CMD_TSK'. Valid: install|update|delete|backup|health" >&2
    exit 1
    ;;
esac