yt-search-play

#!/bin/bash

## Where am I?

script_path=$(realpath $0)
script_dir="$(dirname "$script_path")"
default_configfile="$script_dir/.default.config.json"

## Check dependencies

check_failed=

declare -A deps=( \
  ["rofi"]="The dependency rofi is missing. See https://github.com/davatorium/rofi/blob/next/INSTALL.md for instructions on how to install it." \
  ["jq"]="The dependency jq is missing. Run \"sudo apt install jq\" to install it, or see https://stedolan.github.io/jq/download for further details." \
  ["youtube-dl"]="The dependency youtube-dl is missing. See https://github.com/ytdl-org/youtube-dl/blob/master/README.md#installation for instructions on how to install it." \
  ["mpv"]="The dependency mpv is missing. Run \"sudo apt install mpv\" to install it, or see https://mpv.io/installation for further details." \
  )

for dep in "${!deps[@]}"
do
  type "$dep" &>/dev/null
  if [[ $? -ne 0 ]]
  then
    check_failed=true
    >&2 echo "${deps[$dep]}"
  fi
done

if [[ "$check_failed" == "true" ]]
then
  exit 1
fi

## Check if using node/puppeteer features

USE_PPT=
ppt_check_failed=
type node &>/dev/null
if [[ $? -ne 0 ]]
then
  ppt_check_failed="The dependency node is missing. Run \"sudo apt install node\" to install it, or see https://nodejs.org/en/download/package-manager for further details.\n"
fi

type npm &>/dev/null
if [[ $? -ne 0 ]]
then
  ppt_check_failed="${ppt_check_failed}The dependency npm is missing. Run \"sudo apt install npm\" to install it, or see https://nodejs.org/en/download/package-manager for further details.\n"
fi

if [[ -z "$ppt_check_failed" ]]
then
  if [[ "$(cd "$script_dir/ppt" > /dev/null 2>&1 && npm list --depth 1 puppeteer >/dev/null 2>&1; echo $?)" != "0" ]]
  then
    ppt_check_failed="Can't find the puppeteer module. Have you run \"npm install\" inside the ppt directory?"
  else
    USE_PPT=true
  fi
fi

## Variable defaults

search_size=30
max_history_size=50
max_cache_age=60
data_dir=".data"
force_no_cache=false
force_no_history=false
subs_mode=false
wl_mode=false
reverse_mode=false
use_max_downloads=false
mark_watched=false
use_thumbnails=false
thumbnail_size=2

## Custom rofi exit codes

ROFI_OK=0
ROFI_MORE=10
ROFI_SEND_LL=11
ROFI_SEND_WL=12
ROFI_REFRESH=19
ROFI_MODE_NORMAL=20
ROFI_MODE_SUBS=21
ROFI_MODE_WL=22

## Cleanup on exit
# trap cleanup EXIT
# cleanup () {
  # echo "Exiting..."
# }

## Define functions

dump_config () {
jq -n --arg search_size "$search_size" --arg max_history_size "$max_history_size" --arg max_cache_age "$max_cache_age" --arg data_dir "$data_dir" --arg force_no_cache "$force_no_cache" --arg force_no_history "$force_no_history" --arg subs_mode "$subs_mode" --arg wl_mode "$wl_mode" --arg reverse_mode "$reverse_mode" --arg use_max_downloads "$use_max_downloads" --arg mark_watched "$mark_watched" --arg use_thumbnails "$use_thumbnails" --arg thumbnail_size "$thumbnail_size" '{search_size: $search_size, max_history_size: $max_history_size, max_cache_age: $max_cache_age, data_dir: $data_dir, force_no_cache: $force_no_cache, force_no_history: $force_no_history, subs_mode: $subs_mode, wl_mode: $wl_mode, reverse_mode: $reverse_mode, use_max_downloads: $use_max_downloads, mark_watched: $mark_watched, use_thumbnails: $use_thumbnails, thumbnail_size: $thumbnail_size}'
}

print_help () {
echo -en "The following options are available:
--generate-config: Generate a default configuration file and exit
--dump-config: Print the loaded config file's location and the loaded configuration and exit
--config [file]: Use file as the configuration file
--no-config: Use default configurations
--clear-history: Clear the search history and exit
--clear-cache: Clear the cache and exit
-s|--subs: Fetch videos from a YouTube account's subscription feed instead of using a search query
-wl|--watch-later: Fetch videos from a YouTube account's Watch Later playlist instead of using a search query
-n [num]|--search-size [num]: Fetch a maximum of num videos
--max-cache-age [num]: Set the number of seconds before a cached search expires
-r|--reverse: Reverse the order of the displayed videos
-m|--mark-watched: Mark videos as watched on YouTube
-t|--thumbnails: Download and display video thumbnails
--thumbnail-size [num]: How large to display thumbnails, relative to text size.
--use-max-downloads: This uses youtube-dl's --max-downloads flag instead of --playlist-end internally
-C|--force-no-cache: Prevent the program from reading from or writing to the cache
-H|--force-no-history: Disable the search history
-h|--help: Print help and exit

Pressing Alt+1 switches to normal (search) mode
Pressing Alt+2 switches to subscriptions mode
Pressing Alt+3 switches to watch later mode

In normal mode, pressing Enter will select the currently highlighted history entry, or search if there are no highlighted entries. Pressing Ctrl+Enter will always search.
Video results may be toggled for inclusion in the watch later or liked videos playlists by pressing Alt+w or Alt+l, respectively. You can use Shift+Enter to select multiple videos before toggling.
Pressing Alt+r will refresh the search if there is a stale cache.
Pressing Alt+m will load more video results, if available."
}

find_config () {
  if [[ -s "$passed_configfile" ]]
  then
    echo "$passed_configfile"
  elif [[ -s "$script_dir/config.json" ]]
  then
    echo "$script_dir/config.json"
  elif [[ -s "$HOME/.config/yt-search-play/config.json" ]]
  then
    echo "$HOME/.config/yt-search-play/config.json"
  else
    echo "$default_configfile"
  fi
}

load_config () {
  configfile="$(find_config)"

  if [[ -n "$configfile" ]]
  then
    config=$(jq '.' "$configfile")
    
    if [[ -z "$config" ]]
    then
      >&2 echo "Error reading config in $configfile."
      exit 1
    fi

    search_size=$(echo "$config" | jq -r --arg default "$search_size" '. | if has("search_size") then .search_size else $default end')
    max_history_size=$(echo "$config" | jq -r --arg default "$max_history_size" '. | if has("max_history_size") then .max_history_size else $default end')
    max_cache_age=$(echo "$config" | jq -r --arg default "$max_cache_age" '. | if has("max_cache_age") then .max_cache_age else $default end')
    data_dir=$(echo "$config" | jq -r --arg default "$data_dir" '. | if has("data_dir") then .data_dir else $default end')
    force_no_cache=$(echo "$config" | jq -r --arg default "$force_no_cache" '. | if has("force_no_cache") then .force_no_cache else $default end')
    force_no_history=$(echo "$config" | jq -r --arg default "$force_no_history" '. | if has("force_no_history") then .force_no_history else $default end')
    subs_mode=$(echo "$config" | jq -r --arg default "$subs_mode" '. | if has("subs_mode") then .subs_mode else $default end')
    wl_mode=$(echo "$config" | jq -r --arg default "$wl_mode" '. | if has("wl_mode") then .wl_mode else $default end')
    reverse_mode=$(echo "$config" | jq -r --arg default "$reverse_mode" '. | if has("reverse_mode") then .reverse_mode else $default end')
    use_max_downloads=$(echo "$config" | jq -r --arg default "$use_max_downloads" '. | if has("use_max_downloads") then .use_max_downloads else $default end')
    mark_watched=$(echo "$config" | jq -r --arg default "$mark_watched" '. | if has("mark_watched") then .mark_watched else $default end')
    use_thumbnails=$(echo "$config" | jq -r --arg default "$use_thumbnails" '. | if has("use_thumbnails") then .use_thumbnails else $default end')
    thumbnail_size=$(echo "$config" | jq -r --arg default "$thumbnail_size" '. | if has("thumbnail_size") then .thumbnail_size else $default end')
  fi

  if ! [[ "$search_size" =~ ^[0-9]+$ ]]
  then
    >&2 echo "Error in $configfile: search_size must be an integer."
    exit 1
  fi
  if ! [[ "$max_history_size" =~ ^[0-9]+$ ]]
  then
    >&2 echo "Error in $configfile: max_history_size must be an integer."
    exit 1
  fi
  if ! [[ "$max_cache_age" =~ ^[0-9]+$ ]]
  then
    >&2 echo "Error in $configfile: max_cache_age must be an integer."
    exit 1
  fi
  if ! [[ "$thumbnail_size" =~ ^[0-9]+(\.[0-9]+)?$ ]]
  then
    >&2 echo "Error in $configfile: thumbnail_size must be a number."
    exit 1
  fi

  case $data_dir in
    /*) mkdir -p "$data_dir" ;;
    *) mkdir -p "$script_dir/$data_dir" ;;
  esac
  if [[ $? -ne 0 ]]
  then
    >&2 echo "Error in $configfile: $data_dir is not a valid directory."
    exit 1
  fi

  historyfile="$script_dir/$data_dir/search-history"
  #TODO: allow for passed in cookiefile
  cookiefile="$script_dir/$data_dir/cookies.txt"
  cachefile="$script_dir/$data_dir/cache.json"
  thumb_dir="$script_dir/$data_dir/thumbs"
  
  if [[ ! -f "$historyfile" ]]
  then
    touch -a "$historyfile"
  fi
  history_size=$(< "$historyfile" wc -l 2>/dev/null || echo 0)
  playlist_start=1
}

clear_history () {
  rm -f "$historyfile"
}

clear_cache () {
  rm -f "$cachefile"
  rm -f "$thumb_dir"/*.jpg
}

update_history () {
  if [[ -n "$search_query" ]]
  then
    sed -i '\`'"$search_query"'`d' "$historyfile"
    echo "$search_query" >> "$historyfile"
  fi
  if [[ -n "$url" ]]
  then
    sed -i '\`'"$history_entry"'`d' "$historyfile"
    echo "$history_entry" >> "$historyfile"
  fi

  echo "$(tail -n $max_history_size "$historyfile")" > "$historyfile"
}

purge_expired_cache () {
  local expiry_window="$(($(date +%s)-$max_cache_age))"
  for dead_thumb in $(jq -r --arg time "$expiry_window" --arg td "$thumb_dir" '.[] | select(.time < $time) | "\($td)/\(.cache[].id).jpg"' "$cachefile")
  do
    rm -f "$dead_thumb"
  done
  jq --arg time "$expiry_window" 'del(.[] | select(.time < $time))' "$cachefile" > "$cachefile.tmp" && mv "$cachefile.tmp" "$cachefile"
}

write_to_cache () {
  if [[ ! -s "$cachefile" ]] || [[ -z "$(cat "$cachefile" | tr -d '[:space:]')" ]]
  then
    echo "[]" > "$cachefile"
  fi

  local cache=$(echo "$2" | jq --slurp '.')
  
  if [[ -n "$cache" ]] && [[ "$cache" != "[]" ]]
  then
    jq --arg sel "$1" 'del(.[] | select(.entry == $sel))' "$cachefile" | jq --arg sel "$1" --arg time "$(date +%s)" --argjson cc "$cache" '. += [{entry: $sel, time: $time, cache: $cc, size: ($cc|length)}]' > "$cachefile.tmp" && mv "$cachefile.tmp" "$cachefile"
  fi
}

read_from_cache () {
  if [[ -s "$cachefile" ]]
  then
    purge_expired_cache
    local entry_found=$(jq --arg sel "$1" '.[] | select(.entry == $sel)' "$cachefile")
    if [[ -n "$entry_found" ]] 
    then
      cache_size=$(echo "$entry_found" | jq -r '.size')
      if [[ $cache_size -lt $playlist_end ]]
      then
        partial_cache="true"
        playlist_start=$((cache_size+1))
      fi
      cache_value=$(echo "$entry_found" | jq -r '.cache[]')
    fi
  fi 
}

find_video () {
  local search_url="$1"
  local found_results=
  cache_value=
  partial_cache=
  history_entry=
  selection=
  idx=
  url=

  if [[ "$subs_mode" == "true" ]] || [[ "$wl_mode" == "true" ]]
  then
    # check/update this file if sub feed breaks
    if [[ -s "$cookiefile" ]]
    then
      local include_cookies="--cookies $cookiefile"
    else
      >&2 echo "A valid cookies.txt file must exist in $data_dir."
      exit 1
    fi
  fi
  
  if [[ "$reverse_mode" == "true" ]] || ([[ "$reverse_mode" == "wl" ]] && [[ "$wl_mode" == "true" ]])
  then
    reverse="--playlist-reverse"
  else
    reverse=
  fi

  # this will be deprecated in future
  if [[ "$use_max_downloads" == "true" ]] || ([[ "$use_max_downloads" == "wl" ]] && [[ "$wl_mode" == "true" ]])
  then
    playlist_ender="--max-downloads"
  else
    playlist_ender="--playlist-end"
  fi

  [[ "$refresh_cache" != "true" ]] && [[ "$force_no_cache" != "true" ]] && read_from_cache "$search_url"
  
  refresh_cache=false
  if [[ -z "$cache_value" ]] || [[ "$partial_cache" == "true" ]]
  then
    rofi -e "Loading videos..." &
    mesgpid=$!
    # reverse and max-downloads will get all videos, used due to below
    # https://github.com/ytdl-org/youtube-dl/pull/24487
    local search="$(youtube-dl -i $reverse --playlist-start $playlist_start $playlist_ender $playlist_end $include_cookies -j "$search_url" --flat-playlist | jq -r --unbuffered '. | {title: .title, uploader: .uploader, id: .url, url: ("http://www.youtube.com/watch?v=" + .url)}')"
    { kill $mesgpid && wait $mesgpid; } &>/dev/null
    
    if [[ "$use_thumbnails" == "true" ]]
    then
      get_thumbnails
    fi

    local found_results="$(jq --slurp '.[]' <(echo "$cache_value") <(echo "$search"))"

    [[ "$force_no_cache" != "true" ]] && write_to_cache "$search_url" "$found_results"
  else
    local found_results="$cache_value"
  fi

  if [[ "$use_thumbnails" == "true" ]]
  then
    iconstr="\"\(.title) :: \(.uploader)\\\0icon\\\x1f${thumb_dir}/\(.id).jpg\""
    useicons="-show-icons"
    themestr="element-icon{size:${thumbnail_size}em;}"
  else
    iconstr="\"\(.title) :: \(.uploader)\""
    useicons="-no-show-icons"
  fi

  idx=$(echo -e "$(echo "$found_results" | jq -r "$iconstr" | head -n $playlist_end)" | rofi -dmenu -i -p 'Select Video' -l 10 -scroll-method 0 -format i -async-pre-read 0 -selected-row $((playlist_end-search_size)) -multi-select $useicons -theme-str "$themestr" -kb-custom-1 Alt+m -kb-custom-2 Alt+l -kb-custom-3 Alt+w -kb-custom-10 Alt+r -kb-custom-11 Alt+1 -kb-custom-12 Alt+2 -kb-custom-13 Alt+3; exit $?)
  rofi_exit_code=$?

  if [[ $rofi_exit_code -eq $ROFI_OK ]] && [[ -n "$idx" ]] && idx=$(echo "$idx" | tail -1) && [[ $idx -ge 0 ]]
  then
    selection=$(echo "$found_results" | jq --slurp -r --arg idx "$idx" '.[$idx|tonumber]')
    url=$(echo "$selection" | jq -r '.url')
    history_entry="$(echo "$selection" | jq -r '"\(.title) :: \(.uploader) :: \(.url)"')"
  elif [[ $rofi_exit_code -eq $ROFI_MORE ]]
  then
    playlist_end=$((playlist_end+search_size))
    find_video "$1"
  else
    if [[ $rofi_exit_code -eq $ROFI_SEND_WL ]] && [[ "$idx" =~ ^([0-9]+$'\n')*[0-9]+$ ]]
    then
      send_to_node "$found_results" "$idx" wl
      rofi_exit_code=$ROFI_REFRESH
    elif [[ $rofi_exit_code -eq $ROFI_SEND_LL ]] && [[ "$idx" =~ ^([0-9]+$'\n')*[0-9]+$ ]]
    then
      send_to_node "$found_results" "$idx" like
      rofi_exit_code=$ROFI_REFRESH
    elif [[ $rofi_exit_code -eq $ROFI_MODE_NORMAL ]]
    then
      subs_mode=false
      wl_mode=false
      do_search
    elif [[ $rofi_exit_code -eq $ROFI_MODE_SUBS ]]
    then
      subs_mode=true
      wl_mode=false
      do_search
    elif [[ $rofi_exit_code -eq $ROFI_MODE_WL ]]
    then
      subs_mode=false
      wl_mode=true
      do_search
    fi

    if [[ $rofi_exit_code -eq $ROFI_REFRESH ]]
    then
      playlist_start=1
      playlist_end=$search_size
      refresh_cache=true
      find_video "$1"
    fi
  fi
}

send_to_node () {
  if [[ "$USE_PPT" != "true" ]]
  then
    >&2 echo -en "You seem to be trying to use features you have not enabled:\n$ppt_check_failed"
    return 1
  fi

  local urls=
  local node_pids=
  local i=0
  for ii in $2
  do
    selections=$(echo "$1" | jq --slurp -r --arg idx "$ii" '.[$idx|tonumber]')
    local urls[$((i++))]=$(echo "$selections" | jq -r '.url')
  done

  rofi -e "Updating playlist, this may take a moment..." &
  mesgpid=$!
  local i=0
  for uuu in ${urls[*]}
  do
    node "$script_dir/ppt/index.js" "$uuu" $3 &>/dev/null &
    local node_pids[$((i++))]=$!
  done
  for node_pid in ${node_pids[*]}
  do
    wait $node_pid
  done
  { kill $mesgpid && wait $mesgpid; } &>/dev/null
}

get_thumbnails () {
  local thumb_pids=
  rofi -e "Fetching thumbnails..." &
  mesgpid=$!
  mkdir -p "$thumb_dir"
  local thumb_idx=0
  for id in $(echo "$search" | jq -r '.id')
  do
    if [[ ! -f "$thumb_dir/$id.jpg" ]]
    then
      wget -q --output-document="$thumb_dir/$id.jpg" "https://i.ytimg.com/vi/$id/hqdefault.jpg" &
      local thumb_pids[$((thumb_idx++))]=$!
    fi
  done
  for thumb_pid in ${thumb_pids[*]}
  do
    wait $thumb_pid
  done
  { kill $mesgpid && wait $mesgpid; } &>/dev/null
}

process_args () {
  if [[ $# -gt 0 ]]
  then
    if [[ "$@" =~ --generate-config ]]
    then
      echo "Default config generated in: $default_configfile"
      exit
    elif [[ "$@" =~ --dump-config ]]
    then
      echo "Config file located at: $configfile"
      dump_config
      exit
    elif [[ "$@" =~ --clear-history ]] && [[ "$@" =~ --clear-cache ]]
    then
      clear_history
      clear_cache
      exit
    elif [[ "$@" =~ --clear-history ]]
    then
      clear_history
      exit
    elif [[ "$@" =~ --clear-cache ]]
    then
      clear_cache
      exit
    else
      while [[ $# -gt 0 ]]
      do
        key="$1"
        case $key in
          -s|--subs)
            subs_mode=true
            shift
            ;;
          -wl|--watch-later)
            wl_mode=true
            shift
            ;;
          -r|--reverse)
            reverse_mode=true
            shift
            ;;
          --max-cache-age)
            if ! [[ "$2" =~ ^[0-9]+$ ]]
            then
             >&2 echo "$1 must be an integer."
              exit 1
            fi
            max_cache_age="$2"
            shift
            shift
            ;;
          -m|--mark-watched)
            mark_watched=true
            shift
            ;;
          -t|--thumbnails)
            use_thumbnails=true
            shift
            ;;
          --thumbnail-size)
            if ! [[ "$2" =~ ^[0-9]+(\.[0-9]+)?$ ]]
            then
             >&2 echo "$1 must be a number."
              exit 1
            fi
            thumbnail_size="$2"
            shift
            shift
            ;;
          --use-max-downloads)
            use_max_downloads=true
            shift
            ;;
          -n|--search-size)
            if ! [[ "$2" =~ ^[0-9]+$ ]]
            then
             >&2 echo "$1 must be an integer."
              exit 1
            fi
            search_size="$2"
            shift
            shift
            ;;
          -C|--force-no-cache)
            force_no_cache=true
            shift
            ;;
          -H|--force-no-history)
            force_no_history=true
            shift
            ;;
          --config)
            passed_configfile="$2"
            load_config
            shift
            shift
            ;;
          --no-config)
            passed_configfile="$default_configfile"
            load_config
            shift
            ;;
          -h|--help)
            echo "$(print_help)"
            exit
            ;;
          *)
            echo "Unrecognised option $1" >&2
            exit 1
            ;;
        esac
      done
    fi
  fi
}

get_channel_url () {
  youtube-dl -s --flat-playlist "https://www.youtube.com/c/$1" &>/dev/null && [[ $? -eq 0 ]] && echo "https://www.youtube.com/c/$1/videos" &
  youtube-dl -s --flat-playlist "https://www.youtube.com/channel/$1" &>/dev/null && [[ $? -eq 0 ]] && echo "https://www.youtube.com/channel/$1/videos" &
  youtube-dl -s --flat-playlist "https://www.youtube.com/user/$1" &>/dev/null && [[ $? -eq 0 ]] && echo "https://www.youtube.com/user/$1/videos" &
  wait
}

do_search () {
  playlist_end=$search_size
  if [[ "$subs_mode" == "true" ]]
  then
    # flow=1 -> grid view, as list causes bugs!
    find_video "https://www.youtube.com/feed/subscriptions?flow=1"
  elif [[ "$wl_mode" == "true" ]]
  then
    find_video "https://www.youtube.com/playlist?list=WL"
  else
    if [[ "$force_no_history" != "true" ]]
    then
      local history_entries=$(tac "$historyfile")
    else
      history_size=0
    fi
    search_query=$(echo "$history_entries" | rofi -dmenu -i -p "Search YouTube" -no-show-icons -theme-str 'entry { placeholder: "Enter text or select recent search..."; }' -l $([[ $history_size -lt 10 ]] && echo "$history_size" || echo 10) -scroll-method 0 -kb-custom-1 '' -kb-custom-2 '' -kb-custom-3 '' -kb-custom-11 Alt+1 -kb-custom-12 Alt+2 -kb-custom-13 Alt+3; exit $?)
    sq_exit_code=$?

    if [[ $sq_exit_code -eq $ROFI_OK ]]
    then
      if [[ "$search_query" =~ (https?\:\/\/)?(www\.)?(youtube\.com|youtu\.?be)\/.+$ ]]
      then
        selection="$search_query"
        url="${BASH_REMATCH[0]}"
        history_entry="$search_query"
      elif [[ "$search_query" =~ ^@(.+)$ ]]
      then
        sanitised_query=${BASH_REMATCH[1]// /+}
        find_video "$(get_channel_url "$sanitised_query")"
      elif [[ -n "$search_query" ]]
      then
        sanitised_query=${search_query// /+}
        find_video "ytsearchall:$sanitised_query"
      else
        exit
      fi
    elif [[ $sq_exit_code -eq $ROFI_MODE_NORMAL ]]
    then
      subs_mode=false
      wl_mode=false
      do_search
    elif [[ $sq_exit_code -eq $ROFI_MODE_SUBS ]]
    then
      subs_mode=true
      wl_mode=false
      do_search
    elif [[ $sq_exit_code -eq $ROFI_MODE_WL ]]
    then
      subs_mode=false
      wl_mode=true
      do_search
    else
      exit $sq_exit_code
    fi
  fi
}

play_video () {
  if [[ -n "$url" ]]
  then
    pkill mpv
    mpv --really-quiet "$url" "$([[ "$mark_watched" == "true" ]] && [[ -s "$cookiefile" ]] && echo "--ytdl-raw-options=mark-watched=,cookies=$cookiefile")" &
  fi
}

## (Re)write default config

dump_config > "$default_configfile"

## Main thread

load_config
process_args "$@"
do_search
[[ "$force_no_history" != "true" ]] && update_history
play_video