diff --git a/snap-upgrade.sh b/snap-upgrade.sh
index 82d7fbd..385691f 100644
--- a/snap-upgrade.sh
+++ b/snap-upgrade.sh
@@ -1,39 +1,71 @@
 #!/bin/bash
 
-SNAP_CHANNEL="1.21/stable"
+# 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"
 
-# Download Juju 2.9
-snap download --channel=2.9 juju
+# 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
 
-# Download latest snaps from designated channel
-for snap in $ALL_SNAPS
-do
-  snap download --channel=$SNAP_CHANNEL $snap
-done
+# 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
+}
 
-# Attach new snaps to master units
-for snap in $MASTER_SNAPS
-do
-  juju attach kubernetes-master $snap=`ls ${snap}_*.snap`
-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
+}
 
-# Attach new snaps to worker units
-for snap in $WORKER_SNAPS
-do
-  juju attach kubernetes-worker $snap=`ls ${snap}_*.snap`
-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
+}
 
-# Upgrade to new snaps on masters, one at a time
-for unit in `juju status --format json | jq -r '.applications|.["kubernetes-master"].units | keys[]'`
-do
-  juju run-action $unit upgrade --wait
-done
+# Download Juju
+echo "Downloading Juju from channel $JUJU_CHANNEL..."
+snap download --channel="$JUJU_CHANNEL" juju || exit 1
 
-# Upgrade to new snaps on workers, one at a time
-for unit in `juju status --format json | jq -r '.applications|.["kubernetes-worker"].units | keys[]'`
-do
-  juju run-action $unit upgrade --wait
-done
\ No newline at end of file
+# 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."
\ No newline at end of file