#!/bin/bash

# Kubernetes Shutdown Script
# This script scales down deployments and suspends cronjobs in a specific order

set -euo pipefail

# Default configuration
KUBECONFIG_FILE=""
CLUSTER_UID=""
IS_HELM=false
CLUSTER_NAMESPACE=""  # Set to cluster-<uid> after parsing
SPECTRO_MGMT_DEPLOYMENTS=("cluster-management-agent-lite" "palette-lite-controller-manager")
HUBBLE_SYSTEM_DEPLOYMENTS=("spectrocluster" "spectrocluster-jobs" "timeseries" "cloud" "event" "hashboard" "hutil" "auth" "memstore" "system" "user" "configserver" "foreq" "mgmt" "oci-proxy" "spectro-tunnel" "spectrocluster-reconciler" "spectroclusterop" "spectrossh" "system")
HUBBLE_SYSTEM_CRONJOBS=()

# 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 Shutdown Script (Edge)

This script performs the following operations in order:
1. Scale down deployments in cluster-<uid> namespace
2. Scale down deployments in hubble-system namespace
3. Suspend cronjobs in hubble-system namespace
4. Scale down msg-broker statefulset in hubble-system namespace (last step)

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-down.

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 0 replicas and wait for pods to stop
scale_deployment() {
    local namespace="$1"
    local deployment="$2"
    
    log_info "Scaling deployment $deployment in namespace $namespace to 0 replicas..."
    
    # Scale deployment to 0
    if ! kubectl --kubeconfig="$KUBECONFIG_FILE" scale deployment "$deployment" --replicas=0 -n "$namespace" &>/dev/null; then
        log_error "Failed to scale deployment $deployment in namespace $namespace"
        exit 1
    fi
    
    log_success "Deployment $deployment scaled to 0 replicas"
    
    # Wait for pods to stop
    log_info "Waiting for pods to stop..."
    local timeout=300  # 5 minutes timeout
    local count=0
    
    while [[ $count -lt $timeout ]]; do
        local running_pods=$(kubectl --kubeconfig="$KUBECONFIG_FILE" get pods -n "$namespace" -l app="$deployment" --no-headers 2>/dev/null | grep -v "Completed\|Succeeded" | wc -l | tr -d ' \n' || echo "0")
        
        if [[ "$running_pods" -eq 0 ]]; then
            log_success "All pods for deployment $deployment have stopped"
            return 0
        fi
        
        log_info "Waiting for $running_pods pods to stop... (${count}s elapsed)"
        sleep 5
        count=$((count + 5))
    done
    
    log_warning "Timeout waiting for pods to stop for deployment $deployment"
    return 1
}

# Suspend cronjob
suspend_cronjob() {
    local namespace="$1"
    local cronjob="$2"
    
    log_info "Suspending cronjob $cronjob in namespace $namespace..."
    
    if ! kubectl --kubeconfig="$KUBECONFIG_FILE" patch cronjob "$cronjob" -n "$namespace" -p '{"spec":{"suspend":true}}' &>/dev/null; then
        log_error "Failed to suspend cronjob $cronjob in namespace $namespace"
        exit 1
    fi
    
    log_success "Cronjob $cronjob suspended"
}

# Scale down deployments in cluster-<uid> namespace
scale_spectro_mgmt_deployments() {
    log_info "=== Scaling down deployments in ${CLUSTER_NAMESPACE} namespace ==="
    
    for deployment in "${SPECTRO_MGMT_DEPLOYMENTS[@]}"; do
        scale_deployment "$CLUSTER_NAMESPACE" "$deployment"
    done
    
    log_success "All deployments in ${CLUSTER_NAMESPACE} namespace scaled down"
}

# Scale down deployments in hubble-system namespace
scale_hubble_system_deployments() {
    log_info "=== Scaling down deployments in hubble-system namespace ==="
    
    for deployment in "${HUBBLE_SYSTEM_DEPLOYMENTS[@]}"; do
        scale_deployment "hubble-system" "$deployment"
    done
    
    log_success "All deployments in hubble-system namespace scaled down"
}

# Suspend cronjobs in hubble-system namespace
suspend_hubble_system_cronjobs() {
    log_info "=== Suspending cronjobs in hubble-system namespace ==="
    
    if [[ ${#HUBBLE_SYSTEM_CRONJOBS[@]} -eq 0 ]]; then
        log_info "No cronjobs to suspend in hubble-system namespace"
        return 0
    fi
    
    for cronjob in "${HUBBLE_SYSTEM_CRONJOBS[@]}"; do
        suspend_cronjob "hubble-system" "$cronjob"
    done
    
    log_success "All cronjobs in hubble-system namespace suspended"
}

# Scale down msg-broker statefulset
scale_msg_broker_statefulset() {
    log_info "=== Scaling down msg-broker statefulset in hubble-system namespace ==="
    
    log_info "Scaling statefulset msgbroker in namespace hubble-system to 0 replicas..."
    
    # Scale statefulset to 0
    if ! kubectl --kubeconfig="$KUBECONFIG_FILE" scale statefulset "msgbroker" --replicas=0 -n "hubble-system" &>/dev/null; then
        log_error "Failed to scale statefulset msgbroker in namespace hubble-system"
        exit 1
    fi
    
    log_success "Statefulset msgbroker scaled to 0 replicas"
    
    # Wait for pods to stop
    log_info "Waiting for msg-broker pods to stop..."
    local timeout=300  # 5 minutes timeout
    local count=0
    
    while [[ $count -lt $timeout ]]; do
        local running_pods=$(kubectl --kubeconfig="$KUBECONFIG_FILE" get pods -n "hubble-system" -l app="msg-broker" --no-headers 2>/dev/null | grep -v "Completed\|Succeeded" | wc -l | tr -d ' \n' || echo "0")
        
        if [[ "$running_pods" -eq 0 ]]; then
            log_success "All msg-broker pods have stopped"
            return 0
        fi
        
        log_info "Waiting for $running_pods msg-broker pods to stop... (${count}s elapsed)"
        sleep 5
        count=$((count + 5))
    done
    
    log_warning "Timeout waiting for msg-broker pods to stop"
    return 1
}

# Main function
main() {
    parse_arguments "$@"
    validate_parameters
    check_kubectl
    test_cluster_connection
    
    log_info "Starting Kubernetes shutdown process..."
    
    # Step 1: Scale down cluster-<uid> deployments (skipped for helm)
    if [[ "$IS_HELM" == "true" ]]; then
        log_info "Skipping cluster-<uid> deployment scale-down (--is-helm)"
    else
        scale_spectro_mgmt_deployments
    fi
    
    # Step 2: Scale down hubble-system deployments
    scale_hubble_system_deployments
    
    # Step 3: Suspend hubble-system cronjobs
    suspend_hubble_system_cronjobs
    
    # Step 4: Scale down msg-broker statefulset (last step)
    scale_msg_broker_statefulset
    
    log_success "Kubernetes shutdown process completed successfully!"
}

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