#!/bin/bash

# Kubernetes Startup Script (Edge)
# This script scales up deployments in hubble-system and cluster-<uid> in a specific order

set -euo pipefail

# Default configuration
KUBECONFIG_FILE=""
CLUSTER_UID=""
IS_HELM=false
CLUSTER_NAMESPACE=""  # Set to cluster-<uid> after parsing

# hubble-system: scale to 1 first (all deployments that were scaled down)
HUBBLE_SYSTEM_TO_1=("spectrocluster" "spectrocluster-jobs" "timeseries" "cloud" "event" "hashboard" "hutil" "memstore" "user" "foreq" "mgmt" "spectroclusterop" "oci-proxy" "spectro-tunnel" "spectrocluster-reconciler" "spectrossh")
# hubble-system: then scale these to 2 replicas
HUBBLE_SYSTEM_TO_2=("cloud" "hashboard" "hutil" "spectroclusterop" "user")
# hubble-system: lastly scale these to 3 replicas
HUBBLE_SYSTEM_TO_3=("spectrocluster" "timeseries" "event")

# cluster-<uid>: scale to 1 replica (order: agent first, then controller)
CLUSTER_UID_DEPLOYMENTS=("cluster-management-agent-lite" "palette-lite-controller-manager")

# Colors for output
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

# Logging functions
log_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

log_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

log_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

log_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# Show help message
show_help() {
    cat << EOF
Kubernetes Startup Script (Edge)

This script performs the following operations in order:

hubble-system namespace:
  1. Scale msgbroker statefulset to 2 replicas and wait for pods to become ready
  2. Scale configserver deployment to 1 replica and wait for pods to be Running
  3. Scale all scaled-down deployments to 1 replica
  4. Scale auth deployment to 2 replicas and wait for pods to become ready
  5. Scale system deployment to 2 replicas and wait for pods to become ready
  6. Scale cloud, hashboard, hutil, spectroclusterop, user to 2 replicas
  7. Scale spectrocluster, timeseries, event to 3 replicas

cluster-<uid> namespace:
  8. Scale cluster-management-agent-lite and palette-lite-controller-manager to 1 replica

Usage: $0 --kubeconfig FILE [--uid UID] [--is-helm]

OPTIONS:
    --kubeconfig FILE    Path to kubeconfig file (REQUIRED)
    --uid UID            Cluster UID; namespace will be cluster-<UID> (REQUIRED unless --is-helm)
    --is-helm            Helm mode (default: false). Skips cluster-<uid> deployment scale-up.

EXAMPLES:
    $0 --kubeconfig /path/to/kubeconfig --uid abc123
    $0 --kubeconfig /path/to/kubeconfig --is-helm
EOF
}

# Parse command-line arguments
parse_arguments() {
    while [[ $# -gt 0 ]]; do
        case $1 in
            --kubeconfig)
                KUBECONFIG_FILE="$2"
                shift 2
                ;;
            --uid)
                CLUSTER_UID="$2"
                shift 2
                ;;
            --is-helm)
                IS_HELM=true
                shift
                ;;
            --help)
                show_help
                exit 0
                ;;
            *)
                log_error "Unknown option: $1"
                show_help
                exit 1
                ;;
        esac
    done
    if [[ -n "$CLUSTER_UID" ]]; then
        CLUSTER_NAMESPACE="cluster-${CLUSTER_UID}"
    fi
}

# Validate parameters
validate_parameters() {
    if [[ -z "$KUBECONFIG_FILE" ]]; then
        log_error "kubeconfig file path is required"
        show_help
        exit 1
    fi

    if [[ ! -f "$KUBECONFIG_FILE" ]]; then
        log_error "kubeconfig file does not exist: $KUBECONFIG_FILE"
        exit 1
    fi

    if [[ "$IS_HELM" != "true" && -z "$CLUSTER_UID" ]]; then
        log_error "Cluster UID is required (--uid UID)"
        show_help
        exit 1
    fi
}

# Check if kubectl is available
check_kubectl() {
    if ! command -v kubectl &> /dev/null; then
        log_error "kubectl command not found. Please install kubectl."
        exit 1
    fi
}

# Test cluster connectivity
test_cluster_connection() {
    log_info "Testing cluster connectivity..."
    if ! kubectl --kubeconfig="$KUBECONFIG_FILE" cluster-info &>/dev/null; then
        log_error "Failed to connect to cluster. Please check your kubeconfig and cluster status."
        exit 1
    fi
    log_success "Cluster connection successful"
}

# Scale deployment to N replicas
scale_deployment() {
    local namespace="$1"
    local deployment="$2"
    local replicas="$3"

    log_info "Scaling deployment $deployment in namespace $namespace to $replicas replica(s)..."

    if ! kubectl --kubeconfig="$KUBECONFIG_FILE" scale deployment "$deployment" --replicas="$replicas" -n "$namespace" &>/dev/null; then
        log_error "Failed to scale deployment $deployment in namespace $namespace to $replicas"
        exit 1
    fi

    log_success "Deployment $deployment scaled to $replicas replica(s)"
}

# Scale statefulset to N replicas
scale_statefulset() {
    local namespace="$1"
    local name="$2"
    local replicas="$3"

    log_info "Scaling statefulset $name in namespace $namespace to $replicas replica(s)..."

    if ! kubectl --kubeconfig="$KUBECONFIG_FILE" scale statefulset "$name" --replicas="$replicas" -n "$namespace" &>/dev/null; then
        log_error "Failed to scale statefulset $name in namespace $namespace to $replicas"
        exit 1
    fi

    log_success "Statefulset $name scaled to $replicas replica(s)"
}

# Scale msgbroker statefulset in hubble-system to 2 replicas and wait for readiness
scale_msgbroker_to_2() {
    local desired_replicas=2

    scale_statefulset "hubble-system" "msgbroker" "$desired_replicas"

    log_info "Waiting for msgbroker pods to become ready..."
    local timeout=120  # 2 minutes timeout
    local count=0

    while [[ $count -lt $timeout ]]; do
        local ready_pods
        ready_pods=$(kubectl --kubeconfig="$KUBECONFIG_FILE" get pods -n "hubble-system" -l component=msgbroker --no-headers 2>/dev/null | awk '$3=="Running" {split($2,a,"/"); if(a[1]==a[2] && a[1]>0) c++} END {print c+0}' || echo "0")

        if [[ "$ready_pods" -ge "$desired_replicas" ]]; then
            log_success "All msgbroker pods are ready ($ready_pods/$desired_replicas)"
            return 0
        fi

        log_info "Waiting for msgbroker pods to become ready ($ready_pods/$desired_replicas ready, ${count}s elapsed)..."
        sleep 5
        count=$((count + 5))
    done

    log_error "Timeout waiting for msgbroker pods to become ready"
    exit 1
}

# Scale configserver deployment in hubble-system to 1 replica and wait for Running state
scale_configserver_to_1() {
    local desired_replicas=1

    log_info "Scaling deployment configserver in hubble-system to $desired_replicas replica(s)..."
    scale_deployment "hubble-system" "configserver" "$desired_replicas"

    log_info "Waiting for configserver pods to reach Running state..."
    local timeout=120  # 2 minutes timeout
    local count=0

    while [[ $count -lt $timeout ]]; do
        local running_pods
        running_pods=$(kubectl --kubeconfig="$KUBECONFIG_FILE" get pods -n "hubble-system" -l component=configserver --no-headers 2>/dev/null | awk '$3=="Running" {c++} END {print c+0}' || echo "0")

        if [[ "$running_pods" -ge "$desired_replicas" ]]; then
            log_success "All configserver pods are Running ($running_pods/$desired_replicas)"
            return 0
        fi

        log_info "Waiting for configserver pods to reach Running state ($running_pods/$desired_replicas running, ${count}s elapsed)..."
        sleep 5
        count=$((count + 5))
    done

    log_error "Timeout waiting for configserver pods to reach Running state"
    exit 1
}

# Scale auth deployment in hubble-system to 2 replicas and wait for readiness
scale_auth_to_2() {
    local desired_replicas=2

    log_info "Scaling deployment auth in hubble-system to $desired_replicas replica(s)..."
    scale_deployment "hubble-system" "auth" "$desired_replicas"

    log_info "Waiting for auth pods to become ready..."
    local timeout=120  # 2 minutes timeout
    local count=0

    while [[ $count -lt $timeout ]]; do
        local ready_pods
        ready_pods=$(kubectl --kubeconfig="$KUBECONFIG_FILE" get pods -n "hubble-system" -l component=auth --no-headers 2>/dev/null | awk '$3=="Running" {split($2,a,"/"); if(a[1]==a[2] && a[1]>0) c++} END {print c+0}' || echo "0")

        if [[ "$ready_pods" -ge "$desired_replicas" ]]; then
            log_success "All auth pods are ready ($ready_pods/$desired_replicas)"
            return 0
        fi

        log_info "Waiting for auth pods to become ready ($ready_pods/$desired_replicas ready, ${count}s elapsed)..."
        sleep 5
        count=$((count + 5))
    done

    log_error "Timeout waiting for auth pods to become ready"
    exit 1
}

# Scale system deployment in hubble-system to 2 replicas and wait for readiness
scale_system_to_2() {
    local desired_replicas=2

    log_info "Scaling deployment system in hubble-system to $desired_replicas replica(s)..."
    scale_deployment "hubble-system" "system" "$desired_replicas"

    log_info "Waiting for system pods to become ready..."
    local timeout=120  # 2 minutes timeout
    local count=0

    while [[ $count -lt $timeout ]]; do
        local ready_pods
        ready_pods=$(kubectl --kubeconfig="$KUBECONFIG_FILE" get pods -n "hubble-system" -l component=system --no-headers 2>/dev/null | awk '$3=="Running" {split($2,a,"/"); if(a[1]==a[2] && a[1]>0) c++} END {print c+0}' || echo "0")

        if [[ "$ready_pods" -ge "$desired_replicas" ]]; then
            log_success "All system pods are ready ($ready_pods/$desired_replicas)"
            return 0
        fi

        log_info "Waiting for system pods to become ready ($ready_pods/$desired_replicas ready, ${count}s elapsed)..."
        sleep 5
        count=$((count + 5))
    done


    log_error "Timeout waiting for system pods to become ready"
    exit 1
}

# Step 1: Scale all hubble-system deployments to 1
scale_hubble_system_to_1() {
    log_info "=== Scaling all deployments in hubble-system to 1 replica ==="
    for deployment in "${HUBBLE_SYSTEM_TO_1[@]}"; do
        scale_deployment "hubble-system" "$deployment" 1
        sleep 10
    done

    log_success "All hubble-system deployments scaled to 1 replica"
}

# Step 2: Scale selected hubble-system deployments to 2
scale_hubble_system_to_2() {
    log_info "=== Scaling cloud, hashboard, hutil, spectroclusterop, user to 2 replicas ==="

    for deployment in "${HUBBLE_SYSTEM_TO_2[@]}"; do
        scale_deployment "hubble-system" "$deployment" 2
        sleep 10
    done

    log_success "Deployments scaled to 2 replicas"
}

# Step 3: Scale spectrocluster, timeseries, event to 3
scale_hubble_system_to_3() {
    log_info "=== Scaling spectrocluster, timeseries, event to 3 replicas ==="

    for deployment in "${HUBBLE_SYSTEM_TO_3[@]}"; do
        scale_deployment "hubble-system" "$deployment" 3
        sleep 10
    done

    log_success "Deployments scaled to 3 replicas"
}

# Step 4: Scale cluster-<uid> deployments to 1 replica
scale_cluster_uid_deployments() {
    log_info "=== Scaling deployments in ${CLUSTER_NAMESPACE} to 1 replica ==="

    for deployment in "${CLUSTER_UID_DEPLOYMENTS[@]}"; do
        scale_deployment "$CLUSTER_NAMESPACE" "$deployment" 1
        sleep 10
    done

    log_success "Deployments in ${CLUSTER_NAMESPACE} scaled to 1 replica"
}

# Main function
main() {
    parse_arguments "$@"
    validate_parameters
    check_kubectl
    test_cluster_connection

    log_info "Starting Kubernetes startup process (edge)..."

    # hubble-system: msgbroker to 2, configserver to 1, all deployments to 1, then 2-replica set, then 3-replica set
    scale_msgbroker_to_2
    scale_configserver_to_1
    scale_system_to_2
    scale_auth_to_2
    scale_hubble_system_to_1
    scale_hubble_system_to_2
    scale_hubble_system_to_3

    # cluster-<uid>: both deployments to 1 replica (skipped for helm)
    if [[ "$IS_HELM" == "true" ]]; then
        log_info "Skipping cluster-<uid> deployment scale-up (--is-helm)"
    else
        scale_cluster_uid_deployments
    fi

    log_success "Kubernetes startup process (edge) completed successfully!"
}

# Run main function with all arguments
main "$@"