#!/bin/bash
# Rook-Ceph Preparation Script
# Wipes NVMe drives on all nodes, creates namespace with proper PSS labels,
# cleans up stale resources, and verifies node availability

set -euo pipefail

NS="rook-ceph"
WIPE_IMAGE="quay.io/ceph/ceph:v19.2.0"

# Node to NVMe device mapping
declare -A NODE_NVME=(
  ["talos-cp-01"]="nvme0n1"
  ["talos-cp-02"]="nvme1n1"
  ["talos-cp-03"]="nvme0n1"
  ["talos-cp-04"]="nvme0n1"
)

echo "=== Rook-Ceph Preparation ==="

# Create namespace if it doesn't exist
if ! kubectl get ns "${NS}" >/dev/null 2>&1; then
  echo "Creating namespace ${NS}..."
  kubectl create ns "${NS}"
fi

# Apply privileged Pod Security Standards (required for Ceph)
echo "Applying privileged Pod Security Standards..."
kubectl label namespace "${NS}" pod-security.kubernetes.io/enforce=privileged --overwrite
kubectl label namespace "${NS}" pod-security.kubernetes.io/audit=privileged --overwrite
kubectl label namespace "${NS}" pod-security.kubernetes.io/warn=privileged --overwrite

# Clean up any stale OSD deployments
echo "Cleaning up stale OSD resources..."
for osd_deploy in $(kubectl -n "${NS}" get deployments -o name 2>/dev/null | grep rook-ceph-osd- || true); do
  echo "  Deleting ${osd_deploy}..."
  kubectl -n "${NS}" delete "${osd_deploy}" --ignore-not-found=true
done

# Clean up stale OSD prepare pods
kubectl -n "${NS}" delete pods -l app=rook-ceph-osd-prepare --ignore-not-found=true 2>/dev/null || true

# Clean up any stale finalizers from previous failed deployments
echo "Checking for stale CephCluster..."
if kubectl -n "${NS}" get cephcluster rook-ceph >/dev/null 2>&1; then
  PHASE=$(kubectl -n "${NS}" get cephcluster rook-ceph -o jsonpath='{.status.phase}' 2>/dev/null || echo "Unknown")
  if [ "$PHASE" = "Deleting" ]; then
    echo "Found stale CephCluster in Deleting state, cleaning up finalizers..."
    kubectl -n "${NS}" patch cephcluster rook-ceph --type merge -p '{"metadata":{"finalizers":[]}}' 2>/dev/null || true
  fi
fi

# Function to wipe NVMe on a specific node
wipe_nvme() {
  local node=$1
  local device=$2
  local pod_name="wipe-nvme-${node}"

  echo "Wiping /dev/${device} on ${node}..."

  # Delete any existing wipe pod
  kubectl -n "${NS}" delete pod "${pod_name}" --ignore-not-found=true 2>/dev/null || true

  # Create wipe pod
  cat <<EOF | kubectl apply -f -
apiVersion: v1
kind: Pod
metadata:
  name: ${pod_name}
  namespace: ${NS}
spec:
  restartPolicy: Never
  nodeName: ${node}
  tolerations:
    - operator: Exists
  containers:
    - name: wipe
      image: ${WIPE_IMAGE}
      securityContext:
        privileged: true
      command:
        - /bin/bash
        - -c
        - |
          set -ex
          DEVICE="/dev/${device}"

          # Skip if device doesn't exist
          if [ ! -b "\${DEVICE}" ]; then
            echo "Device \${DEVICE} not found, skipping"
            exit 0
          fi

          echo "=== Wiping \${DEVICE} on ${node} ==="

          # Remove any Ceph LVM volumes
          for vg in \$(vgs --noheadings -o vg_name 2>/dev/null | grep -i ceph || true); do
            echo "Removing VG: \${vg}"
            vgremove -ff "\${vg}" 2>/dev/null || true
          done

          # Remove PV if exists
          pvremove -ff "\${DEVICE}" 2>/dev/null || true

          # Remove device mapper entries
          dmsetup remove_all -f 2>/dev/null || true

          # Wipe filesystem signatures
          wipefs -af "\${DEVICE}"

          # Zap GPT/MBR
          sgdisk --zap-all "\${DEVICE}"

          # Zero first 100MB (clears any remaining metadata)
          dd if=/dev/zero of="\${DEVICE}" bs=1M count=100 conv=fsync

          # Zero last 100MB (clears backup GPT)
          SECTORS=\$(blockdev --getsz "\${DEVICE}")
          dd if=/dev/zero of="\${DEVICE}" bs=1M count=100 seek=\$((SECTORS/2048 - 100)) conv=fsync

          echo "=== Wipe complete for \${DEVICE} ==="
          lsblk -f "\${DEVICE}"
          blkid "\${DEVICE}" || echo "No signatures (clean)"
      volumeMounts:
        - name: dev
          mountPath: /dev
  volumes:
    - name: dev
      hostPath:
        path: /dev
EOF

  # Wait for pod to complete
  echo "  Waiting for wipe to complete on ${node}..."
  if ! kubectl -n "${NS}" wait --for=condition=Ready pod/"${pod_name}" --timeout=30s 2>/dev/null; then
    # Pod might have completed already
    true
  fi

  # Wait for completion (up to 2 minutes)
  local timeout=120
  local elapsed=0
  while [ $elapsed -lt $timeout ]; do
    local phase=$(kubectl -n "${NS}" get pod "${pod_name}" -o jsonpath='{.status.phase}' 2>/dev/null || echo "Unknown")
    if [ "$phase" = "Succeeded" ]; then
      echo "  ✓ Wipe completed on ${node}"
      kubectl -n "${NS}" logs "${pod_name}" 2>/dev/null | tail -5
      kubectl -n "${NS}" delete pod "${pod_name}" --ignore-not-found=true
      return 0
    elif [ "$phase" = "Failed" ]; then
      echo "  ✗ Wipe failed on ${node}"
      kubectl -n "${NS}" logs "${pod_name}" 2>/dev/null | tail -20
      kubectl -n "${NS}" delete pod "${pod_name}" --ignore-not-found=true
      return 1
    fi
    sleep 5
    elapsed=$((elapsed + 5))
  done

  echo "  ✗ Wipe timed out on ${node}"
  kubectl -n "${NS}" delete pod "${pod_name}" --ignore-not-found=true
  return 1
}

# Verify nodes are available
echo "Verifying node availability..."
NODE_COUNT=$(kubectl get nodes --no-headers 2>/dev/null | wc -l)
if [ "${NODE_COUNT}" -lt 1 ]; then
  echo "ERROR: No nodes available in cluster"
  exit 1
fi
echo "Found ${NODE_COUNT} node(s) available"

# Check for Ready nodes
READY_NODES=$(kubectl get nodes --no-headers 2>/dev/null | grep -c " Ready" || echo "0")
if [ "${READY_NODES}" -lt 1 ]; then
  echo "ERROR: No Ready nodes available in cluster"
  exit 1
fi
echo "Found ${READY_NODES} Ready node(s)"

# Check if Ceph OSDs already exist (skip wipe if so, unless forced)
FORCE_WIPE="${FORCE_WIPE:-false}"
if [ "$FORCE_WIPE" != "true" ]; then
  OSD_COUNT=$(kubectl -n "${NS}" get pods -l app=rook-ceph-osd --no-headers 2>/dev/null | grep -c Running || echo "0")
  if [ "$OSD_COUNT" -gt 0 ]; then
    echo ""
    echo "WARNING: Found ${OSD_COUNT} running OSD(s). Skipping disk wipe to protect existing data."
    echo "To force wipe, run: FORCE_WIPE=true ./scripts/rook-ceph-prepare.sh"
    echo ""
    echo "=== Rook-Ceph preparation complete (no wipe) ==="
    exit 0
  fi
fi

# Wipe NVMe drives on all configured nodes
echo ""
echo "=== Wiping NVMe drives ==="
WIPE_FAILED=0
for node in "${!NODE_NVME[@]}"; do
  device="${NODE_NVME[$node]}"

  # Check if node exists in cluster
  if ! kubectl get node "${node}" >/dev/null 2>&1; then
    echo "Node ${node} not found in cluster, skipping..."
    continue
  fi

  if ! wipe_nvme "${node}" "${device}"; then
    echo "WARNING: Failed to wipe ${device} on ${node}"
    WIPE_FAILED=1
  fi
done

if [ $WIPE_FAILED -eq 1 ]; then
  echo ""
  echo "WARNING: Some wipe operations failed. Check logs above."
  echo "Continuing with Ceph setup anyway..."
fi

# List nodes for informational purposes
echo ""
echo "Cluster nodes:"
kubectl get nodes -o wide --no-headers 2>/dev/null | while read -r line; do
  echo "  $line"
done

echo ""
echo "=== Rook-Ceph preparation complete ==="
echo "NVMe drives have been wiped and namespace is ready."
echo "Run 'helmfile apply' to deploy Ceph."