#!/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)

# Shared option parsing for common core scripts
# shellcheck disable=SC1090
source "$SCRIPT_DIR/common_core_lib.sh"

common_core_preparse_config "$@"

# shellcheck disable=SC1090
source "$SCRIPT_DIR/prole_cfg.sh"

set -- "${COMMON_CORE_ARGS[@]}"
common_core_parse_args "$@"

if [[ "${COMMON_CORE_HELP:-0}" == 1 ]]; then
  common_core_usage "$0"
  exit 0
fi

if [[ -n "${COMMON_CORE_PARSE_ERROR:-}" ]]; then
  echo "ERROR: ${COMMON_CORE_PARSE_ERROR}" >&2
  common_core_usage "$0"
  exit 2
fi

ACTION="$COMMON_CORE_ACTION"

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

RESOLVED_NAMESPACE="$(common_core_resolve_namespace "default")"
common_core_apply_namespace "$RESOLVED_NAMESPACE"

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}
GARAGE_NAMESPACE=${GARAGE_NAMESPACE:-$RESOLVED_NAMESPACE}
NAMESPACE="$GARAGE_NAMESPACE"

# 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/storageclass-prole-iscsi.yaml"
  "$GARAGE_MANIFEST_DIR/iscsi-pvs.yaml"
  "$GARAGE_MANIFEST_DIR/garage-configmap.yaml"
  "$GARAGE_MANIFEST_DIR/garage-statefulset.yaml"
  "$GARAGE_MANIFEST_DIR/garage-service.yaml"
)
if [[ "${PROLE_MODE:-}" == "k3d" ]]; then
  GARAGE_FILES=(
    "$GARAGE_MANIFEST_DIR/garage-configmap.yaml"
    "$GARAGE_MANIFEST_DIR/garage-statefulset.yaml"
    "$GARAGE_MANIFEST_DIR/garage-service.yaml"
  )
fi

GARAGE_APPLY_CHANGED=0
GARAGE_CONFIG_CHANGED=0
GARAGE_STATEFULSET_CHANGED=0

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' ..."
  # Wait for apiserver to be ready
  local count=0
  while ! kubectl get nodes >/dev/null 2>&1; do
      if [ $count -ge 30 ]; then
          echo "ERROR: apiserver not ready after 60s" >&2
          exit 1
      fi
      echo "Waiting for apiserver... ($count/30)"
      sleep 2
      count=$((count + 1))
  done

  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"
}

k3d_cleanup_pending_pvc() {
  if [[ "${PROLE_MODE:-}" != "k3d" ]]; then
    return 0
  fi
  if ! kubectl get pvc data-garage-0 -n "$NAMESPACE" >/dev/null 2>&1; then
    return 0
  fi
  local phase selector
  phase=$(kubectl get pvc data-garage-0 -n "$NAMESPACE" -o jsonpath='{.status.phase}' 2>/dev/null || true)
  if [[ "$phase" == "Bound" ]]; then
    return 0
  fi
  selector=$(kubectl get pvc data-garage-0 -n "$NAMESPACE" -o jsonpath='{.spec.selector}' 2>/dev/null || true)
  if printf '%s' "$selector" | grep -q "prole.storage/"; then
    echo "Removing pending Garage PVC with selector for k3d local-path ..."
    kubectl delete statefulset "$GARAGE_NAME" -n "$NAMESPACE" --ignore-not-found >/dev/null 2>&1 || true
    kubectl delete pvc data-garage-0 -n "$NAMESPACE" --ignore-not-found >/dev/null 2>&1 || true
  fi
}

apply_manifests() {
  GARAGE_APPLY_CHANGED=0
  GARAGE_CONFIG_CHANGED=0
  GARAGE_STATEFULSET_CHANGED=0
  for f in "${GARAGE_FILES[@]}"; do
    if [[ -f "$f" ]]; then
      local output=""
      local rendered=""
      local diff_out=""
      local diff_rc=0
      rendered=$(prole_render_manifest "$f")
      if diff_out=$(printf '%s' "$rendered" | kubectl diff -n "$NAMESPACE" -f - 2>&1); then
        diff_rc=0
      else
        diff_rc=$?
      fi
      if [[ $diff_rc -eq 0 ]]; then
        continue
      fi
      if [[ $diff_rc -ne 1 ]]; then
        echo "$diff_out" >&2
        exit 1
      fi
      if output=$(printf '%s' "$rendered" | kubectl apply --validate=false -n "$NAMESPACE" -f - 2>&1); then
        printf '%s\n' "$output"
        GARAGE_APPLY_CHANGED=1
        case "$(basename "$f")" in
          garage-configmap.yaml) GARAGE_CONFIG_CHANGED=1 ;;
          garage-statefulset.yaml) GARAGE_STATEFULSET_CHANGED=1 ;;
        esac
      else
        if [[ "${PROLE_MODE:-}" == "k3d" && "$(basename "$f")" == "garage-statefulset.yaml" ]] \
          && echo "$output" | grep -q "updates to statefulset spec"; then
          echo "WARN: Garage StatefulSet immutable in k3d; skipping apply."
          continue
        fi
        echo "$output" >&2
        exit 1
      fi
    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
}

pin_garage_node_k3d() {
  if [[ "${PROLE_MODE:-}" != "k3d" ]]; then
    return 0
  fi
  if [[ -z "${GARAGE_PIN_NODE:-}" ]]; then
    return 0
  fi
  if ! kubectl get statefulset "$GARAGE_NAME" -n "$NAMESPACE" >/dev/null 2>&1; then
    return 0
  fi
  local target=""
  target=$(kubectl get nodes -l node-role.kubernetes.io/master -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)
  if [[ -z "$target" ]]; then
    target=$(kubectl get nodes -l node-role.kubernetes.io/control-plane -o jsonpath='{.items[0].metadata.name}' 2>/dev/null || true)
  fi
  if [[ -z "$target" ]]; then
    return 0
  fi
  if [[ "$target" != k3d-* ]]; then
    return 0
  fi
  echo "Pinning Garage StatefulSet to node '$target' for k3d exec access ..."
  kubectl patch statefulset "$GARAGE_NAME" -n "$NAMESPACE" --type merge \
    -p "{\"spec\":{\"template\":{\"spec\":{\"nodeSelector\":{\"kubernetes.io/hostname\":\"$target\"}}}}}" >/dev/null || true
}

restart_statefulset() {
  if [[ "${GARAGE_CONFIG_CHANGED:-0}" -eq 0 ]]; then
    return 0
  fi
  if [[ "${GARAGE_STATEFULSET_CHANGED:-0}" -eq 1 ]]; then
    return 0
  fi
  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=${ROLLOUT_TIMEOUT:-300s}; 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
    
    # Re-fetch layout and apply staged changes using the NEXT layout version
    local layout_after version next_version apply_ok
    layout_after=$(kubectl exec -n "$NAMESPACE" "$pod" -- /garage layout show 2>/dev/null | grep -v "INFO" || true)
    # Extract "Current cluster layout version: X"
    version=$(echo "$layout_after" | awk -F: '/Current cluster layout version/ {gsub(/[[:space:]]/,"",$2); print $2; exit}' || true)
    
    if echo "$layout_after" | grep -qi "staged"; then
      if [[ -n "$version" ]]; then
        next_version=$((version + 1))
        echo "Applying Garage layout (current version $version, next version $next_version) ..."
        if kubectl exec -n "$NAMESPACE" "$pod" -- /garage layout apply --version "$next_version"; then
          apply_ok=1
        else
          apply_ok=0
        fi
      else
        apply_ok=0
      fi
      if [[ "${apply_ok:-0}" -eq 0 ]]; then
        echo "WARN: layout apply with explicit version failed; retrying without version..."
        kubectl exec -n "$NAMESPACE" "$pod" -- /garage layout apply || true
      fi
    else
      echo "No staged layout changes detected; skipping layout apply."
    fi
  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|initialize|update|reload|restart)
    ensure_tools
    ensure_namespace
    ensure_secrets
    k3d_cleanup_pending_pvc
    echo "Applying Garage manifests in namespace '$NAMESPACE'..."
    apply_manifests
    ensure_container_command
    pin_garage_node_k3d
    restart_statefulset
    wait_ready
    init_layout
    prole_register_port_forward "garage" "${NAMESPACE:-default}" "svc/garage" "3900" "3900" "0.0.0.0" "TCP" "Garage S3 API"
    ;;
  stop)
    ensure_tools
    echo "Deleting Garage manifests from namespace '$NAMESPACE'..."
    delete_manifests
    ;;
  status)
    status
    ;;
  *)
    common_core_usage "$0"
    exit 1
    ;;
esac