#!/usr/bin/env bash

set -euo pipefail

# init_garage_store.sh
# Purpose:
# - Deploy Garage (S3-compatible object store) in Kubernetes
# - Initialize single-node layout for immediate use

# Initialize SCRIPT_DIR
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:-}

if [[ -n "${GARAGE_INIT_LOG:-}" ]]; then
  mkdir -p "$(dirname "$GARAGE_INIT_LOG")"
  exec > >(tee -a "$GARAGE_INIT_LOG") 2>&1
fi

GARAGE_NAME=${GARAGE_NAME:-garage}
GARAGE_SECRET_NAME=${GARAGE_SECRET_NAME:-garage-secrets}
GARAGE_NODE_CAPACITY=${GARAGE_NODE_CAPACITY:-10GB}
GARAGE_ZONE=${GARAGE_ZONE:-local}

# Support both PROLE_HOME/k8s and sibling k8s directory
if [[ -d "$SCRIPT_DIR/../k8s/prole" ]]; then
  GARAGE_MANIFEST_DIR="$SCRIPT_DIR/../k8s/prole"
elif [[ -n "${PROLE_HOME:-}" && -d "$PROLE_HOME/k8s/prole" ]]; then
  GARAGE_MANIFEST_DIR="$PROLE_HOME/k8s/prole"
else
  GARAGE_MANIFEST_DIR="$SCRIPT_DIR/../k8s/prole"
fi

GARAGE_FILES=(
  "$GARAGE_MANIFEST_DIR/garage-configmap.yaml"
  "$GARAGE_MANIFEST_DIR/garage-statefulset.yaml"
  "$GARAGE_MANIFEST_DIR/garage-service.yaml"
)

usage() {
  cat <<USAGE
Usage: $0 [start|stop|restart|status]

Actions:
  start     Apply Garage manifests and initialize layout
  stop      Delete Garage manifests
  restart   Re-apply Garage manifests and initialize layout
  status    Show Garage resources
USAGE
  exit 1
}

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

ensure_namespace() {
  if ! kubectl get namespace "$NAMESPACE" >/dev/null 2>&1; then
    echo "Creating namespace '$NAMESPACE' ..."
    kubectl create namespace "$NAMESPACE" >/dev/null 2>&1 || true
  fi
}

ensure_secrets() {
  if kubectl get secret "$GARAGE_SECRET_NAME" -n "$NAMESPACE" >/dev/null 2>&1; then
    return 0
  fi

  echo "Creating Garage secrets in namespace '$NAMESPACE' ..."
  local rpc_secret admin_token metrics_token
  rpc_secret=$(openssl rand -hex 32)
  admin_token=$(openssl rand -base64 32)
  metrics_token=$(openssl rand -base64 32)

  kubectl create secret generic "$GARAGE_SECRET_NAME" -n "$NAMESPACE" \
    --from-literal=rpc_secret="$rpc_secret" \
    --from-literal=admin_token="$admin_token" \
    --from-literal=metrics_token="$metrics_token"
}

apply_manifests() {
  for f in "${GARAGE_FILES[@]}"; do
    if [[ -f "$f" ]]; then
      kubectl apply -n "$NAMESPACE" -f "$f"
    else
      echo "ERROR: Missing manifest: $f" >&2
      exit 1
    fi
  done
}

ensure_container_command() {
  if ! kubectl get statefulset "$GARAGE_NAME" -n "$NAMESPACE" >/dev/null 2>&1; then
    return 0
  fi

  local cmd args
  cmd=$(kubectl get statefulset "$GARAGE_NAME" -n "$NAMESPACE" -o jsonpath='{.spec.template.spec.containers[?(@.name=="garage")].command}' 2>/dev/null || true)
  args=$(kubectl get statefulset "$GARAGE_NAME" -n "$NAMESPACE" -o jsonpath='{.spec.template.spec.containers[?(@.name=="garage")].args}' 2>/dev/null || true)

  if [[ -z "$cmd" || "$cmd" == "[]" || "$cmd" != *"/garage"* || -z "$args" || "$args" == "[]" || "$args" != *"server"* ]]; then
    echo "Patching Garage container command/args ..."
    kubectl patch statefulset "$GARAGE_NAME" -n "$NAMESPACE" --type merge -p '{
      "spec": {
        "template": {
          "spec": {
            "containers": [
              {
                "name": "garage",
                "command": ["/garage"],
                "args": ["server"]
              }
            ]
          }
        }
      }
    }' >/dev/null || true
  fi
}

restart_statefulset() {
  if kubectl get statefulset "$GARAGE_NAME" -n "$NAMESPACE" >/dev/null 2>&1; then
    echo "Restarting Garage StatefulSet to pick up config changes ..."
    kubectl rollout restart statefulset/$GARAGE_NAME -n "$NAMESPACE" || true
  fi
}

delete_manifests() {
  for f in "${GARAGE_FILES[@]}"; do
    if [[ -f "$f" ]]; then
      kubectl delete -n "$NAMESPACE" -f "$f" --ignore-not-found
    fi
  done
}

dump_debug() {
  echo "---- Garage debug ----"
  kubectl get statefulset "$GARAGE_NAME" -n "$NAMESPACE" -o wide || true
  kubectl get pods -n "$NAMESPACE" -l "app=$GARAGE_NAME" -o wide || true
  kubectl get svc "$GARAGE_NAME" -n "$NAMESPACE" -o wide || true
  kubectl get events -n "$NAMESPACE" --sort-by=.metadata.creationTimestamp | tail -n 50 || true

  local pod
  pod=$(kubectl get pods -n "$NAMESPACE" -l "app=$GARAGE_NAME" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)
  if [[ -n "$pod" ]]; then
    echo "--- pod: $pod (describe) ---"
    kubectl describe pod "$pod" -n "$NAMESPACE" || true
    echo "--- logs (current) ---"
    kubectl logs -n "$NAMESPACE" "$pod" --tail=200 || true
    echo "--- logs (previous) ---"
    kubectl logs -n "$NAMESPACE" "$pod" --previous --tail=200 || true
  fi
  echo "---- Garage debug end ----"
}

wait_ready() {
  echo "Waiting for Garage StatefulSet to become ready ..."
  if ! kubectl rollout status statefulset/$GARAGE_NAME -n "$NAMESPACE" --timeout=180s; then
    echo "ERROR: Garage StatefulSet did not become ready in time." >&2
    dump_debug
    return 1
  fi
}

init_layout() {
  local pod
  pod=$(kubectl get pods -n "$NAMESPACE" -l "app=$GARAGE_NAME" -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)
  if [[ -z "$pod" ]]; then
    echo "ERROR: Garage pod not found in namespace '$NAMESPACE'." >&2
    exit 1
  fi

  local node_id=""
  echo "Determining Garage node ID..."
  for i in {1..30}; do
    node_id=$(kubectl exec -n "$NAMESPACE" "$pod" -- /garage node id -q 2>/dev/null | awk '{print $1}' || true)
    if [[ -n "$node_id" ]]; then
      break
    fi
    echo "Waiting for Garage node ID to be available... ($i/30)"
    sleep 2
  done

  if [[ -z "$node_id" ]]; then
    echo "ERROR: Unable to determine Garage node ID after 30 attempts." >&2
    dump_debug
    exit 1
  fi

  # Extract the short ID for better matching
  local short_id
  short_id=$(echo "$node_id" | cut -d'@' -f1)
  echo "Garage node ID: $short_id"

  local layout
  layout=$(kubectl exec -n "$NAMESPACE" "$pod" -- /garage layout show 2>/dev/null || true)
  if ! echo "$layout" | grep -q "$short_id"; then
    echo "Assigning Garage node role (capacity: $GARAGE_NODE_CAPACITY, zone: $GARAGE_ZONE) ..."
    # Retry assigning role as it might fail if node is not yet fully ready in the cluster logic
    local max_assign_retries=10
    local assign_count=0
    while ! kubectl exec -n "$NAMESPACE" "$pod" -- /garage layout assign -z "$GARAGE_ZONE" -c "$GARAGE_NODE_CAPACITY" "$short_id"; do
      if [[ $assign_count -ge $max_assign_retries ]]; then
        echo "ERROR: Failed to assign Garage node role after $max_assign_retries retries." >&2
        exit 1
      fi
      echo "Retrying Garage layout assign... ($((assign_count + 1))/$max_assign_retries)"
      sleep 2
      assign_count=$((assign_count + 1))
    done
    
    # Get current version for apply
    local version
    version=$(echo "$layout" | grep "Version:" | awk '{print $2}' || echo "0")
    if [[ -z "$version" ]]; then version=0; fi
    local next_version=$((version + 1))
    echo "Applying Garage layout (version $next_version) ..."
    kubectl exec -n "$NAMESPACE" "$pod" -- /garage layout apply --version "$next_version" || true
  else
    echo "Garage layout already assigned for node $short_id."
  fi
}

status() {
  ensure_tools
  echo "Garage status in namespace '$NAMESPACE':"
  kubectl get statefulset "$GARAGE_NAME" -n "$NAMESPACE" || true
  kubectl get pods -n "$NAMESPACE" -l "app=$GARAGE_NAME" || true
  kubectl get svc "$GARAGE_NAME" -n "$NAMESPACE" || true
}

case "$ACTION" in
  start)
    ensure_tools
    ensure_namespace
    ensure_secrets
    echo "Applying Garage manifests in namespace '$NAMESPACE'..."
    apply_manifests
    ensure_container_command
    restart_statefulset
    wait_ready
    init_layout
    ;;
  stop)
    ensure_tools
    echo "Deleting Garage manifests from namespace '$NAMESPACE'..."
    delete_manifests
    ;;
  restart)
    ensure_tools
    ensure_namespace
    ensure_secrets
    echo "Re-applying Garage manifests in namespace '$NAMESPACE'..."
    apply_manifests
    ensure_container_command
    restart_statefulset
    wait_ready
    init_layout
    ;;
  status)
    status
    ;;
  *)
    usage
    ;;
esac