#!/bin/bash

function raise_err() {
  local msg="$1"
  echo -e "$msg"
  exit 1
}

function get_latest_tag() {
  local owner=$1
  local repo=$2

  if [ -z "$owner" ] || [ -z "$repo" ]; then
    echo "Usage: $0 <owner> <repository>"
    exit 1
  fi

  # Get latest tag using GitHub API
  latest_tag=$(curl -s "https://api.github.com/repos/$owner/$repo/tags" | grep '"name":' | head -n 1 | cut -d'"' -f4)

  if [ -z "$latest_tag" ]; then
    echo "No tags found or error accessing repository"
    exit 1
  fi

  echo "$latest_tag"
}

function upgrade_ollama() {
  local LATEST_VER="$1"
  local CURRENT_VER

  CURRENT_VER=$(ollama --version | awk '{print $4}')
  CURRENT_VER="v$CURRENT_VER"

  if [ "$LATEST_VER" != "$CURRENT_VER" ]; then
    echo "ollama"
    echo "- latest version: $LATEST_VER"
    echo "- installed version: $CURRENT_VER"
    curl -fsSL https://ollama.com/install.sh | sh
  else
    echo "Ollama is the latest version. $LATEST_VER"
  fi

  ollama list | tail -n +2 | awk '{print $1}' | while read -r model; do
    ollama pull $model
  done
}

function upgrade_openweb_ui() {
  local INSTALL_PATH="$1"
  local VENV_ACTIVE_CMD=$2
  local LATEST_VER="$3"
  local CURRENT_VER

  pushd "$INSTALL_PATH" >/dev/null || raise_err "OpenWebUI dir not exists: $INSTALL_PATH"

  if [ -n "$VENV_ACTIVE_CMD" ]; then
    eval "$VENV_ACTIVE_CMD"
  fi
  CURRENT_VER=$(pip list | grep open-webui | awk '{print $2}')
  CURRENT_VER="v$CURRENT_VER"

  if [ "$LATEST_VER" != "$CURRENT_VER" ]; then
    echo "OpenWebUI"
    echo "- latest version: $LATEST_VER"
    echo "- installed version: $CURRENT_VER"
    pip install --upgrade open-webui
  else
    echo "OpenWebUI is the latest version. $LATEST_VER"
  fi

  popd >/dev/null || raise_err "$PWD Removed"
}

function find_git_repos() {
  find . -type d -name ".git" -exec dirname {} \;
}

function upgrade_comfyui() {
  local INSTALL_PATH="$1"
  local VENV_ACTIVE_CMD=$2
  local LATEST_VER="$3"
  local CURRENT_VER

  pushd "$INSTALL_PATH" >/dev/null || raise_err "ComfyUI dir not exists: $INSTALL_PATH"

  if [ -n "$VENV_ACTIVE_CMD" ]; then
    eval "$VENV_ACTIVE_CMD"
  fi

  #  CURRENT_VER=$(python3 -c"import comfyui_version;print (comfyui_version.__version__)")
  #  CURRENT_VER="v$CURRENT_VER"

  #  if [ "$LATEST_VER" != "$CURRENT_VER" ]; then
  #    echo "ComfyUI"
  #    echo "- latest version: $LATEST_VER"
  #    echo "- installed version: $CURRENT_VER"
  #    #pip install --upgrade open-webui
  #  else
  #    echo "ComfyUI is the latest version. $LATEST_VER"
  #  fi

  # comfyui upgrade by git
  echo "ComfyUI Upgrade by git"
  git_repos=$(find_git_repos)
  if [ -n "$git_repos" ]; then
    for git_repo in $git_repos; do
      pushd "$git_repo" >/dev/null || raise_err "$(readlink -f "$git_repo") dir call error"
      git pull --recurse-submodules >/dev/null || pip install -r requirements.txt --upgrade
      popd >/dev/null || raise_err "$(readlink -f "$PWD") dir call error"
    done
  fi

  CURRENT_VER=$(python3 -c"import comfyui_version;print (comfyui_version.__version__)")
  echo "- latest version: $LATEST_VER"
  echo "- installed version: $CURRENT_VER"

  popd >/dev/null || raise_err "$PWD Removed"
}

upgrade_ollama "$(get_latest_tag "ollama" "ollama")"
upgrade_openweb_ui /mnt/ai/home/vraptor/open-webui "source venv/bin/activate" "$(get_latest_tag "open-webui" "open-webui")"
upgrade_comfyui /mnt/ai/home/vraptor/comfyui/workspace "source /mnt/ai/home/vraptor/comfyui/venv/bin/activate" "$(get_latest_tag "comfyanonymous" "ComfyUI")"

#echo -e "\nChecking Python package versions..."
#pip list | grep -E "celery|click|django|jinja2|pyopenssl|pytz|pyyaml|redis|requests|six|tornado|twisted|wheel"

# pushd [dir] / popd
# dirs -l -v / dirs +N / dirs -N
# dirs +1 / dirs -1
