#!/usr/bin/env bash

# ***This script was made by Clay Gomera (Drake)***
#  - Description: A simple screenshot menu script for rofi/dmenu/wofi
#  - Dependencies: {rofi||dmenu||wofi}, wayshot, slurp, wf-recorder

###########################
## Screenshots Directory ##
###########################
SHOTDIR="$MEDIA_DIR/Pictures/Screenshots" # $MEDIA_DIR reffers to a global variable in .xinitrc/.bash_profile
VCAPDIR="$MEDIA_DIR/Videos/Screencasts" # $MEDIA_DIR reffers to a global variable in .xinitrc/.bash_profile
[ ! -d "$SHOTDIR" ] && [ ! -d "$VCAPDIR" ] && mkdir -p "$SHOTDIR" "$VCAPDIR" || echo

#####
## This variable will store the
## available audio input devices
#####
devices=$(pactl list sources | grep "Name" | awk '{print $2}')

#######################
## Main menu choices ##
#######################
mcho1="󰹑  Take an screenshot"
mcho2="󰻃  Record the screen"
mcho3="  Exit"
mchos="$mcho1\n$mcho2\n$mcho3"

#############################
## Screenshot menu choices ##
#############################
scho1="󱣴  Entire screen"
scho2="󱎫  Entire screen with delay"
scho3="󱕻  Select region"
scho4="󰖯  Active window"
scho5="󱎘  Exit"
schos="$scho1\n$scho2\n$scho3\n$scho4\n$scho5"

#############################
## Screenshot menu choices ##
#############################
vcho1="󰕧  Record the screen"
vcho2="󰕩  Record region"
vcho3="󰕧  Record the screen and audio"
vcho4="󱎘  Exit"
vchos="$vcho1\n$vcho2\n$vcho3\n$vcho4"

########################
## Screenshot submenu ##
########################
sscho1="󰆏  Copy to clipboard"
sscho2="󰠘  Save to ~/Media/Pictures/Screenshots"
sscho3="󱎘  Exit"
sschos="$sscho1\n$sscho2\n$sscho3"

#################################
## Screenshot delay subsubmenu ##
#################################
del1="󱑀  3 sec delay"
del2="󱑂  5 sec delay"
del3="󱑇  10 sec delay"
del4="󱎘  Exit"
dels="$del1\n$del2\n$del3\n$del4"

#####
## This function uses the sschos variable
## to ask the user what to do with the
## screenshot
#####
## param: none
## return: string
#####
fsschos() {
    sschoice=$(echo -e "$sschos" | $RUNNER -i -L 4 -p "  What do you want to do with this screenshot?")
}

#####
## This function uses the dels variable
## to ask the user which delay option to
## choose
#####
## param: none
## return: string
#####
fdel() {
    del=$(echo -e "$dels" | $RUNNER -L 5 -i -p " Select Delay")
}

# countdown function
countdown() {
  notify-send "Screenshot" "Executing in 3 seconds" -t 1000
  sleep 1
  notify-send "Screenshot" "Executing in 2 seconds" -t 1000
  sleep 1
  notify-send "Screenshot" "Executing in 1 seconds" -t 1000
  sleep 2
}

#####
## This function does a full screen
## screenshot without delay, depending on
## what the user chooses on the fsschos
## function, the screenshot will be saved
## to the clipboard or to $SHOTDIR
#####
## param: none
## return: void
#####
shot_screen() {
    fsschos
    if [ "$sschoice" = "$sscho1" ]; then
        sleep 0.5 && wayshot --stdout | wl-copy && notify-send "Screenshot copied to clipboard"
    elif [ "$sschoice" = "$sscho2" ]; then
        sleep 0.5 && wayshot -f "$SHOTDIR/$(date +%s).jpg" && notify-send "Screenshot saved to $SHOTDIR"
    else
        exit 0
    fi
}

#####
## This function does a full screen
## screenshot with delay, depending on
## what the user chooses on the fsschos
## function, the screenshot will be saved
## to the clipboard or to $SHOTDIR. And
## depending on what the user chooses on
## the fdel function, the delay will be
## between 3 and 10 seconds
#####
## param: none
## return: void
#####
shot_screen_delay() {
    fsschos;
    if [ "$sschoice" = "$sscho1" ]; then
        fdel;
        case $del in
            "$del1")
                sleep 3 && wayshot --stdout | wl-copy && notify-send "Screenshot saved to clipboard";
                ;;
            "$del2")
                sleep 5 && wayshot --stdout | wl-copy && notify-send "Screenshot saved to clipboard";
                ;;
            "$del3")
                sleep 10 && wayshot --stdout | wl-copy && notify-send "Screenshot saved to clipboard";
                ;;
            "$del4")
                exit 0;
        esac
    elif [ "$sschoice" = "$sscho2" ]; then
        fdel;
        case $del in
            "$del1")
                sleep 3 && wayshot -f "$SHOTDIR/$(date +%s)" && notify-send "Screenshot saved to $SHOTDIR";
                ;;
            "$del2")
                sleep 5 && wayshot -f "$SHOTDIR/$(date +%s)" && notify-send "Screenshot saved to $SHOTDIR";
                ;;
            "$del3")
                sleep 10 && wayshot -f "$SHOTDIR/$(date +%s)" && notify-send "Screenshot saved to $SHOTDIR";
                ;;
            "$del4")
                exit 0;
        esac
    else
        exit 0
    fi
}


#####
## This function allows the user to select
## the area on screen to screenshot
## depending on what the user chooses on
## the fsschos function, the screenshot
## will be saved to the clipboard or to
## $SHOTDIR
#####
## param: none
## return: void
#####
shot_area() {
    fsschos;
    if [ "$sschoice" = "$sscho1" ]; then
        sleep 0.5 && wayshot -s "$(slurp)" --stdout | wl-copy && notify-send "Screenshot saved to clipboard";
    elif [ "$sschoice" = "$sscho2" ]; then
        sleep 0.5 && wayshot -f "$SHOTDIR/$(date +%s).jpg" && notify-send "Screenshot saved to $SHOTDIR";
    else
        exit 0
    fi
}

#####
## This function does an screnshot of the
## currently active window, depending on
## what the user chooses on the fsschos
## function, the screenshot will be saved
## to the clipboard or to $SHOTDIR
#####
## param: none
## return: void
#####
shot_window() {
    fsschos;
    if [ "$sschoice" = "$sscho1" ]; then
        wayshot -s "$(hyprctl activewindow -j | jq -c -r '.at,.size | .[]' | tr '\n' ' ')" --stdout | wl-copy && notify-send "Screenshot saved to clipboard";
    elif [ "$sschoice" = "$sscho2" ]; then
        wayshot -s "$(hyprctl activewindow -j | jq -c -r '.at,.size | .[]' | tr '\n' ' ')" --stdout -f "$SHOTDIR/$(date +%s).jpg" && notify-send "Screenshot saved to $SHOTDIR";
    else
        exit 0
    fi
}

#####
## This function starts recording the whole screen,
## with mp4 format and no audio. The recording will
## continue until the user stops it with the -s
## option, when that happens, the video will be
## saved to $VCAPDIR.
#####
## param: none
## return: void
#####
vshot_screen() {
    for ((i = 3; i >= 1; i--)); do
        notify-send -t 1000 "Starting in $i seconds";
        sleep 1;
    done
    wf-recorder --codec=h264_vaapi -d /dev/dri/renderD128 -f "$VCAPDIR/$(date +%s).mp4"
    notify-send "Screenshot" "Recording saved to $VCAPDIR"
}

#####
## This function will execute slurp to get the area,
## to record, and then starts the recording with mp4
## format and no audio. The recording will continue
## until the user stops it with the -s option, when
## that happens, the video will be saved to $VCAPDIR.
#####
## param: none
## return: void
#####
vshot_area() {
    wf-recorder --codec=h264_vaapi -d /dev/dri/renderD128 -g "$(slurp)" -f "$VCAPDIR/$(date +%s).mp4"
    notify-send "Screenshot" "Recording saved to $VCAPDIR"
}

#####
## This function starts recording the whole screen,
## with mp4 format and with audio. The recording will
## continue until the user stops it with the -s
## option, when that happens, the video will be
## saved to $VCAPDIR.
#####
## param: none
## return: void
#####
vshot_screen_audio() {
    chosendevice=$(echo -e "$devices" | $RUNNER -i -L 10 -p "Select audio input:")
    if [ -n "$chosendevice" ]; then
        device="$chosendevice"
        for ((i = 3; i >= 1; i--)); do
            notify-send -t 1000 "Starting in $i seconds";
            sleep 1;
        done
        wf-recorder --audio="$device" --codec=h264_vaapi -d /dev/dri/renderD128 -f "$VCAPDIR/$(date +%s).mp4"
    else
        exit 1
    fi
    notify-send "Screenshot" "Recording saved to $VCAPDIR"
}

# show the help output with --help or -h arguments 
if [ "$1" == '--help' ] || [ "$1" = '-h' ]; then
    echo "rs-scrot"
    echo "USAGE: rs-scrot [OPTION]"
    echo -e "(no option)\tshow the screenshot menu"
    echo -e "-s, --stop\tstop recording"
    echo -e "-h, --help\tthis screen"
    exit 1
fi

# stop recording with -s or --stop arguments
if [ "$1" = '--stop' ] || [ "$1" = '-s' ]; then
    killall -s SIGINT wf-recorder
    exit 1
fi

##########
## main ##
##########
mchoice=$(echo -e "$mchos" | $RUNNER -i -L 5 -p "  Screenshot/Screen-Capture Utility") # main menu prompt
case $mchoice in
    "$mcho1")
        schoice=$(echo -e "$schos" | $RUNNER -i -L 7 -p "  Sreenshot Menu") # screenshot menu prompt
        case $schoice in
            "$scho1")
                shot_screen;
                ;;
            "$scho2")
                shot_screen_delay;
                ;;
            "$scho3")
                shot_area;
                ;;
            "$scho4")
                shot_window;
                ;;
            "$scho5")
                exit 0
        esac
        ;;
    "$mcho2")
        vchoice=$(echo -e "$vchos" | $RUNNER -i -L 6 -p "󰕧 Screen-Capture Menu")
        case $vchoice in
            "$vcho1")
                vshot_screen;
                ;;
            "$vcho2")
                vshot_area;
                ;;
            "$vcho3")
                vshot_screen_audio;
                ;;
            "$vcho4")
                exit 0;
        esac
        ;;
    "$mcho3")
        exit 0
        ;;
esac