#!/bin/bash

# Constants
ALL_SNAPS="kube-apiserver kube-scheduler kube-controller-manager kube-proxy kubectl kubelet cdk-addons"
MASTER_SNAPS="kube-apiserver kube-scheduler kube-controller-manager kube-proxy kubectl cdk-addons"
WORKER_SNAPS="kube-proxy kubelet kubectl"

# Prompt for channels
read -p "Enter the SNAP_CHANNEL (e.g. 1.21/stable): " SNAP_CHANNEL
read -p "Enter the JUJU_CHANNEL (e.g. 2.9/stable): " JUJU_CHANNEL

# Function to download snaps from a specific channel
download_snaps() {
  local snaps=$1
  local channel=$2
  for snap in $snaps; do
    echo "Downloading $snap from channel $channel..."
    snap download --channel="$channel" "$snap" || exit 1
  done
}

# Function to attach snaps to units
attach_snaps_to_units() {
  local snaps=$1
  local service=$2
  for snap in $snaps; do
    local snap_file
    snap_file=$(ls "${snap}"_*.snap | head -n 1)
    if [[ -f "$snap_file" ]]; then
      echo "Attaching $snap_file to $service..."
      # juju attach "$service" "$snap"="$snap_file" || exit 1
    else
      echo "Error: $snap_file not found."
      exit 1
    fi
  done
}

# Function to upgrade units of a service
upgrade_service_units() {
  local service=$1
  local unit_names
  unit_names=$(juju status --format json | jq -r ".applications|.[\"$service\"].units | keys[]")
  for unit in $unit_names; do
    echo "Upgrading $unit..."
    # juju run-action "$unit" upgrade --wait || exit 1
  done
}

# Download Juju
echo "Downloading Juju from channel $JUJU_CHANNEL..."
snap download --channel="$JUJU_CHANNEL" juju || exit 1

# Download Kubernetes snaps
download_snaps "$ALL_SNAPS" "$SNAP_CHANNEL"

echo "Downloading Juju from channel $JUJU_CHANNEL..."
snap download --channel="$JUJU_CHANNEL" juju || exit 1

# Download Kubernetes snaps
download_snaps "$ALL_SNAPS" "$SNAP_CHANNEL"

# Attach new snaps to Kubernetes master and worker units
attach_snaps_to_units "$MASTER_SNAPS" "kubernetes-master"
attach_snaps_to_units "$WORKER_SNAPS" "kubernetes-worker"

# Upgrade Kubernetes master and worker units
upgrade_service_units "kubernetes-master"
upgrade_service_units "kubernetes-worker"

echo "Upgrade process completed successfully."