#!/bin/bash

# musidl - download music from YouTube
# Copyright (C) 2021 TriMill
# 
# This program is free software: you can redistribute it and/or modify
# it under the terms of the GNU General Public License as published by
# the Free Software Foundation, either version 3 of the License, or
# (at your option) any later version.
# 
# This program is distributed in the hope that it will be useful,
# but WITHOUT ANY WARRANTY; without even the implied warranty of
# MERCHANTABILITY or FITNESS FOR A PARTICULAR PURPOSE.  See the
# GNU General Public License for more details.
# 
# You should have received a copy of the GNU General Public License
# along with this program.  If not, see <https://www.gnu.org/licenses/>.


# check to make sure all commands are available before running
if ! (command -v "yt-dlp" && command -v "jq" && command -v "eyeD3") &>/dev/null; then
    >&2 echo "Missing dependencies."
    exit 1
fi

# this is printed on -h/--help
helpmsg="Usage: musidl [options]... [url]... (-- [yt-dlp options])  
Download music from YouTube using yt-dlp and add ID3v2 metadata automatically

Options:
 -h --help              Print this message and exit
 -v --version           Print the version and exit
 -q --quiet             Only print errors
 -f --format <fmt>      The audio format to use (default: mp3)
 -a --artist <name>     Set the artist (default: detect from YouTube)
 -A --album  <name>     Set the album title (default: detect from YouTube)
 -t --song   <name>     Set the song title (default: detect from YouTube)
 -g --genre  <name>     Set the genre (default: none)
 -y --year   <year>     Set the year (default: detect from YouTube)
 -T --track  <track>    Set the track (default: detect from playlist, otherwise none)
 -s --show              Print information to stdout instead of downloading to file
 -o --output            Set the template used for filenames. See the yt-dlp documentation for more info
 --color     <when>     Enable or disable colored output. <when> must be always, auto (default), or never

Any arguments after '--' will be sent to yt-dlp directly. These are not officially
supported and some may cause problems."

versionmsg="1.0.0"

# Add a prefix to error messages
function printStderr {
    read -r input
    [ -n "$input" ] && >&2 echo -e "[$1] $input" 
}

urls=() # The list of URLs to download
arg_format="mp3" # The file format (-f/--format)
output_format=()
# Process the arguments
while [[ $# -gt 0 ]]; do
  case "$1" in
    -h|--help)
      echo "$helpmsg"; exit 0 ;;
    -v|--version)
      echo "musidl $versionmsg"; exit 0 ;;
    -q|--quiet)
      quiet="true"; shift 1 ;;
    -f|--format)
      arg_format="$2"; shift 2 ;;
    -a|--artist)
      i_artist="$2"; shift 2 ;;
    -A|--album)
      i_album="$2"; shift 2 ;;
    -t|--song)
      i_song="$2"; shift 2 ;;
    -g|--genre)
      genre="$2"; shift 2 ;;
    -y|--year)
      i_year="$2"; shift 2 ;;
    -T|--track)
      i_track="$2"; shift 2 ;;
    -s|--show)
      arg_show="true"; shift ;;
    -o|--output)
      output_format=( "-o" "$2" ); shift 2 ;;
    --color)
      color_mode="$2"; shift 2 ;;
    --)
      shift; break ;;
    *)
      urls+=("$1"); shift ;;
  esac
done

case "$color_mode" in
    always)
        color_enabled="true" ;;
    never)
        ;;
    auto|"")
        if tput colors &>/dev/null; then
            color_enabled="true"
        fi ;;
    *)
        >&2 echo "[musidl] invalid argument for --color: $color_mode" && exit 1
esac

if [ -n "$color_enabled" ]; then
    pref_musidl=$"\x1b[32m[musidl]\x1b[0m"
    pref_ytdl=$"\x1b[36m[yt-dlp]\x1b[0m"
    pref_jq=$"\x1b[35m[jq]\x1b[0m"
    errorstr=$"\x1b[31;1mERROR:\x1b[0m"
else
    pref_musidl="[musidl]"
    pref_ytdl="[yt-dlp]"
    pref_jq="[jq]"
    errorstr="ERROR:"
fi

# Exit if no URLs are given
[ ${#urls[@]} -eq 0 ] && >&2 echo -e "$pref_musidl $errorstr no URLs provided. Run '$0 --help' for help." && exit 1

ytdl_args=( "$@" )

IFS=$'\n'

# Download metadata from all the URLs
[ -z "$quiet" ] && echo -e "$pref_musidl downloading data..."
data=($(yt-dlp -j "${ytdl_args[@]}" "${urls[@]}" 2> >(printStderr "$pref_ytdl")))
[ -z "$data" ] && >&2 echo -e "$pref_musidl $errorstr yt-dlp returned nothing, exiting" && exit 1

[ -z "$quiet" ] && echo -e "$pref_musidl downloading videos..."

# Download the video,
# Print all errors,
# Add a prefix to the messages,
# And only print messages if quiet is off
if [ -z "$show" ]; then
    stdbuf -o0 yt-dlp --no-progress --extract-audio -f bestaudio \
        --audio-format="$arg_format" "${ytdl_args[@]}" ${output_format[@]} "${urls[@]}" \
        2> >(printStderr "$pref_ytdl") \
        | stdbuf -o0 sed -e 's/^/'"$pref_ytdl"' /' \
        | ( if [ -z "$quiet" ]; then stdbuf -o0 cat; else cat > /dev/null; fi )
fi

# Get the filenames
filenames=($(yt-dlp --get-filename --extract-audio -f bestaudio --audio-format="$arg_format" "${ytdl_args[@]}" ${output_format[@]} "${urls[@]}"))

# Run for each video to tag
for i in "${!data[@]}"; do
    line="${data[i]}"
    filename="${filenames[i]%.*}"."$arg_format"

    if ! [ -f "$filename" ]; then
        >&2 echo -e "$pref_musidl $errorstr failed to find $filename, skipping..."
        continue
    fi
    # Use jq to extract data from the metadata. The following options are considered in order:
    # Artist: artist, channel name
    # Album: album, playlist name, none
    # Song name: track, video title
    # Year: release year, upload date
    # Track number: playlist index, none
    # Video URL
    readarray -t songdata <<< "$(echo "$line" | jq -r \
        'if .artist then .artist else .channel end, if .album then .album else if .playlist_title then .playlist_title else "" end end, if .track then .track else .title end, if .release_year then .release_year else .upload_date end, if .playlist_index then .playlist_index else "" end, .webpage_url' \
        2> >(printStderr "$pref_jq"))"

    [ "${#songdata[@]}" -eq 0 ] && >&2 echo -e "$pref_musidl $errorstr could not parse JSON, skipping" && continue

    # If arguments are set manually prefer them, otherwise use ones from YouTube
    if [ -n "$i_artist" ]; then artist="$i_artist"; else
        artist="${songdata[0]}"
    fi
    if [ -n "$i_album" ]; then album="$i_album"; else
        album="${songdata[1]}"
    fi
    if [ -n "$i_song" ]; then title="$i_song"; else
        title="${songdata[2]}"
    fi
    if [ -n "$i_year" ]; then year="$i_year"; else
        year="$(echo "${songdata[3]}" | head -c 4)"
    fi
    if [ -n "$i_track" ]; then track="$i_track"; else
        track="${songdata[4]}"
    fi
    url="${songdata[5]}"

    # If -s/--show is passed, show the data and continue
    if [ -n "$arg_show" ]; then
        echo "URL: $url"
        [ -n "$artist" ] && echo "Artist: $artist"
        [ -n "$album" ] && echo "Album: $album"
        [ -n "$title" ] && echo "Title: $title"
        [ -n "$genre" ] && echo "Genre: $genre"
        [ -n "$year" ] && echo "Year: $year"
        [ -n "$track" ] && echo "Track: $track"
        continue
    fi

    [ -z "$quiet" ] && echo -e "$pref_musidl tagging $title - $artist ($url)"
    [ -n "$artist" ] && eyeD3 --artist "$artist" "$filename" --encoding utf8 >/dev/null
    [ -n "$album" ] && eyeD3 --album "$album" "$filename" --encoding utf8 >/dev/null
    [ -n "$title" ] && eyeD3 --title "$title" "$filename" --encoding utf8 >/dev/null
    [ -n "$genre" ] && eyeD3 --genre "$genre" "$filename" --encoding utf8 >/dev/null
    [ -n "$year" ] && eyeD3 --release-year "$year" "$filename" >/dev/null
    [ -n "$track" ] && eyeD3 --track "$track" "$filename" >/dev/null
done

[ -z "$quiet" ] && echo -e "$pref_musidl done!"