#!/usr/bin/env bash
set -e
if [[ $EUID -ne 0 ]]; then
   echo "This script must be run as root"
   exit 1
fi
SCRIPT_PATH="$( cd "$(dirname "$0")" ; pwd -P )"
ARCH=$(uname -m)
LEDGER_FILE_LOC=/opt/kasm/current/plugins
YQ="$(dirname $0)/yq_${ARCH}"


# Get formatted timestamp for ledger
function get_ledger_timestamp() {
  echo "$(TZ=UTC date -Iseconds | awk -F "+" '{print $1}')Z"
}

# If there is a newer stable version than is installed, returns that version tag.
function get_update_string() {
  local latest_stable

  latest_stable=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --latest)
  if [[ -z ${latest_stable} ]]; then
    >&2 echo "* ERROR: Unable to determine latest plugin version. Current plugin will not be modified."
    return 0
  fi

  local installed_stable
  local should_update
  local reason
  if [[ ! -f ${LEDGER_FILE} ]]; then
    >&2 echo "* No ledger file for plugin"
    should_update=1
  else
    installed_stable=$(tail -n 1 "${LEDGER_FILE}" | awk -F',' '{print $2}')
    if [[ -z ${installed_stable} ]]; then
      >&2 echo "* WARNING: Current ledger file for plugin is corrupt"
      mv "${LEDGER_FILE}" "${LEDGER_FILE}.$(get_ledger_timestamp)" || :
      should_update=1
    elif [[ "${installed_stable}" != "${latest_stable}" ]]; then
      if [[ $(printf "${installed_stable}\n${latest_stable}" | sort -V | tail -n 1) == "${latest_stable}" ]]; then
        should_update=1
      else
        reason="version is ahead of registry"
      fi
    else
      reason="currently at the latest version"
    fi
  fi

  if [[ -z ${should_update} ]]; then
      >&2 echo "* Will not update rolling plugin: ${reason}"
      return 0
  fi

  echo "${latest_stable}"
}

function update_ledger() {
  local installed_digest="$1"
  >&2 echo "* Recording new plugin version in ledger"
  if [[ ! -e "${LEDGER_FILE}" ]]; then
    touch "${LEDGER_FILE}"
    chown "${KASM_UID:-$(id -u kasm)}:${KASM_GID:-$(id -g kasm)}" "${LEDGER_FILE}"
  fi
  echo "$(get_ledger_timestamp),${NEWER_PLUGIN_VERSION},${installed_digest}" >> "${LEDGER_FILE}"
}

function run_upgrade() {
  set +e
  local output
  output=$(docker plugin disable "${PLUGIN_NAME}" 2>&1)
  if [[ $? -ne 0 ]] && [[ ! "${output}" =~ "plugin is already disabled" ]]; then
    >&2 echo "* ERROR: Unable to disable plugin, skipping upgrade"
    >&2 echo "${output}"
    set -e
  else
    set -e
    >&2 echo "* Running plugin upgrade"
    output=$(docker plugin upgrade "${PLUGIN_NAME}" --grant-all-permissions)
    if [[ $? -eq 0 ]]; then
      >&2 echo "${output}"
      local digest
      if [[ "${output}" =~ Digest:[[:space:]]+(sha256:[[:xdigit:]]+) ]]; then
        digest="${BASH_REMATCH[1]}"
      fi
      update_ledger "$digest"
    else
      >&2 echo "* ${output}"
      >&2 echo "* ERROR: Upgrade failed"
    fi
    docker plugin enable "${PLUGIN_NAME}" > /dev/null
  fi
}

function network_plugin_prologue() {
  local containers
  local container_name
  local container_state
  local net_name
  local net_id
  local net_driver

  # Identify all networks that use this plugin as their driver
  while read -r NETWORK_LINE; do
    net_id=$(echo "${NETWORK_LINE}" | awk -F',' '{ print $1 }' )
    net_name=$(echo "${NETWORK_LINE}" | awk -F',' '{ print $2 }' )
    net_driver=$(echo "${NETWORK_LINE}" | awk -F',' '{ print $3 }' )
    if [[ "${net_driver}" != "${PLUGIN_NAME}" ]]; then
      continue
    fi

    # ensure there are no containers using the network
    containers=$(docker ps -a --filter "network=${net_id}" --format "{{.Names}},{{.State}}")
    for container in ${containers}; do
      container_name=$(echo "${container}" | awk -F, '{print $1}')
      container_state=$(echo "${container}" | awk -F, '{print $2}')
      if [[ "${container_state}" == "running" ]] || [[ "${container_state}" == "restarting" ]] || [[ "${container_state}" == "paused" ]]; then
        >&2 echo "* ERROR: Network plugin has one or more active containers, please stop local Kasm services and Kasms before upgrading"
        return 1
      fi
      >&2 echo "* Disconnecting dependent container ${container_name}"
      docker network disconnect "${net_name}" "${container_name}" > /dev/null
      NETWORK_CONTAINERS+=("${container_name},${net_name}")
    done

    >&2 echo "* Removing dependent docker network ${net_name}"
    docker network rm "${net_name}" > /dev/null
    NETWORKS+=("${net_name}")
    # Remove old api_hostname file if it was not provided at install time
    if [[ "${net_name}" == kasm_sidecar_network ]] && [[ ! -e /var/run/kasm-sidecar/api_hostname_external ]]; then
      rm -f /var/run/kasm-sidecar/api_hostname
    fi
  done < <(docker network ls --format '{{.ID}},{{.Name}},{{.Driver}}')
}

function network_plugin_epilogue() {
  local container_name
  local net_name

  for net_name in "${NETWORKS[@]}"; do
    >&2 echo "* Recreating docker network ${net_name}"
    docker network create --driver "${PLUGIN_NAME}" ${net_name} > /dev/null
  done

  # Register the kasm_proxy first. It needs to be present to configure other
  # containers on the sidecar network.
  if [[ "${NETWORK_CONTAINERS[*]}" =~ kasm_proxy ]]; then
    >&2 echo "* Reconnecting kasm_proxy"
    docker network connect "${net_name}" kasm_proxy
  fi
  for container in "${NETWORK_CONTAINERS[@]}"; do
    container_name=$(echo "${container}" | awk -F, '{print $1}')
    net_name=$(echo "${container}" | awk -F, '{print $2}')
    if [[ "${container_name}" == kasm_proxy ]]; then
      continue
    fi
    >&2 echo "* Reconnecting dependent container ${container_name}"
    docker network connect "${net_name}" "${container_name}"
  done
}

function volume_plugin_prologue() {
  local containers

  while read -r VOLUME_LINE; do
    volume_name=$(echo "${VOLUME_LINE}" | awk -F, '{print $1}')
    volume_driver=$(echo "${VOLUME_LINE}" | awk -F, '{print $2}')

    if [[ "${volume_driver}" != "${PLUGIN_NAME}" ]]; then
      continue
    fi
    # check if there are containers using the volume
    containers=$(docker ps -a --filter "volume=${volume_name}" --format "{{.Names}},{{.State}}")
    if [[ -n ${containers} ]]; then
      >&2 echo "* ERROR: Volume plugin has one or more active containers, please stop local Kasm services and Kasms before upgrading"
        return 1
    fi
    >&2 echo "* ERROR: Volume plugin has one or more volumes, please remove them to enable upgrade"
    return 1
  done < <(docker volume ls --format '{{.Name}},{{.Driver}}')
}

function set_plugin_types() {
  IS_NETWORK_PLUGIN=
  IS_VOLUME_PLUGIN=

  while read -r plugin_type; do
    case "${plugin_type}" in
      "docker.networkdriver/1.0")
        >&2 echo "* Type: network plugin"
        IS_NETWORK_PLUGIN=1
        ;;
      "docker.volumedriver/1.0")
        >&2 echo "* Type: volume plugin"
        IS_VOLUME_PLUGIN=1
        ;;
      "docker.logdriver/1.0")
        >&2 echo "* Type: log plugin"
        # No precautions are necessary in the management of logging plugins.
        # Logging plugins can't have active containers when disabling.
        # Attempting to do so while containers are using it will result in failure, which we report.
        ;;
      *)
        >&2 "* ERROR: Plugin type ${plugin_type} is not supported for upgrade"
        return 0
        ;;
    esac
  done < <(docker plugin inspect "${PLUGIN_NAME}" --format '{{json .Config.Interface.Types}}' | ${YQ} -r '.[]')
}

function upgrade_plugin() {
  if [[ -z ${PLUGIN_NAME} ]]; then
    >&2 echo "* ERROR: PLUGIN_NAME is not set"
    return 1
  fi
  if [[ -z ${PLUGIN_ID} ]]; then
    >&2 echo "* ERROR: PLUGIN_ID is not set"
    return 1
  fi
  if [[ -z ${PLUGIN_REF} ]]; then
    >&2 echo "* ERROR: PLUGIN_REF is not set"
    return 1
  fi
  if [[ -z ${LEDGER_FILE} ]]; then
    >&2 echo "* ERROR: LEDGER_FILE is not set"
    return 1
  fi

  NETWORKS=()
  NETWORK_CONTAINERS=()

  set_plugin_types

  local ref_is_release
  ref_is_release=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --is-release)
  if [[ ${ref_is_release} == false ]]; then
    >&2 echo "* Updating plugin with unversioned release"
    NEWER_PLUGIN_VERSION=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --tag)
  elif [[ -n ${FORCE} ]]; then
    NEWER_PLUGIN_VERSION=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --latest)
    if [[ -z ${NEWER_PLUGIN_VERSION} ]]; then
      >&2 echo "* Unable to find patch-level tag for plugin ledger, will use ${NEWER_PLUGIN_VERSION}"
      NEWER_PLUGIN_VERSION=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --tag)
    fi
  else
    NEWER_PLUGIN_VERSION=$(get_update_string)
    if [[ -z ${NEWER_PLUGIN_VERSION} ]]; then
      return 0
    fi
  fi

  local run_upgrade=1
  if [[ -n ${IS_NETWORK_PLUGIN} ]]; then
    network_plugin_prologue
    if [[ $? -ne 0 ]]; then
      unset run_upgrade
    fi
  fi
  if [[ -n ${IS_VOLUME_PLUGIN} ]]; then
    volume_plugin_prologue
    if [[ $? -ne 0 ]]; then
      unset run_upgrade
    fi
  fi

  if [[ -n ${run_upgrade} ]]; then
    run_upgrade
  fi

  if [[ -n ${IS_NETWORK_PLUGIN} ]]; then
    network_plugin_epilogue
  fi
}

function reinstall_plugin() {
  local registry_host
  local org_name
  local repo
  local new_ref
  local output
  local digest

  PLUGIN_REF=$(sudo docker plugin inspect "${PLUGIN_NAME}" --format "{{.PluginReference}}" | sed -r 's#docker\.io/##')
  registry_host=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --host)
  org_name=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --org)
  repo=$(bash ${SCRIPT_PATH}/plugin_helper ${PLUGIN_REF} --repo)
  new_ref="${org_name}/${repo}:${PLUGIN_VERSION}"
  if [[ -n ${registry_host} ]] && [[ ${registry_host} != registry\.hub\.docker\.com ]]; then
    new_ref="${registry_host}/${new_ref}"
  fi

  set_plugin_types

  >&2 echo -e "* Will install\t${new_ref}"
  local run_install=1
  if [[ $(bash "${SCRIPT_PATH}/plugin_helper" ${new_ref} --exists) != true ]]; then
    >&2 echo "* ERROR: Specified plugin/version does not exist"
    return 1
  fi

  if [[ -n ${IS_NETWORK_PLUGIN} ]]; then
    network_plugin_prologue
    if [[ $? -ne 0 ]]; then
      unset run_install
    fi
  fi
  if [[ -n ${IS_VOLUME_PLUGIN} ]]; then
    volume_plugin_prologue
    if [[ $? -ne 0 ]]; then
      unset run_install
    fi
  fi

  if [[ -n ${run_install} ]]; then
    >&2 echo "* Disabling and removing old plugin"
    docker plugin disable "${PLUGIN_NAME}"
    docker plugin rm "${PLUGIN_NAME}"
    >&2 echo "* Installing and enabling new plugin"
    if [[ -z ${PLUGIN_ALIAS} ]]; then
      output=$(docker plugin install --grant-all-permissions ${new_ref})
      PLUGIN_NAME="${new_ref}"
    else
      output=$(docker plugin install --grant-all-permissions --alias "${PLUGIN_ALIAS}" ${new_ref})
      PLUGIN_NAME="${PLUGIN_ALIAS}"
    fi
    if [[ $? -ne 0 ]]; then
      >&2 echo "${output}"
      >&2 echo "* ERROR: Failed to install new plugin version"
      return 1
    fi
    >&2 echo "${output}"
    >&2 echo "* Note: Plugin name may differ from previous installed version if no --alias was provided"
    if [[ "${output}" =~ Digest:[[:space:]]+(sha256:[[:xdigit:]]+) ]]; then
      digest="${BASH_REMATCH[1]}"
    fi
    if [[ "${PLUGIN_VERSION}" =~ ^(amd64|arm64)-[0-9]+\.[0-9]+\.[0-9]{14}$ ]]; then
      NEWER_PLUGIN_VERSION="${PLUGIN_VERSION}"
    else
      NEWER_PLUGIN_VERSION="$(${SCRIPT_PATH}/plugin_helper ${new_ref} --latest)"
    fi

    if [[ -z ${NEWER_PLUGIN_VERSION} ]]; then
      >&2 echo "* WARNING: Unable to determine version string for ledger entry"
    fi
    LEDGER_FILE="${LEDGER_FILE_LOC}/$(echo "${PLUGIN_NAME}" | sed -r 's#[./:-]#_#g').csv"
    update_ledger "$digest"
  fi

  if [[ -n ${IS_NETWORK_PLUGIN} ]]; then
    network_plugin_epilogue
  fi
}

function display_help() {
  echo -e "Usage:\t$0"
  echo -e "\t$0 --name PLUGIN_NAME"
  echo -e "\t$0 --rollback --name PLUGIN_NAME [--alias PLUGIN_ALIAS]"
  echo -e "\t$0 --reinstall --name PLUGIN_NAME --version PLUGIN_TAG [--alias PLUGIN_ALIAS]"
  echo
  echo -e "Without options: Upgrades each installed Kasm plugin with a rolling tag to its latest patch-level version."
  echo -e "With only --name: Upgrades the specified plugin."
  echo -e "With options: Rolls back / reinstalls the specified plugin. Rollback mode targets the most recent previously installed version, and only if that information is available."
  echo
  echo -e "--name\t\tName of the plugin, as seen in the output of \`docker plugin ls\`"
  echo -e "--version\tDesired plugin tag to install (--reinstall mode only)."
  echo -e "--alias\t\tName that the installed plugin will use. Defaults to the plugin registry reference value, ex: kasmweb/kasm-network-plugin:amd64-0.1.20250607080910"
}

# Main
[[ ! -e ${YQ} ]] && {
  >&2 echo "ERROR: Dependency yq_${ARCH} not found"
  exit 1
}

if [[ ! -e "${LEDGER_FILE_LOC}" ]]; then
  mkdir -p "${LEDGER_FILE_LOC}"
  chown "${KASM_UID:-$(id -u kasm)}:${KASM_GID:-$(id -g kasm)}" "${LEDGER_FILE_LOC}"
fi

MODE=UPGRADE

ARGS=("$@")
for index in "${!ARGS[@]}"; do
  case ${ARGS[index]} in
    -h|--help)
        display_help
        exit 0
        ;;
    -i|--reinstall)
        MODE=REINSTALL
        ;;
    -r|--rollback)
        MODE=ROLLBACK
        ;;
    -n|--name)
        PLUGIN_NAME="${ARGS[index+1]}"
        if [[ -z ${PLUGIN_NAME} ]]; then
          display_help
          exit 1
        fi
        ;;
    --version)
        PLUGIN_VERSION="${ARGS[index+1]}"
        if [[ -z ${PLUGIN_VERSION} ]]; then
          display_help
          exit 1
        fi
        ;;
    --alias)
        PLUGIN_ALIAS="${ARGS[index+1]}"
        ;;
    -*|--*)
        echo "Unknown option ${ARGS[index]}"
        display_help
      	exit 1
        ;;
  esac
done

if [[ "${MODE}" == ROLLBACK ]]; then
  >&2 echo -e "Rolling back plugin\t${PLUGIN_NAME}"
  LEDGER_FILE="${LEDGER_FILE_LOC}/$(echo "${PLUGIN_NAME}" | sed -r 's#[./:-]#_#g').csv"
  >&2 echo "* Pulling previous version from ledger file"
  if [[ $(wc -l "${LEDGER_FILE}" | awk '{ print $1 }') -lt 2 ]]; then
    >&2 echo "* ERROR: No previous version recorded"
    exit 1
  fi
  PLUGIN_VERSION=$(tail -n 2 "${LEDGER_FILE}" | sed '1q;d' | awk -F, '{ print $2}')
  >&2 echo "* Using version ${PLUGIN_VERSION}"
  reinstall_plugin
elif [[ "${MODE}" == REINSTALL ]]; then
  if [[ -z ${PLUGIN_VERSION} ]]; then
    display_help
    exit 1
  fi
  >&2 echo "* Targeting version ${PLUGIN_VERSION}"
  reinstall_plugin
else
  if [[ -n "${PLUGIN_NAME}" ]]; then
    >&2 echo "Upgrading plugin: ${PLUGIN_NAME}"
    _inspect_output=$(docker plugin inspect "${PLUGIN_NAME}" --format '{{.ID}} {{.PluginReference}}')
    FORCE=1
    PLUGIN_ID=$(echo "${_inspect_output}" | awk '{ print $1 }')
    PLUGIN_REF=$(echo "${_inspect_output}" | awk '{ print $2 }')
    LEDGER_FILE="${LEDGER_FILE_LOC}/$(echo "${PLUGIN_NAME}" | sed -r 's#[./:-]#_#g').csv"
    upgrade_plugin || :
  else
    >&2 echo "Upgrading rolling plugins"
    docker plugin ls --format '{{.Name}} {{.ID}} {{.PluginReference}}' | while read -r PLUGIN_LINE; do
      PLUGIN_NAME=$(echo "${PLUGIN_LINE}" | awk '{ print $1 }')
      PLUGIN_ID=$(echo "${PLUGIN_LINE}" | awk '{ print $2 }')
      PLUGIN_REF=$(echo "${PLUGIN_LINE}" | awk '{ print $3 }')
      LEDGER_FILE="${LEDGER_FILE_LOC}/$(echo "${PLUGIN_NAME}" | sed -r 's#[./:-]#_#g').csv"
      if [[ "${PLUGIN_NAME}" =~ kasmweb/.+:.+-rolling(-.*)? ]]; then
        >&2 echo -e "Inspecting plugin:\t${PLUGIN_NAME}"
        upgrade_plugin || :
      fi
    done
  fi
fi