#!/usr/bin/env bash
set -euo pipefail

CM_NAMESPACE="observability"
CM_NAME="o11y-cr-templates"
KUBERNETES_AGENT_KEY="kubernetesagent.yaml"
SHARED_POOL_AGENT_KEY="sharedpoolagent.yaml"
DRY_RUN="true"
ROLLBACK="false"

LOGGING_NAMESPACE="logging-serverless"
VECTOR_AGENT_DAEMONSET="o11y-vectoragent"
MONITORING_NAMESPACE="monitoring-serverless"
VMAGENT_RESOURCE="vmagents.operator.victoriametrics.com"
VMAGENT_STATEFULSET_SELECTOR="app.kubernetes.io/name=vmagent"
KUBERNETES_AGENT_RESOURCE="kubernetesagent.agent.observability.pingcap.com"
VICTORIA_METRICS_OPERATOR_DEPLOYMENT="o11y-victoria-metrics-operator"

ORIGINAL_REPLICAS_ANNOTATION="tidbcloud.com/original-replicas"
ORIGINAL_DS_SELECTOR_ANNOTATION="tidbcloud.com/o11y-vectoragent-disabled-original"
DISABLE_DS_SELECTOR_KEY="tidbcloud.com/o11y-vectoragent-disabled"
DISABLE_DS_SELECTOR_VALUE="true"
ANNOTATION_ABSENT_VALUE="__ABSENT__"

usage() {
  cat <<'USAGE'
Usage: create-o11y-crs.sh [options]

Read CR templates from a ConfigMap, stop monitoring workloads, and apply KubernetesAgent + SharedPoolAgent CRs.

Options:
  -n, --namespace <namespace>                ConfigMap namespace (default: observability)
  -c, --configmap <name>                     ConfigMap name (default: o11y-cr-templates)
      --kubernetes-agent-key <key>           Key for KubernetesAgent template (default: kubernetesagent.yaml)
      --sharedpool-agent-key <key>           Key for SharedPoolAgent template (default: sharedpoolagent.yaml)
      --dry-run                              Validate only, do not persist resources (default)
      --execute                              Apply resources to cluster
      --rollback                             Delete KubernetesAgent + SharedPoolAgent and resume previously stopped resources
  -h, --help                                 Show this help
USAGE
}

while [[ $# -gt 0 ]]; do
  case "$1" in
    -n|--namespace)
      CM_NAMESPACE="$2"
      shift 2
      ;;
    -c|--configmap)
      CM_NAME="$2"
      shift 2
      ;;
    --kubernetes-agent-key)
      KUBERNETES_AGENT_KEY="$2"
      shift 2
      ;;
    --sharedpool-agent-key)
      SHARED_POOL_AGENT_KEY="$2"
      shift 2
      ;;
    --dry-run)
      DRY_RUN="true"
      shift
      ;;
    --execute)
      DRY_RUN="false"
      shift
      ;;
    --rollback)
      ROLLBACK="true"
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      echo "Unknown argument: $1" >&2
      usage
      exit 1
      ;;
  esac
done

require_commands() {
  local missing=0
  local cmd
  for cmd in "$@"; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
      echo "$cmd is required but not found in PATH." >&2
      missing=1
    fi
  done
  if [[ "$missing" -ne 0 ]]; then
    exit 1
  fi
}

require_commands cat kubectl jq awk sleep

run_mutating_kubectl() {
  if [[ "$DRY_RUN" == "true" ]]; then
    echo "+ kubectl $* --dry-run=client"
    kubectl "$@" --dry-run=client
  else
    echo "+ kubectl $*"
    kubectl "$@"
  fi
}

list_vmagent_resources() {
  local vmagents
  if ! vmagents="$(kubectl -n "$MONITORING_NAMESPACE" get "$VMAGENT_RESOURCE" -o name 2>/dev/null)"; then
    echo "Skip: resource type '$VMAGENT_RESOURCE' not found in namespace $MONITORING_NAMESPACE." >&2
    return
  fi

  vmagents="$(printf '%s\n' "$vmagents" | awk 'NF && !seen[$0]++')"
  if [[ -n "$vmagents" ]]; then
    echo "VMAgent resources in namespace $MONITORING_NAMESPACE:" >&2
    printf '%s\n' "$vmagents" >&2
  fi

  printf '%s\n' "$vmagents"
}

list_vmagent_statefulsets() {
  local vmagent_statefulsets
  if ! vmagent_statefulsets="$(
    kubectl -n "$MONITORING_NAMESPACE" get statefulset -l "$VMAGENT_STATEFULSET_SELECTOR" -o name 2>/dev/null
  )"; then
    echo "Skip: failed to list StatefulSet resources in namespace $MONITORING_NAMESPACE." >&2
    return
  fi

  vmagent_statefulsets="$(printf '%s\n' "$vmagent_statefulsets" | awk 'NF && !seen[$0]++')"
  if [[ -n "$vmagent_statefulsets" ]]; then
    echo "VMAgent StatefulSets in namespace $MONITORING_NAMESPACE (selector: $VMAGENT_STATEFULSET_SELECTOR):" >&2
    printf '%s\n' "$vmagent_statefulsets" >&2
  fi

  printf '%s\n' "$vmagent_statefulsets"
}

delete_vmagent_statefulsets_for_rollback() {
  local vmagent_statefulsets
  vmagent_statefulsets="$(list_vmagent_statefulsets)"
  if [[ -z "$vmagent_statefulsets" ]]; then
    echo "No VMAgent StatefulSet resources found in namespace $MONITORING_NAMESPACE."
    return
  fi

  local non_zero_replicas=""
  while IFS= read -r vmagent_statefulset; do
    [[ -z "$vmagent_statefulset" ]] && continue
    local spec_replicas
    local status_replicas
    spec_replicas="$(kubectl -n "$MONITORING_NAMESPACE" get "$vmagent_statefulset" -o jsonpath='{.spec.replicas}' 2>/dev/null || true)"
    status_replicas="$(kubectl -n "$MONITORING_NAMESPACE" get "$vmagent_statefulset" -o jsonpath='{.status.replicas}' 2>/dev/null || true)"
    # Keep delete checks deterministic when fields are missing: treat empty spec as non-zero (safe),
    # and empty status as zero (status is often absent before controller reports).
    if [[ -z "$spec_replicas" ]]; then
      spec_replicas="1"
    fi
    if [[ -z "$status_replicas" ]]; then
      status_replicas="0"
    fi
    if [[ "$spec_replicas" != "0" || "$status_replicas" != "0" ]]; then
      non_zero_replicas+="$vmagent_statefulset spec.replicas=$spec_replicas status.replicas=$status_replicas"$'\n'
    fi
  done <<< "$vmagent_statefulsets"

  if [[ -n "$non_zero_replicas" ]]; then
    echo "Refuse to delete VMAgent StatefulSets: all StatefulSets must have replicas=0 first." >&2
    printf '%s' "$non_zero_replicas" >&2
    exit 1
  fi

  echo "Deleting VMAgent StatefulSets in namespace $MONITORING_NAMESPACE ..."
  while IFS= read -r vmagent_statefulset; do
    [[ -z "$vmagent_statefulset" ]] && continue
    echo "Clearing finalizers on $vmagent_statefulset before deletion ..."
    run_mutating_kubectl -n "$MONITORING_NAMESPACE" patch "$vmagent_statefulset" --type merge \
      -p '{"metadata":{"finalizers":[]}}'
    echo "Deleting $vmagent_statefulset ..."
    run_mutating_kubectl -n "$MONITORING_NAMESPACE" delete --ignore-not-found "$vmagent_statefulset"
  done <<< "$vmagent_statefulsets"
}

wait_for_vmagent_status_replicas_zero() {
  local vmagents="$1"

  if [[ "$DRY_RUN" == "true" ]]; then
    echo "Dry-run mode: skip waiting for vmagent status.replicas to become 0."
    return
  fi

  if [[ -z "$vmagents" ]]; then
    echo "No VMAgent resources to wait for."
    return
  fi

  local timeout_seconds="${VMAGENT_SCALE_DOWN_TIMEOUT_SECONDS:-300}"
  local poll_interval_seconds=5
  local elapsed_seconds=0

  echo "Waiting for vmagent status.replicas=0 in namespace $MONITORING_NAMESPACE (timeout: ${timeout_seconds}s) ..."
  while (( elapsed_seconds < timeout_seconds )); do
    local pending_statuses=""

    while IFS= read -r vmagent; do
      [[ -z "$vmagent" ]] && continue
      local status_replicas
      status_replicas="$(kubectl -n "$MONITORING_NAMESPACE" get "$vmagent" -o jsonpath='{.status.replicas}' 2>/dev/null || true)"
      if [[ "$status_replicas" != "0" ]]; then
        if [[ -z "$status_replicas" ]]; then
          status_replicas="<empty>"
        fi
        pending_statuses+="$vmagent status.replicas=$status_replicas"$'\n'
      fi
    done <<< "$vmagents"

    if [[ -z "$pending_statuses" ]]; then
      echo "All vmagent resources report status.replicas=0."
      return
    fi

    echo "Still waiting for vmagent resources:"
    printf '%s' "$pending_statuses"
    sleep "$poll_interval_seconds"
    elapsed_seconds=$((elapsed_seconds + poll_interval_seconds))
  done

  echo "Timeout waiting for vmagent status.replicas=0 in namespace $MONITORING_NAMESPACE." >&2
  while IFS= read -r vmagent; do
    [[ -z "$vmagent" ]] && continue
    local status_replicas
    status_replicas="$(kubectl -n "$MONITORING_NAMESPACE" get "$vmagent" -o jsonpath='{.status.replicas}' 2>/dev/null || true)"
    if [[ -z "$status_replicas" ]]; then
      status_replicas="<empty>"
    fi
    echo "$vmagent status.replicas=$status_replicas" >&2
  done <<< "$vmagents"
  exit 1
}

stop_vectoragent_daemonset() {
  if ! kubectl -n "$LOGGING_NAMESPACE" get daemonset "$VECTOR_AGENT_DAEMONSET" >/dev/null 2>&1; then
    echo "Skip: daemonset $LOGGING_NAMESPACE/$VECTOR_AGENT_DAEMONSET not found."
    return
  fi

  local original_selector_value
  original_selector_value="$(
    kubectl -n "$LOGGING_NAMESPACE" get daemonset "$VECTOR_AGENT_DAEMONSET" \
      -o go-template="{{if .spec.template.spec.nodeSelector}}{{index .spec.template.spec.nodeSelector \"$DISABLE_DS_SELECTOR_KEY\"}}{{end}}"
  )"
  if [[ -z "$original_selector_value" || "$original_selector_value" == "<no value>" ]]; then
    original_selector_value="$ANNOTATION_ABSENT_VALUE"
  fi

  echo "Stopping daemonset $LOGGING_NAMESPACE/$VECTOR_AGENT_DAEMONSET ..."
  run_mutating_kubectl -n "$LOGGING_NAMESPACE" annotate daemonset "$VECTOR_AGENT_DAEMONSET" \
    "$ORIGINAL_DS_SELECTOR_ANNOTATION=$original_selector_value" --overwrite
  run_mutating_kubectl -n "$LOGGING_NAMESPACE" patch daemonset "$VECTOR_AGENT_DAEMONSET" --type merge \
    -p "{\"spec\":{\"template\":{\"spec\":{\"nodeSelector\":{\"$DISABLE_DS_SELECTOR_KEY\":\"$DISABLE_DS_SELECTOR_VALUE\"}}}}}"
}

resume_vectoragent_daemonset() {
  if ! kubectl -n "$LOGGING_NAMESPACE" get daemonset "$VECTOR_AGENT_DAEMONSET" >/dev/null 2>&1; then
    echo "Skip: daemonset $LOGGING_NAMESPACE/$VECTOR_AGENT_DAEMONSET not found."
    return
  fi

  local original_selector_value
  original_selector_value="$(
    kubectl -n "$LOGGING_NAMESPACE" get daemonset "$VECTOR_AGENT_DAEMONSET" \
      -o go-template="{{if .metadata.annotations}}{{index .metadata.annotations \"$ORIGINAL_DS_SELECTOR_ANNOTATION\"}}{{end}}"
  )"
  if [[ -z "$original_selector_value" || "$original_selector_value" == "<no value>" ]]; then
    echo "Skip: daemonset $LOGGING_NAMESPACE/$VECTOR_AGENT_DAEMONSET missing annotation $ORIGINAL_DS_SELECTOR_ANNOTATION."
    return
  fi

  echo "Resuming daemonset $LOGGING_NAMESPACE/$VECTOR_AGENT_DAEMONSET ..."
  if [[ "$original_selector_value" == "$ANNOTATION_ABSENT_VALUE" ]]; then
    run_mutating_kubectl -n "$LOGGING_NAMESPACE" patch daemonset "$VECTOR_AGENT_DAEMONSET" --type merge \
      -p "{\"spec\":{\"template\":{\"spec\":{\"nodeSelector\":{\"$DISABLE_DS_SELECTOR_KEY\":null}}}}}"
  else
    run_mutating_kubectl -n "$LOGGING_NAMESPACE" patch daemonset "$VECTOR_AGENT_DAEMONSET" --type merge \
      -p "{\"spec\":{\"template\":{\"spec\":{\"nodeSelector\":{\"$DISABLE_DS_SELECTOR_KEY\":\"$original_selector_value\"}}}}}"
  fi
  run_mutating_kubectl -n "$LOGGING_NAMESPACE" annotate daemonset "$VECTOR_AGENT_DAEMONSET" \
    "${ORIGINAL_DS_SELECTOR_ANNOTATION}-"
}

stop_vmagent_workloads() {
  local vmagents
  vmagents="$(list_vmagent_resources)"
  if [[ -z "$vmagents" ]]; then
    echo "No VMAgent resources found in namespace $MONITORING_NAMESPACE."
  else
    echo "Setting VMAgent spec.replicaCount to 0 in namespace $MONITORING_NAMESPACE:"

    while IFS= read -r vmagent; do
      [[ -z "$vmagent" ]] && continue
      local current_replica_count
      current_replica_count="$(kubectl -n "$MONITORING_NAMESPACE" get "$vmagent" -o jsonpath='{.spec.replicaCount}')"
      if [[ -z "$current_replica_count" ]]; then
        current_replica_count="1"
      fi

      echo "Patching $vmagent spec.replicaCount from $current_replica_count to 0 ..."
      run_mutating_kubectl -n "$MONITORING_NAMESPACE" patch "$vmagent" --type merge \
        -p '{"spec":{"replicaCount":0}}'
    done <<< "$vmagents"
  fi
  wait_for_vmagent_status_replicas_zero "$vmagents"

  local operator_workload="deployment.apps/$VICTORIA_METRICS_OPERATOR_DEPLOYMENT"
  if ! kubectl -n "$MONITORING_NAMESPACE" get deployment "$VICTORIA_METRICS_OPERATOR_DEPLOYMENT" >/dev/null 2>&1; then
    echo "Skip: deployment $MONITORING_NAMESPACE/$VICTORIA_METRICS_OPERATOR_DEPLOYMENT not found."
    return
  fi

  local original_replicas
  original_replicas="$(kubectl -n "$MONITORING_NAMESPACE" get "$operator_workload" -o jsonpath='{.spec.replicas}')"
  if [[ -z "$original_replicas" ]]; then
    original_replicas="1"
  fi

  echo "Scaling $operator_workload from replicas=$original_replicas to 0 ..."
  run_mutating_kubectl -n "$MONITORING_NAMESPACE" annotate "$operator_workload" \
    "$ORIGINAL_REPLICAS_ANNOTATION=$original_replicas" --overwrite
  run_mutating_kubectl -n "$MONITORING_NAMESPACE" scale "$operator_workload" --replicas=0
}

resume_vmagent_workloads() {
  local operator_workload="deployment.apps/$VICTORIA_METRICS_OPERATOR_DEPLOYMENT"
  if kubectl -n "$MONITORING_NAMESPACE" get deployment "$VICTORIA_METRICS_OPERATOR_DEPLOYMENT" >/dev/null 2>&1; then
    local original_replicas
    local has_original_replicas="true"
    original_replicas="$(
      kubectl -n "$MONITORING_NAMESPACE" get "$operator_workload" \
        -o go-template="{{if .metadata.annotations}}{{index .metadata.annotations \"$ORIGINAL_REPLICAS_ANNOTATION\"}}{{end}}"
    )"
    if [[ -z "$original_replicas" || "$original_replicas" == "<no value>" ]]; then
      has_original_replicas="false"
      original_replicas="1"
      echo "Annotation $ORIGINAL_REPLICAS_ANNOTATION missing on $operator_workload, defaulting to replicas=1."
    fi

    echo "Restoring $operator_workload to replicas=$original_replicas ..."
    run_mutating_kubectl -n "$MONITORING_NAMESPACE" scale "$operator_workload" --replicas="$original_replicas"
    if [[ "$has_original_replicas" == "true" ]]; then
      run_mutating_kubectl -n "$MONITORING_NAMESPACE" annotate "$operator_workload" \
        "${ORIGINAL_REPLICAS_ANNOTATION}-"
    fi
  else
    echo "Skip: deployment $MONITORING_NAMESPACE/$VICTORIA_METRICS_OPERATOR_DEPLOYMENT not found."
  fi

  local vmagents
  vmagents="$(list_vmagent_resources)"
  if [[ -z "$vmagents" ]]; then
    echo "No VMAgent resources found in namespace $MONITORING_NAMESPACE."
  else
    echo "Restoring VMAgent spec.replicaCount to 1 in namespace $MONITORING_NAMESPACE:"

    while IFS= read -r vmagent; do
      [[ -z "$vmagent" ]] && continue
      echo "Patching $vmagent spec.replicaCount to 1 ..."
      run_mutating_kubectl -n "$MONITORING_NAMESPACE" patch "$vmagent" --type merge \
        -p '{"spec":{"replicaCount":1}}'
    done <<< "$vmagents"
  fi
}

get_template_content() {
  local key="$1"
  local content

  content="$(kubectl -n "$CM_NAMESPACE" get configmap "$CM_NAME" -o go-template="{{index .data \"$key\"}}")"
  if [[ -z "$content" || "$content" == "<no value>" ]]; then
    echo "Template key '$key' not found in ConfigMap $CM_NAMESPACE/$CM_NAME." >&2
    exit 1
  fi

  printf '%s\n' "$content"
}

list_template_resource_targets() {
  local content="$1"

  printf '%s\n' "$content" | kubectl create --dry-run=client -f - -o json | jq -r '
    def to_resource_type:
      (.kind | ascii_downcase) +
      (if (.apiVersion | contains("/")) then "." + (.apiVersion | split("/")[0]) else "" end);
    if .kind == "List" then .items[] else . end
    | [to_resource_type, .metadata.name, (.metadata.namespace // "")] | @tsv
  '
}

apply_template() {
  local key="$1"
  local content

  content="$(get_template_content "$key")"

  echo "Template key '$key' from $CM_NAMESPACE/$CM_NAME:"
  echo "----- BEGIN CR ($key) -----"
  printf '%s\n' "$content"
  echo "----- END CR ($key) -----"
  echo "Applying template key '$key' ..."
  if [[ "$DRY_RUN" == "true" ]]; then
    printf '%s\n' "$content" | kubectl apply --dry-run=client -f -
  else
    printf '%s\n' "$content" | kubectl apply -f -
  fi
}

delete_template() {
  local key="$1"
  local content
  local targets

  content="$(get_template_content "$key")"

  echo "Template key '$key' from $CM_NAMESPACE/$CM_NAME:"
  echo "----- BEGIN CR ($key) -----"
  printf '%s\n' "$content"
  echo "----- END CR ($key) -----"
  targets="$(list_template_resource_targets "$content")"
  if [[ -z "$targets" ]]; then
    echo "No resources found in template key '$key'."
    return
  fi

  echo "Deleting resources from template key '$key' by kind/name/namespace ..."
  while IFS=$'\t' read -r resource_type resource_name resource_namespace; do
    [[ -z "$resource_type" || -z "$resource_name" ]] && continue
    if [[ -n "$resource_namespace" ]]; then
      echo "Deleting $resource_type/$resource_name in namespace $resource_namespace ..."
      run_mutating_kubectl -n "$resource_namespace" delete --ignore-not-found "$resource_type" "$resource_name"
    else
      echo "Deleting $resource_type/$resource_name ..."
      run_mutating_kubectl delete --ignore-not-found "$resource_type" "$resource_name"
    fi
  done <<< "$targets"
}

print_rollback_stuck_cleanup_hint() {
  cat <<EOF
[Rollback hint] If KubernetesAgent deletion is stuck, ask operator to clear finalizers and delete resources manually:

# 1) VMAgent CRs in $MONITORING_NAMESPACE
kubectl -n "$MONITORING_NAMESPACE" get "$VMAGENT_RESOURCE" -o name
kubectl -n "$MONITORING_NAMESPACE" get "$VMAGENT_RESOURCE" -o name | while read -r res; do
  [ -z "\$res" ] && continue
  kubectl -n "$MONITORING_NAMESPACE" patch "\$res" --type merge -p '{"metadata":{"finalizers":[]}}'
  kubectl -n "$MONITORING_NAMESPACE" delete "\$res" --ignore-not-found
done

# 2) VMAgent StatefulSets in $MONITORING_NAMESPACE (selector: $VMAGENT_STATEFULSET_SELECTOR)
kubectl -n "$MONITORING_NAMESPACE" get statefulset -l "$VMAGENT_STATEFULSET_SELECTOR" -o name
kubectl -n "$MONITORING_NAMESPACE" get statefulset -l "$VMAGENT_STATEFULSET_SELECTOR" -o name | while read -r sts; do
  [ -z "\$sts" ] && continue
  kubectl -n "$MONITORING_NAMESPACE" patch "\$sts" --type merge -p '{"metadata":{"finalizers":[]}}'
  kubectl -n "$MONITORING_NAMESPACE" delete "\$sts" --ignore-not-found
done

# 3) KubernetesAgent CRs (all namespaces)
kubectl get "$KUBERNETES_AGENT_RESOURCE" -A
kubectl get "$KUBERNETES_AGENT_RESOURCE" -A -o jsonpath='{range .items[*]}{.metadata.namespace}{"\\t"}{.metadata.name}{"\\n"}{end}' | while IFS=$'\\t' read -r ns name; do
  [ -z "\$ns" ] && continue
  [ -z "\$name" ] && continue
  kubectl -n "\$ns" patch "$KUBERNETES_AGENT_RESOURCE" "\$name" --type merge -p '{"metadata":{"finalizers":[]}}'
  kubectl -n "\$ns" delete "$KUBERNETES_AGENT_RESOURCE" "\$name" --ignore-not-found
done
EOF
}

if [[ "$ROLLBACK" == "true" ]]; then
  if ! kubectl -n "$CM_NAMESPACE" get configmap "$CM_NAME" >/dev/null 2>&1; then
    echo "ConfigMap $CM_NAMESPACE/$CM_NAME not found." >&2
    exit 1
  fi
  print_rollback_stuck_cleanup_hint
  delete_template "$SHARED_POOL_AGENT_KEY"
  delete_template "$KUBERNETES_AGENT_KEY"
  delete_vmagent_statefulsets_for_rollback
  resume_vmagent_workloads
  resume_vectoragent_daemonset
  echo "Done."
  exit 0
fi

if ! kubectl -n "$CM_NAMESPACE" get configmap "$CM_NAME" >/dev/null 2>&1; then
  echo "ConfigMap $CM_NAMESPACE/$CM_NAME not found." >&2
  exit 1
fi

stop_vectoragent_daemonset
stop_vmagent_workloads
apply_template "$KUBERNETES_AGENT_KEY"
apply_template "$SHARED_POOL_AGENT_KEY"

echo "Done."