#!/bin/bash

source /usr/lib/ujust/ujust.sh

IMAGE_INFO="/usr/share/ublue-os/image-info.json"
DEFAULT_IMAGE=$(jq -r '."image-name"' < $IMAGE_INFO)
DEFAULT_BRANCH=stable

# Cache skopeo output if possible (only if bkt already exists)
if command -v bkt >/dev/null 2>&1; then
    skopeo() { bkt --ttl=30m --stale=5m -- command skopeo "$@"; }
fi

helptext="====== Bazzite Rollback Helper Util ======

This tool helps with rollbacks and rebases. Warns when switching between KDE and GNOME variants.

Usage: $(basename "$0") [OPTION] [ARGUMENT]

Options:
  list [BRANCH]      List available images (default: $DEFAULT_BRANCH)
  rollback           Rollback to previous image
  current            Show current image
  rebase [TARGET]    Rebase to target (format: TAG, IMAGE:TAG, or full ostree ref)

Examples:
  $(basename "$0") list stable
  $(basename "$0") rebase testing
  $(basename "$0") rebase bazzite-gnome:stable

For help: https://discord.bazzite.gg"

VARIANTS=(
  "bazzite"
  "bazzite-deck"
  "bazzite-nvidia"
  "bazzite-nvidia-open"
  "bazzite-deck-nvidia"
  "bazzite-gnome"
  "bazzite-gnome-nvidia"
  "bazzite-gnome-nvidia-open"
  "bazzite-deck-gnome"
  "bazzite-dx"
  "bazzite-dx-gnome"
  "bazzite-dx-nvidia"
  "bazzite-dx-nvidia-gnome"
)

# Utility functions
confirm() { 
  local msg="$1"
  echo -n "$msg ${yellow}(Y/n)${normal} "
  read -r reply
  [[ -z "$reply" || "${reply,,}" =~ ^y ]]
}

press_any_key() {
  echo "Press Enter to continue..."
  read -r
}

choose() { Choose "$@"; }

current_de() {
  if [[ "$XDG_CURRENT_DESKTOP" == *"GNOME"* ]] || [[ "$DESKTOP_SESSION" == *"gnome"* ]] || command -v gnome-shell >/dev/null 2>&1; then
    echo "gnome"
  elif [[ "$XDG_CURRENT_DESKTOP" == *"KDE"* ]] || [[ "$DESKTOP_SESSION" == *"plasma"* ]] || command -v plasmashell >/dev/null 2>&1; then
    echo "kde"
  else
    echo "unknown" # if either of the above are not detected, we can still serve warnings for custom image rebases
  fi
}

target_de() {
  [[ "$1" == *"gnome"* ]] && echo "GNOME" || echo "KDE/Plasma"
}

signing_scheme() {
  local current=$(rpm-ostree status -b --json | jq -r '.deployments[]["container-image-reference"]')
  [[ "$current" == *"ostree-image-signed"* ]] && echo "ostree-image-signed:docker://ghcr.io" || echo "ostree-unverified-registry:ghcr.io"
}

warn_if_de_mismatch() {
  local target_img="$1"
  local current=$(current_de)
  local target=$(target_de "$target_img")
  
  local warn=false
  local msg=""
  
  case "$current" in
    "gnome") [[ "$target_img" != *"gnome"* ]] && warn=true && msg="You're running GNOME but switching to $target ($target_img)" ;;
    "kde") [[ "$target_img" == *"gnome"* ]] && warn=true && msg="You're running KDE/Plasma but switching to $target ($target_img)" ;;
  esac
  
  if [[ "$warn" == true ]]; then
    echo
    echo "${red}${bold}⚠️  DESKTOP ENVIRONMENT WARNING ⚠️${normal}"
    echo "${yellow}$msg${normal}"
    echo "${yellow}Switching DEs may cause issues and is unsupported. Consider backing up data.${normal}"
    echo
    confirm "Continue anyway?" || return 1
  fi
  return 0
}

skopeo_tags() {
  local provider="$1"
  local image="$2"
  local filter="${3:-all}"
  
  local regex='.Tags[]'
  case "$filter" in
    stable) regex="$regex | select(test(\"^stable\"))" ;;
    testing) regex="$regex | select(test(\"^testing\"))" ;;
  esac
  
  skopeo list-tags "docker://ghcr.io/$provider/$image" | jq -r "$regex" | sort
}

available_variants() {
  local provider="${1:-ublue-os}"
  echo >&2 "Checking variants from $provider..."
  
  for variant in "${VARIANTS[@]}"; do
    skopeo list-tags "docker://ghcr.io/$provider/$variant" >/dev/null 2>&1 && echo "$variant"
  done
}

print_status() {
  echo "${bold}Current System Status:${normal}"
  local current=$(rpm-ostree status --json | jq -r '.deployments[] | select(.booted == true) | .["container-image-reference"]')
  # Only treat deployments with .staged==true as pending.
  local pending=$(rpm-ostree status --json | jq -r '.deployments[] | select(.staged == true) | .["container-image-reference"]' | head -1)
  
  [[ -n "$current" ]] && echo "  ${bold}Active:${normal} ${green}$current${normal}"
  [[ -n "$pending" && "$pending" != "null" ]] && echo "  ${bold}Pending:${normal} ${yellow}$pending${normal} ${yellow}(next boot)${normal}"
  echo
}

rebase_to() {
  local target="$1"
  local skip_confirm="${2:-false}"
  
  # Extract image name for DE checking
  local img_name
  case "$target" in
    ostree-*) img_name=$(echo "$target" | sed 's/.*\/\([^:]*\).*/\1/') ;;
    *:*) img_name=$(echo "$target" | cut -d':' -f1) ;;
    *) img_name="$DEFAULT_IMAGE" ;;
  esac
  
  warn_if_de_mismatch "$img_name" || return 1
  
  # Build full ostree ref if needed
  local full_ref="$target"
  case "$target" in
    ostree-*) ;;
    *:*) 
      local scheme=$(signing_scheme)
      full_ref="$scheme/ublue-os/$target"
      ;;
    *)
      local scheme=$(signing_scheme)
      full_ref="$scheme/ublue-os/$DEFAULT_IMAGE:$target"
      ;;
  esac
  
  echo "Rebasing to: ${bold}$full_ref${normal}"
  
  if [[ "$skip_confirm" != true ]]; then
    confirm "Continue?" || { echo "Cancelled."; return 1; }
  fi
  
  echo "Performing rebase..."
  if rpm-ostree rebase "$full_ref"; then
    echo "${green}Rebase successful!${normal} Reboot to apply changes."
  else
    echo "${red}Rebase failed.${normal} Image may not exist or be accessible."
    return 1
  fi
}

pick_image() {
  local provider="${1:-ublue-os}"
  local variants=$(available_variants "$provider")
  
  if [[ -z "$variants" ]]; then
    echo >&2 "${red}No variants found for provider: $provider${normal}"
    return 1
  fi
  
  local options=()
  while IFS= read -r line; do
    [[ -n "$line" ]] && options+=("$line")
  done <<< "$variants"
  options+=("Back")
  
  echo >&2 "${bold}Available images from ${green}$provider${normal}:"
  choose "${options[@]}"
}

pick_version() {
  local provider="$1"
  local image="$2"
  local filter="${3:-all}"
  
  local versions=$(skopeo_tags "$provider" "$image" "$filter")
  if [[ -z "$versions" ]]; then
    echo >&2 "${red}No $filter versions found for $image${normal}"
    return 1
  fi
  
  local options=()
  while IFS= read -r line; do
    [[ -n "$line" ]] && options+=("$line")
  done <<< "$versions"
  options+=("Back")
  
  echo >&2 "${bold}$filter versions for ${green}$image${normal}:"
  choose "${options[@]}"
}

choose_filter() {
  local provider="$1"
  local image="$2"
  local options=()

  # Only include Stable/Testing if tags exist for that image
  if [[ -n "$(skopeo_tags "$provider" "$image" stable 2>/dev/null)" ]]; then
    options+=("Stable")
  fi
  if [[ -n "$(skopeo_tags "$provider" "$image" testing 2>/dev/null)" ]]; then
    options+=("Testing")
  fi

  options+=("All" "Back")
  echo >&2 "${bold}Available tracks for ${green}$image${normal}:"
  choose "${options[@]}"
}

# Command functions
list_images() {
  local branch="${1:-$DEFAULT_BRANCH}"
  
  if [[ "$branch" == "all" ]]; then
    echo >&2 "Listing all available images..."
    skopeo_tags "ublue-os" "bazzite" "all"
  else
    echo >&2 "Listing $branch images..."
    skopeo_tags "ublue-os" "bazzite" "$branch"
  fi
}

current() {
  rpm-ostree status -vb
}

rollback() {
  local skip_confirm=false

  if [[ "$1" == "-y" || "$1" == "--yes" ]]; then
    skip_confirm=true
  fi

  if $skip_confirm || confirm "Rollback to previous image?"; then
    rpm-ostree rollback && echo >&2 "Reboot for changes to take effect"
  else
    echo "Cancelled."
  fi
}

rebase() {
  local target="$1"
  local skip_confirm=false
  
  if [[ "$2" == "-y" || "$2" == "--yes" ]]; then
    skip_confirm=true
  fi
  
  if [[ -z "$target" ]]; then
    echo >&2 "Usage: rebase TARGET"
    echo >&2 "Examples: rebase stable, rebase bazzite-gnome:testing"
    return 1
  fi
  
  rebase_to "$target" "$skip_confirm"
}

interactive_menu() {
  while true; do
    clear
    echo "$helptext"
    echo
    print_status
    
    local action=$(choose "List Images" "Show Current" "Rollback" "Rebase" "Custom Provider" "Exit")
    
    case "$action" in
      "List Images")
        clear
        local variants=$(available_variants "ublue-os")
        local img=$(pick_image "ublue-os")
        [[ "$img" != "Back" ]] && { 
          clear
          local filter=$(choose_filter "ublue-os" "$img")
          [[ "$filter" != "Back" ]] && {
            clear
            skopeo_tags "ublue-os" "$img" "${filter,,}"
            press_any_key
          }
        }
        ;;
        
      "Show Current")
        clear
        current
        press_any_key
        ;;
        
      "Rollback")
        clear
        rollback
        press_any_key
        ;;
        
      "Rebase")
        clear
        local img=$(pick_image "ublue-os")
        if [[ "$img" != "Back" ]]; then
          local filter=$(choose_filter "ublue-os" "$img")
          if [[ "$filter" != "Back" ]]; then
            local version=$(pick_version "ublue-os" "$img" "${filter,,}")
            [[ "$version" != "Back" && -n "$version" ]] && {
              clear
              rebase_to "$(signing_scheme)/ublue-os/$img:$version"
              press_any_key
            }
          fi
        fi
        ;;
        
      "Custom Provider")
        clear
        echo "Enter provider name (default: ublue-os):"
        read -r provider
        provider="${provider:-ublue-os}"
        provider=$(echo "$provider" | tr '[:upper:]' '[:lower:]')
        
        local img=$(pick_image "$provider")
        if [[ "$img" != "Back" ]]; then
          echo "Enter image:tag or press Enter to browse versions:"
          read -r manual_target
          if [[ -n "$manual_target" ]]; then
            rebase_to "$(signing_scheme)/$provider/$manual_target"
          else
            local filter=$(choose_filter "$provider" "$img")
            if [[ "$filter" != "Back" ]]; then
              local version=$(pick_version "$provider" "$img" "${filter,,}")
              [[ "$version" != "Back" && -n "$version" ]] && rebase_to "$(signing_scheme)/$provider/$img:$version"
            fi
          fi
        fi
        press_any_key
        ;;
        
      "Exit")
        echo "Exiting..."
        exit 0
        ;;
    esac
  done
}

case "$1" in
  "list") list_images "$2" ;;
  "rollback") rollback "$2" ;;
  "current") current ;;
  "rebase") rebase "$2" "$3" ;;
  "help"|"-h"|"--help"|"") interactive_menu ;;
  *) echo "Unknown option: $1. Run without arguments for interactive mode." ;;
esac