#!/usr/bin/env bash

set -euo pipefail

OLD_CONTEXT_DEFAULT="prole-k3s"
NEW_CONTEXT_DEFAULT="knoe.dev.prole.org"

usage() {
  cat <<'EOF'
Usage: k3s-cluster-rename.sh [--kubeconfig PATH] [--from OLD] [--to NEW] [--no-backup]

Rename/repair a k3s kubeconfig context so kubectx can switch to the new name.

Defaults:
  --from prole-k3s
  --to   knoe.dev.prole.org

Notes:
  - If --from is missing but kubeconfig has exactly one context, that context is used.
  - Default kubeconfig target is the active shell config source:
      - first path in $KUBECONFIG (if set)
      - otherwise ~/.kube/config
  - By default a timestamped backup is written next to kubeconfig.
EOF
}

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

die() {
  log "ERROR: $*"
  exit 1
}

require_cmd() {
  local cmd="$1"
  command -v "$cmd" >/dev/null 2>&1 || die "required command not found: $cmd"
}

KUBECONFIG_PATH=""
OLD_CONTEXT="$OLD_CONTEXT_DEFAULT"
NEW_CONTEXT="$NEW_CONTEXT_DEFAULT"
CREATE_BACKUP=1

while [[ $# -gt 0 ]]; do
  case "$1" in
    --kubeconfig)
      [[ $# -ge 2 ]] || die "missing value for --kubeconfig"
      KUBECONFIG_PATH="$2"
      shift 2
      ;;
    --from)
      [[ $# -ge 2 ]] || die "missing value for --from"
      OLD_CONTEXT="$2"
      shift 2
      ;;
    --to)
      [[ $# -ge 2 ]] || die "missing value for --to"
      NEW_CONTEXT="$2"
      shift 2
      ;;
    --no-backup)
      CREATE_BACKUP=0
      shift
      ;;
    -h|--help)
      usage
      exit 0
      ;;
    *)
      die "unknown argument: $1"
      ;;
  esac
done

require_cmd kubectl
require_cmd python3

if [[ -z "$KUBECONFIG_PATH" ]]; then
  if [[ -n "${KUBECONFIG:-}" ]]; then
    KUBECONFIG_PATH="${KUBECONFIG%%:*}"
    [[ -n "$KUBECONFIG_PATH" ]] || KUBECONFIG_PATH="${HOME}/.kube/config"
  else
    KUBECONFIG_PATH="${HOME}/.kube/config"
  fi
fi

[[ -n "$KUBECONFIG_PATH" ]] || die "could not resolve kubeconfig path"
[[ -f "$KUBECONFIG_PATH" ]] || die "kubeconfig not found: $KUBECONFIG_PATH"
[[ -r "$KUBECONFIG_PATH" ]] || die "kubeconfig is not readable: $KUBECONFIG_PATH"
[[ -w "$KUBECONFIG_PATH" ]] || die "kubeconfig is not writable: $KUBECONFIG_PATH"

tmp_input=$(mktemp "${TMPDIR:-/tmp}/k3s-cluster-rename-input.XXXXXX.json")
tmp_output=$(mktemp "${TMPDIR:-/tmp}/k3s-cluster-rename-output.XXXXXX.json")
cleanup() {
  rm -f "$tmp_input" "$tmp_output"
}
trap cleanup EXIT

KUBECONFIG="$KUBECONFIG_PATH" kubectl config view --raw -o json >"$tmp_input"

python3 - "$OLD_CONTEXT" "$NEW_CONTEXT" "$tmp_input" "$tmp_output" <<'PY'
import json
import sys

source_name, target_name, in_path, out_path = sys.argv[1:5]

with open(in_path, "r", encoding="utf-8") as fh:
    data = json.load(fh)

contexts = data.get("contexts") or []
context_names = [c.get("name", "") for c in contexts if isinstance(c, dict)]

def get_context_obj(name: str):
    for ctx in contexts:
        if isinstance(ctx, dict) and ctx.get("name") == name:
            return ctx
    return None

primary_name = None
if source_name in context_names:
    primary_name = source_name
elif target_name in context_names:
    primary_name = target_name
elif len(context_names) == 1:
    primary_name = context_names[0]
    sys.stderr.write(
        f"INFO: context '{source_name}' not found, using only context '{primary_name}' for repair.\n"
    )
else:
    sys.stderr.write(
        f"ERROR: context '{source_name}' not found and auto-repair is ambiguous (contexts={context_names}).\n"
    )
    sys.exit(2)

target_exists = target_name in context_names
if not target_exists and primary_name != target_name:
    primary_ctx = get_context_obj(primary_name)
    if primary_ctx is None:
        sys.stderr.write(f"ERROR: could not locate primary context '{primary_name}'.\n")
        sys.exit(2)
    primary_ctx["name"] = target_name
    if data.get("current-context") == primary_name:
        data["current-context"] = target_name

canonical_ctx = get_context_obj(target_name)
if canonical_ctx is None:
    canonical_ctx = get_context_obj(primary_name)

if canonical_ctx is None:
    sys.stderr.write("ERROR: failed to identify context to normalize.\n")
    sys.exit(2)

ctx_body = canonical_ctx.get("context") or {}
old_cluster = ctx_body.get("cluster")
old_user = ctx_body.get("user")

clusters = data.get("clusters") or []
users = data.get("users") or []
contexts = data.get("contexts") or []

cluster_names = {c.get("name") for c in clusters if isinstance(c, dict)}
user_names = {u.get("name") for u in users if isinstance(u, dict)}

if old_cluster and old_cluster != target_name and target_name not in cluster_names:
    for cluster in clusters:
        if isinstance(cluster, dict) and cluster.get("name") == old_cluster:
            cluster["name"] = target_name
            break
    for ctx in contexts:
        body = ctx.get("context") if isinstance(ctx, dict) else None
        if isinstance(body, dict) and body.get("cluster") == old_cluster:
            body["cluster"] = target_name
elif old_cluster:
    ctx_body["cluster"] = target_name if target_name in cluster_names else old_cluster

if old_user and old_user != target_name and target_name not in user_names:
    for user in users:
        if isinstance(user, dict) and user.get("name") == old_user:
            user["name"] = target_name
            break
    for ctx in contexts:
        body = ctx.get("context") if isinstance(ctx, dict) else None
        if isinstance(body, dict) and body.get("user") == old_user:
            body["user"] = target_name
elif old_user:
    ctx_body["user"] = target_name if target_name in user_names else old_user

data["current-context"] = target_name

with open(out_path, "w", encoding="utf-8") as fh:
    json.dump(data, fh, indent=2)
    fh.write("\n")
PY

found_context=0
while IFS= read -r ctx_name; do
  if [[ "$ctx_name" == "$NEW_CONTEXT" ]]; then
    found_context=1
    break
  fi
done < <(KUBECONFIG="$tmp_output" kubectl config get-contexts -o name)

[[ "$found_context" -eq 1 ]] || die "rename failed: context '$NEW_CONTEXT' not present in transformed kubeconfig"

if [[ "$CREATE_BACKUP" -eq 1 ]]; then
  backup_path="${KUBECONFIG_PATH}.bak.$(date +%Y%m%d%H%M%S)"
  cp "$KUBECONFIG_PATH" "$backup_path"
  log "Backup written: $backup_path"
fi

cp "$tmp_output" "$KUBECONFIG_PATH"
chmod 600 "$KUBECONFIG_PATH" 2>/dev/null || true

log "Kubeconfig updated: $KUBECONFIG_PATH"
log "Context/cluster/user target: $NEW_CONTEXT"

if command -v kubectx >/dev/null 2>&1; then
  if KUBECONFIG="$KUBECONFIG_PATH" kubectx "$NEW_CONTEXT" >/dev/null 2>&1; then
    log "kubectx switch verification succeeded for '$NEW_CONTEXT' (using KUBECONFIG=$KUBECONFIG_PATH)."
  else
    die "kubectx could not switch to '$NEW_CONTEXT' for kubeconfig '$KUBECONFIG_PATH'"
  fi
else
  log "kubectx not found; skipped kubectx switch verification."
fi

log "Done."