#!/bin/bash
#
# Install nerd-dictation offline speech2text tool script using VOSK backend/model.
#
# Follow the installer repo at:
# - https://forge.switnet.net/Ark74/simple-nerd-dictation-installer
#
# Nerd Dictation repo at:
# - https://github.com/ideasman42/nerd-dictation

set -e
# Reset
Color_Off='\e[0m'       # Text Reset
# Regular Colors
Black='\e[0;30m'        # Black
Red='\e[0;31m'          # Red
Green='\e[0;32m'        # Green
Yellow='\e[0;33m'       # Yellow
Blue='\e[0;34m'         # Blue
Purple='\e[0;35m'       # Purple
Cyan='\e[0;36m'         # Cyan

printwc() {
    printf "%b$2%b" "$1" "${Color_Off}"
}
print_title() {
printwc "${Blue}" "\n#--------------------------------------------------"
printwc "${Blue}" "\n# $1"
printwc "${Blue}" "\n#--------------------------------------------------\n"
}
printf '
########################################################################
                    Welcome to Nerd-Dictation Installer
########################################################################
# Install nerd-dictation offline speech2text tool script using VOSK backend/model.
#
# Follow the installer repo at:
# - https://forge.switnet.net/Ark74/simple-nerd-dictation-installer
#
# Nerd Dictation repo at:
# - https://github.com/ideasman42/nerd-dictation
\n'

if [ "$(id -u)" = 0 ]; then
   printwc "${Yellow}" "Please don't use with root privileges, sudo will be used only when necessary!"
   exit 0
fi
#--------------------------------------------------
print_title "Check for supported arch."
#--------------------------------------------------
if dpkg-architecture -e amd64 ;then
    printwc "${Green}" " amd64: supported architecture.\n"
else
    printwc "${Red}" " Warning: Only run this script on amd64 arch.\n"
    exit
fi

sudo apt-get update -q2

ENGLISH_VOSK_MODEL="https://alphacephei.com/vosk/models/vosk-model-en-us-0.22-lgraph.zip"
SPANISH_VOSK_MODEL="https://alphacephei.com/vosk/models/vosk-model-small-es-0.42.zip"
FRENCH_VOSK_MODEL="https://alphacephei.com/vosk/models/vosk-model-small-fr-0.22.zip"
GERMAN_VOSK_MODEL="https://alphacephei.com/vosk/models/vosk-model-small-de-0.15.zip"
CHINESE_VOSK_MODEL="https://alphacephei.com/vosk/models/vosk-model-small-cn-0.22.zip"
NERD_DIC_GIT_REPO="https://github.com/ideasman42/nerd-dictation.git"
OPT_VOSK="/opt/vosk-python"

#--------------------------------------------------
print_title "> Vosk Model Language selection."
#--------------------------------------------------
PS3='Select one of the available language models to use on nerd-dictation: '
options=("English" "Spanish" "French" "German" "Chinese" "Exit")
select opt in "${options[@]}"
do
    case $opt in
        "English")
            printf "\n  > Selecting English model.\n"
            VOSK_MODEL="$ENGLISH_VOSK_MODEL"
            break
            ;;
        "Spanish")
            printf "\n  > Selecting Spanish model.\n"
            VOSK_MODEL="$SPANISH_VOSK_MODEL"
            break
            ;;
        "French")
            printf "\n  > Selecting French model.\n"
            VOSK_MODEL="$FRENCH_VOSK_MODEL"
            break
            ;;
        "German")
            printf "\n  > Selecting German model.\n"
            VOSK_MODEL="$GERMAN_VOSK_MODEL"
            break
            ;;
        "Chinese")
            printf "\n  > Selecting Chinese model.\n"
            VOSK_MODEL="$CHINESE_VOSK_MODEL"
            break
            ;;
        "Exit")
            printf "\n  > Exit installer.\n"
            exit 0
            break
            ;;
        *) printwc "${Yellow}" "Invalid option $REPLY, choose one number of the above.";;
    esac
done
VOSK_FOLDER="$(basename $VOSK_MODEL|sed 's|.zip||')"

#--------------------------------------------------
print_title "Install dependencies"
#--------------------------------------------------
sudo apt-get install -y git ssh python3-setuptools pulseaudio-utils xdotool

#--------------------------------------------------
print_title "Install vosk python3 module."
#--------------------------------------------------
[ -d "$OPT_VOSK" ] && sudo rm -rf "$OPT_VOSK"
sudo mkdir "$OPT_VOSK"
cat << EOF  | sudo tee "$OPT_VOSK"/setup.py
from setuptools import setup

setup(
    install_requires=['vosk'],
    dependency_links=[
        'https://files.pythonhosted.org/packages/fc/ca/83398cfcd557360a3d7b2d732aee1c5f6999f68618d1645f38d53e14c9ff/vosk-0.3.45-py3-none-manylinux_2_12_x86_64.manylinux2010_x86_64.whl',
    ],
)
EOF
cd "$OPT_VOSK"
sudo python3 "$OPT_VOSK"/setup.py install
if [ $? = 0 ]; then
    sudo rm "$OPT_VOSK"/setup.py
else
    printwc "${Red}" "Something went wrong, please check setup.py file\n"
    exit
fi

#--------------------------------------------------
print_title "Cleaning existing local directories and binary link."
#--------------------------------------------------
rm -rf "$HOME"/.local/nerd-dictation
rm -f  "$HOME"/.local/bin/nerd-dictation
rm -rf "$HOME"/.config/nerd-dictation

#--------------------------------------------------
print_title "Set repo as local package."
#--------------------------------------------------
cd "$HOME"/.local
git clone "$NERD_DIC_GIT_REPO"
cd nerd-dictation

#--------------------------------------------------
print_title "Setup language model"
#--------------------------------------------------
wget "$VOSK_MODEL"
unzip "$(basename $VOSK_MODEL)"
if [ $? = 0 ]; then
    rm "$(basename $VOSK_MODEL)"
else
    printwc "${Red}" "Something went wrong, please check $(basename $VOSK_MODEL) file\n"
    exit
fi
mv "$VOSK_FOLDER" model
mkdir "$HOME"/.config/nerd-dictation
mv model "$HOME"/.config/nerd-dictation/

#--------------------------------------------------
print_title "Setup local bin for nerd-dictation"
#--------------------------------------------------
[ ! -d "$HOME"/.local/bin ] && mkdir "$HOME"/.local/bin
ln -s "$HOME"/.local/nerd-dictation/nerd-dictation "$HOME"/.local/bin

if [ $? = 0 ]; then
    printwc "${Green}" "Completed script!\n"
    printwc "${Blue}"  "Usage:\n"
    printwc "${Purple}" "    nerd-dictation begin --vosk-model-dir=./model &\n"
    printwc "${Cyan}"   "    # Start speaking.\n"
    printwc "${Purple}" "    nerd-dictation end\n"
    su - $USER
else
    printwc "${Red}" "Something went wrong please check the previous output.\n"
fi