#!/usr/bin/env bash

set -euo pipefail

# init_prole-db-reset.sh
# Purpose:
# - Orderly shutdown of prole-db with mandatory backup/export
# - Reset namespace resources
# - Recreate prole-db from backup if available, else fresh init
#
# Usage:
#   ./init_prole-db-reset.sh reset      # shutdown + reset + recreate
#   ./init_prole-db-reset.sh shutdown   # orderly shutdown only
#   ./init_prole-db-reset.sh recreate   # recreate from backup if available

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"

if [[ "${1:-}" == "--mode" || "${1:-}" == "-m" ]]; then
  prole_set_mode "${2:-}"
  shift 2
elif [[ "${1:-}" == --mode=* || "${1:-}" == -m=* ]]; then
  prole_set_mode "${1#*=}"
  shift
fi

ACTION=${1:-reset}

NAMESPACE=${NAMESPACE:-default}
CNPG_CLUSTER_NAME=${CNPG_CLUSTER_NAME:-prole-db}
PROLE_HOME=${PROLE_HOME:-$(cd "$SCRIPT_DIR/.." && pwd)}

ADMIN_USER=${PROLE_DB_ADMIN_USER:-postgres}
ADMIN_DB=${PROLE_DB_ADMIN_DB:-postgres}

GARAGE_NAME=${GARAGE_NAME:-garage}
SERVICE_NAMESPACE=${SERVICE_NAMESPACE:-}
if [[ -z "${GARAGE_NAMESPACE:-}" ]]; then
  if [[ -n "$SERVICE_NAMESPACE" ]]; then
    GARAGE_NAMESPACE="$SERVICE_NAMESPACE"
  else
    GARAGE_NAMESPACE="$NAMESPACE"
  fi
fi
GARAGE_BACKUP_BUCKET=${GARAGE_BACKUP_BUCKET:-prole-db-backups}
GARAGE_BACKUP_SECRET_NAME=${GARAGE_BACKUP_SECRET_NAME:-prole-db-barman-s3}
GARAGE_S3_ENDPOINT=${GARAGE_S3_ENDPOINT:-http://$GARAGE_NAME.$GARAGE_NAMESPACE.svc.cluster.local:3900}

BACKUP_DIR=${BACKUP_DIR:-$PROLE_HOME/prole/backup}
BACKUP_WAIT_TIMEOUT=${BACKUP_WAIT_TIMEOUT:-1800}
EXPORT_WAIT_TIMEOUT=${EXPORT_WAIT_TIMEOUT:-1800}
ALLOW_NO_BACKUP=${ALLOW_NO_BACKUP:-0}

RECOVERY_TEMPLATE="$SCRIPT_DIR/../k8s/prole/prole-db-recovery.yaml.tpl"

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

usage() {
  cat <<EOF
Usage: $0 {reset|shutdown|recreate}

Actions:
  reset      Orderly shutdown + reset namespace + recreate (default)
  shutdown   Orderly shutdown only (backup + export + reset namespace)
  recreate   Recreate from backup if available, else fresh init

Env overrides:
  BACKUP_DIR, BACKUP_WAIT_TIMEOUT, EXPORT_WAIT_TIMEOUT, ALLOW_NO_BACKUP
  PROLE_DB_ADMIN_USER, PROLE_DB_ADMIN_DB
  GARAGE_BACKUP_BUCKET, GARAGE_BACKUP_SECRET_NAME, GARAGE_S3_ENDPOINT
EOF
}

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

cluster_exists() {
  kubectl -n "$NAMESPACE" get cluster "$CNPG_CLUSTER_NAME" >/dev/null 2>&1
}

get_primary_pod() {
  local primary
  primary=$(kubectl -n "$NAMESPACE" get cluster "$CNPG_CLUSTER_NAME" -o jsonpath='{.status.currentPrimary}' 2>/dev/null || true)
  if [[ -z "$primary" ]]; then
    primary=$(kubectl -n "$NAMESPACE" get pods -l "cnpg.io/cluster=$CNPG_CLUSTER_NAME" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)
  fi
  printf '%s' "$primary"
}


terminate_non_admin_connections() {
  local pod
  pod=$(get_primary_pod)
  if [[ -z "$pod" ]]; then
    log "No primary CNPG pod detected; skipping connection cleanup."
    return 0
  fi

  log "Terminating non-admin database connections (admin=$ADMIN_USER) ..."
  kubectl -n "$NAMESPACE" exec "$pod" -c postgres -- \
    psql -U "$ADMIN_USER" -d "$ADMIN_DB" -v ON_ERROR_STOP=1 -c \
    "SELECT pg_terminate_backend(pid)
     FROM pg_stat_activity
     WHERE usename <> '${ADMIN_USER}'
       AND pid <> pg_backend_pid();" || true
}

latest_backup_name() {
  kubectl -n "$NAMESPACE" get backup \
    --sort-by=.metadata.creationTimestamp \
    -o jsonpath='{range .items[*]}{.metadata.name}{"\n"}{end}' 2>/dev/null | tail -n 1
}

latest_completed_backup_name() {
  kubectl -n "$NAMESPACE" get backup \
    --sort-by=.metadata.creationTimestamp \
    -o jsonpath='{range .items[*]}{.metadata.name}{"|"}{.status.phase}{"\n"}{end}' 2>/dev/null | \
    awk -F'|' '{p=tolower($2); if (p=="completed" || p=="succeeded") {name=$1}} END {print name}'
}

wait_for_backup() {
  local backup_name="$1"
  local start_time now phase phase_lc
  start_time=$(date +%s)

  while true; do
    phase=$(kubectl -n "$NAMESPACE" get backup "$backup_name" -o jsonpath='{.status.phase}' 2>/dev/null || true)
    phase_lc=$(printf '%s' "$phase" | tr '[:upper:]' '[:lower:]')

    case "$phase_lc" in
      completed|succeeded)
        log "Backup $backup_name completed."
        return 0
        ;;
      failed|error)
        err "Backup $backup_name failed (phase=$phase)."
        return 1
        ;;
    esac

    now=$(date +%s)
    if (( now - start_time > BACKUP_WAIT_TIMEOUT )); then
      err "Timed out waiting for backup $backup_name."
      return 1
    fi

    log "Waiting for backup $backup_name to complete (phase=${phase:-unknown}) ..."
    sleep 10
  done
}

run_barman_backup() {
  if ! cluster_exists; then
    err "Cluster '$CNPG_CLUSTER_NAME' not found; cannot run backup."
    return 1
  fi

  if [[ -x "$SCRIPT_DIR/init_prole-db-backup.sh" ]]; then
    log "Running barman backup (Garage bucket: $GARAGE_BACKUP_BUCKET) ..."
    NAMESPACE="$NAMESPACE" \
    GARAGE_BACKUP_BUCKET="$GARAGE_BACKUP_BUCKET" \
    GARAGE_BACKUP_SECRET_NAME="$GARAGE_BACKUP_SECRET_NAME" \
    GARAGE_S3_ENDPOINT="$GARAGE_S3_ENDPOINT" \
    RUN_FIRST_BACKUP=1 "$SCRIPT_DIR/init_prole-db-backup.sh" start
  else
    err "init_prole-db-backup.sh not found; cannot run backup."
    return 1
  fi

  sleep 2
  local backup_name
  backup_name=$(latest_backup_name)
  if [[ -z "$backup_name" ]]; then
    err "No backup resource detected after triggering backup."
    return 1
  fi

  wait_for_backup "$backup_name"
}

export_backup_bucket() {
  if [[ -z "$GARAGE_BACKUP_BUCKET" ]]; then
    err "GARAGE_BACKUP_BUCKET is not set; cannot export backup."
    return 1
  fi
  if ! kubectl -n "$NAMESPACE" get secret "$GARAGE_BACKUP_SECRET_NAME" >/dev/null 2>&1; then
    err "Garage backup secret '$GARAGE_BACKUP_SECRET_NAME' not found; cannot export backup."
    return 1
  fi

  local image pod_name export_subdir
  image=$(kubectl -n "$NAMESPACE" get cluster "$CNPG_CLUSTER_NAME" -o jsonpath='{.spec.imageName}' 2>/dev/null || true)
  image=${image:-${PROLE_DB_IMAGE:-prole-db:latest}}

  pod_name="prole-barman-export-$(date +%s)"
  export_subdir="${BACKUP_DIR}/garage-export-$(date +%Y%m%d%H%M%S)"

  log "Creating export pod $pod_name using image $image ..."
  cat <<EOF | kubectl apply -n "$NAMESPACE" -f -
apiVersion: v1
kind: Pod
metadata:
  name: ${pod_name}
  namespace: ${NAMESPACE}
spec:
  restartPolicy: Never
  containers:
    - name: exporter
      image: ${image}
      command: ["sleep","3600"]
      env:
        - name: AWS_ACCESS_KEY_ID
          valueFrom:
            secretKeyRef:
              name: ${GARAGE_BACKUP_SECRET_NAME}
              key: ACCESS_KEY_ID
        - name: AWS_SECRET_ACCESS_KEY
          valueFrom:
            secretKeyRef:
              name: ${GARAGE_BACKUP_SECRET_NAME}
              key: SECRET_ACCESS_KEY
        - name: AWS_DEFAULT_REGION
          value: us-east-1
        - name: GARAGE_S3_ENDPOINT
          value: ${GARAGE_S3_ENDPOINT}
        - name: GARAGE_BACKUP_BUCKET
          value: ${GARAGE_BACKUP_BUCKET}
        - name: EXPORT_DIR
          value: /backup
      volumeMounts:
        - name: backup
          mountPath: /backup
  volumes:
    - name: backup
      emptyDir: {}
EOF
  trap 'kubectl -n "$NAMESPACE" delete pod "$pod_name" --ignore-not-found >/dev/null 2>&1 || true' RETURN

  log "Waiting for export pod to become ready..."
  if ! kubectl -n "$NAMESPACE" wait --for=condition=Ready pod/"$pod_name" --timeout="${EXPORT_WAIT_TIMEOUT}s" >/dev/null 2>&1; then
    err "Export pod did not become ready."
    kubectl -n "$NAMESPACE" describe pod "$pod_name" || true
    return 1
  fi

  log "Exporting Garage bucket to pod filesystem..."
  kubectl -n "$NAMESPACE" exec "$pod_name" -- sh -c 'python3 - <<PY
import os
import sys
try:
    import boto3
    from botocore.config import Config
except Exception as exc:
    print(f"boto3 is required in the export image: {exc}", file=sys.stderr)
    sys.exit(3)

endpoint = os.environ.get("GARAGE_S3_ENDPOINT")
bucket_name = os.environ.get("GARAGE_BACKUP_BUCKET")
export_dir = os.environ.get("EXPORT_DIR", "/backup")

if not endpoint or not bucket_name:
    print("Missing GARAGE_S3_ENDPOINT or GARAGE_BACKUP_BUCKET", file=sys.stderr)
    sys.exit(2)

os.makedirs(export_dir, exist_ok=True)

config = Config(s3={"addressing_style": "path"})
s3 = boto3.resource("s3", endpoint_url=endpoint, config=config)
bucket = s3.Bucket(bucket_name)

count = 0
for obj in bucket.objects.all():
    key = obj.key
    if key.endswith("/"):
        continue
    dest = os.path.join(export_dir, key)
    os.makedirs(os.path.dirname(dest), exist_ok=True)
    bucket.download_file(key, dest)
    count += 1

print(f"Downloaded {count} objects from {bucket_name} to {export_dir}")
PY'

  log "Copying exported backup to $export_subdir ..."
  mkdir -p "$export_subdir"
  kubectl -n "$NAMESPACE" cp "$pod_name":/backup "$export_subdir"

  log "Cleaning up export pod $pod_name ..."
  kubectl -n "$NAMESPACE" delete pod "$pod_name" --ignore-not-found
  trap - RETURN
}

shutdown_database() {
  if [[ -x "$SCRIPT_DIR/init_cloudnative_pg.sh" ]]; then
    log "Stopping prole-db cluster ..."
    "$SCRIPT_DIR/init_cloudnative_pg.sh" stop || true
  else
    log "init_cloudnative_pg.sh not found; deleting cluster resource directly."
    kubectl -n "$NAMESPACE" delete cluster "$CNPG_CLUSTER_NAME" --ignore-not-found
  fi
}

reset_namespace() {
  if [[ -x "$SCRIPT_DIR/../scripts/reset-ns.sh" ]]; then
    log "Resetting namespace $NAMESPACE ..."
    "$SCRIPT_DIR/../scripts/reset-ns.sh" -n "$NAMESPACE"
  else
    err "scripts/reset-ns.sh not found."
    return 1
  fi
}

recreate_from_backup() {
  local backup_name manifest_tmp
  backup_name=$(latest_completed_backup_name)
  if [[ -z "$backup_name" ]]; then
    return 1
  fi

  if [[ ! -f "$RECOVERY_TEMPLATE" ]]; then
    err "Recovery template not found: $RECOVERY_TEMPLATE"
    return 1
  fi

  manifest_tmp="/tmp/prole-db-recovery-$$.yaml"
  sed "s/{{BACKUP_NAME}}/${backup_name}/g" "$RECOVERY_TEMPLATE" > "$manifest_tmp"

  log "Recreating prole-db from backup $backup_name ..."
  CNPG_MANIFEST_OVERRIDE="$manifest_tmp" "$SCRIPT_DIR/init_cloudnative_pg.sh" initialize

  rm -f "$manifest_tmp"
}

recreate_fresh() {
  log "Recreating prole-db from scratch ..."
  "$SCRIPT_DIR/init_cloudnative_pg.sh" initialize
}

shutdown() {
  ensure_tools
  terminate_non_admin_connections

  if run_barman_backup; then
    export_backup_bucket
  else
    if [[ "$ALLOW_NO_BACKUP" == "1" ]]; then
      err "Backup failed, but ALLOW_NO_BACKUP=1 set; proceeding without export."
    else
      err "Backup failed and ALLOW_NO_BACKUP is not set; aborting shutdown."
      exit 1
    fi
  fi

  shutdown_database
  reset_namespace
}

recreate() {
  ensure_tools
  if recreate_from_backup; then
    log "Recovery initiated from existing backup."
  else
    log "No completed backups found; starting fresh."
    recreate_fresh
  fi
}

case "$ACTION" in
  shutdown)
    shutdown
    ;;
  recreate)
    recreate
    ;;
  reset)
    shutdown
    recreate
    ;;
  *)
    usage
    exit 2
    ;;
esac