#!/usr/bin/env bash


usage() {
    cat <<EOF
Usage: $(basename "$0") REPLICAS [RESOURCE...] [OPTIONS]

Scale up or down one or several deployments/statefulsets.

Arguments:
  REPLICAS        Number of replicas to scale to
  RESOURCE        Resource to scale in "namespace/name" format, or just "name"
                  (uses current context namespace). Can be specified multiple times.

Commands:
  --all           Scale all deployments and statefulsets in the namespace
  --all-namespaces
                  Scale all deployments and statefulsets across all namespaces

Options:
  -h, --help      Show this help message
  -n, --namespace NAMESPACE
                  Namespace to use (default: current context namespace)
  --dry-run       Print what would be scaled without making changes
  -v              Pass -v to kubectl (minimal output)
  -vv             Pass -vv to kubectl (more output)
  -vvv            Pass -vvv to kubectl (debug output)

Examples:
  $(basename "$0") 1 mynamespace/mydeployment
  $(basename "$0") 1 mynamespace/mydeployment mydeployment2
  $(basename "$0") 1 --all
  $(basename "$0") 1 --all --namespace mynamespace
  $(basename "$0") 0 --all-namespaces --dry-run
EOF
}

KUBECTL_V=""
NAMESPACE=""
DRY_RUN=false
REPLICAS=""
RESOURCES=()
ALL=false
ALL_NAMESPACES=false

while [[ $# -gt 0 ]]; do
    case "$1" in
        -h|--help)
            usage
            exit 0
            ;;
        -n|--namespace)
            NAMESPACE="$2"
            shift 2
            ;;
        --dry-run)
            DRY_RUN=true
            shift
            ;;
        -v|-vv|-vvv)
            KUBECTL_V="$1"
            shift
            ;;
        --all)
            ALL=true
            shift
            ;;
        --all-namespaces)
            ALL_NAMESPACES=true
            shift
            ;;
        -*)
            echo "Error: Unknown option: $1" >&2
            usage >&2
            exit 1
            ;;
        *)
            if [[ -z "$REPLICAS" ]]; then
                REPLICAS="$1"
            else
                RESOURCES+=("$1")
            fi
            shift
            ;;
    esac
done

if [[ -z "$REPLICAS" ]]; then
    echo "Error: REPLICAS is required" >&2
    usage >&2
    exit 1
fi

if [[ "$ALL" == false && "$ALL_NAMESPACES" == false && ${#RESOURCES[@]} -eq 0 ]]; then
    echo "Error: Must specify --all, --all-namespaces, or at least one RESOURCE" >&2
    usage >&2
    exit 1
fi

NAMESPACE_ARG=()
if [[ -n "$NAMESPACE" ]]; then
    NAMESPACE_ARG=("-n" "$NAMESPACE")
fi

DRY_RUN_ARG=()
if [[ "$DRY_RUN" == true ]]; then
    DRY_RUN_ARG=("--dry-run=client")
fi

KUBECTL_BASE=(kubectl)
if [[ -n "$KUBECTL_V" ]]; then
    KUBECTL_BASE+=( "$KUBECTL_V" )
fi
KUBECTL_BASE+=( "${NAMESPACE_ARG[@]}" )
KUBECTL_BASE+=( "${DRY_RUN_ARG[@]}" )

scale_resource() {
    local resource="$1"
    local ns name

    if [[ "$resource" == */* ]]; then
        ns="${resource%%/*}"
        name="${resource#*/}"
    else
        ns="${NAMESPACE:-$(kubectl "${NAMESPACE_ARG[@]}" config view --minify --output jsonpath='{.contexts[0].context.namespace}' 2>/dev/null || echo "default")}"
        name="$resource"
    fi

    for kind in deployment statefulset; do
        if "${KUBECTL_BASE[@]}" get "$kind" "$name" -n "$ns" &>/dev/null; then
            echo "Scaling $kind/$ns/$name to $REPLICAS replicas${DRY_RUN:+ (dry-run)}"
            "${KUBECTL_BASE[@]}" scale "$kind" "$name" -n "$ns" --replicas="$REPLICAS"
            return 0
        fi
    done

    echo "Error: Resource '$resource' not found as deployment or statefulset" >&2
    return 1
}

get_resources() {
    local ns_flag=()
    if [[ "$ALL_NAMESPACES" == true ]]; then
        ns_flag=("--all-namespaces")
    elif [[ -n "$NAMESPACE" ]]; then
        ns_flag=("-n" "$NAMESPACE")
    fi

    "${KUBECTL_BASE[@]}" get "${ns_flag[@]}" deployment,statefulset -o jsonpath='{range .items[*]}{.metadata.namespace}/{.kind}/{.metadata.name}{"\n"}{end}' 2>/dev/null | while IFS=/ read -r ns kind name; do
        echo "$ns/$name"
    done
}

if [[ "$ALL" == true || "$ALL_NAMESPACES" == true ]]; then
    while IFS= read -r resource; do
        [[ -n "$resource" ]] && scale_resource "$resource"
    done < <(get_resources)
else
    for resource in "${RESOURCES[@]}"; do
        scale_resource "$resource"
    done
fi