#!/usr/bin/env bash

# ***This script was made by Clay Gomera (Drake)***
#  - Description: A simple screenshot rofi script
#  - Dependencies: main, slop, ffmpeg, xdotool, xdpyinfo, rofi, notify-send

## Screenshot & Screencast Directory
SHOTDIR="$XDG_PICTURES_DIR/Screenshots"
CASTDIR="$XDG_VIDEOS_DIR/Screencasts"
mkdir -p "$SHOTDIR" "$CASTDIR"

## Format Variables
image_formats="jpg\npng"
video_formats="mp4\nmkv"

## Main menu choices
mcho1="󰹑 Take a screenshot"
mcho2="󰐍 Record the screen"
mcho3="󰙧 Stop recording"
mcho4="󱎘 Exit"
mchos="$mcho1\n$mcho2\n$mcho3\n$mcho4"

## Main submenu choices
mscho1="󱣴 Capture the entire screen"
mscho2="󱎫 Capture the entire screen (With delay)"
mscho3="󱕻 Capture area"
mscho4="󰖯 Capture active window"
mscho5="󱎘 Exit"
mschos="$mscho1\n$mscho2\n$mscho3\n$mscho4\n$mscho5"

## Save/clip submenu choices
svcho1="󰆏 Copy to clipboard"
svcho2="󰠘 Save to $SHOTDIR"
svchos="$svcho1\n$svcho2"

## Delay subsubmenu
del1="󱑀 3 sec delay"
del2="󱑂 5 sec delay"
del3="󱑇 10 sec delay"
dels="$del1\n$del2\n$del3"

#####
## This function will check for dependencies
#####
check_dependencies() {
    local dependencies=("maim" "ffmpeg" "rofi" "notify-send" "xclip" "xdotool" "xdpyinfo")
    for dep in "${dependencies[@]}"; do
        if ! command -v "$dep" &>/dev/null; then
            echo "$dep is required but not installed. Please install it and try again."
            exit 1
        fi
    done
}

#####
## This function uses the svchos variable to ask the user what to do with the
## screenshot
#####
fsvchos() {
    svchoice=$(echo -e "$svchos" | rofi -dmenu -i -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
#####
fdel() {
    del=$(echo -e "$dels" | rofi -dmenu -i -p " Select Delay")
}

#####
## This function allows the user to choose the screenshot format
#####
fimage_format() {
    image_format=$(echo -e "$image_formats" | rofi -dmenu -i -p " Choose screenshot format")
}

#####
## This function allows the user to choose the recording format
#####
fvideo_format() {
    video_format=$(echo -e "$video_formats" | rofi -dmenu -i -p " Choose video format")
}

#####
## This function does a full screen screenshot without delay, depending on what
## the user chooses on the fsvchos function, the screenshot will be saved to the
## clipboard or to $SHOTDIR
#####
shot_screen() {
    fimage_format
    fsvchos
    if [ "$svchoice" = "$svcho1" ]; then
        maim -f "$image_format" | xclip -selection clipboard -t image/"$image_format" && notify-send "Screenshot saved to clipboard"
    elif [ "$svchoice" = "$svcho2" ]; then
        maim -f "$image_format" "$SHOTDIR/$(date +%s).$image_format" && notify-send "Screenshot saved"
    else
        exit 0
    fi
}

#####
## This function does a full screen screenshot with delay, depending on what the
## user chooses on the fsvchos 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
#####
shot_screen_delay() {
    fimage_format
    fsvchos
    fdel
    case $del in
        "$del1")
            for i in 3 2 1; do
                notify-send -t 1000 "Capturing in $i.."
                sleep 1
            done
            ;;
        "$del2")
            for i in 5 4 3 2 1; do
                notify-send -t 1000 "Capturing in $i.."
                sleep 1
            done
            ;;
        "$del3")
            for i in 10 9 8 7 6 5 4 3 2 1; do
                notify-send -t 1000 "Capturing in $i.."
                sleep 1
            done
            ;;
    esac

    sleep 1
    if [ "$svchoice" = "$svcho1" ]; then
        maim -f "$image_format" | xclip -selection clipboard -t image/"$image_format" && notify-send "Screenshot saved to clipboard";
    elif [ "$svchoice" = "$svcho2" ]; then
        maim -f "$image_format" "$SHOTDIR/$(date +%s).$image_format" && notify-send "Screenshot saved to $SHOTDIR";
    else
        exit 0
    fi
}

#####
## This function allows the user to select the area on the screen to screenshot
## depending on what the user chooses on the fsvchos function, the screenshot will
## be saved to the clipboard or to $SHOTDIR
#####
shot_area() {
    fimage_format
    fsvchos
    if [ "$svchoice" = "$svcho1" ]; then
        maim -f "$image_format" -s | xclip -selection clipboard -t image/"$image_format" && notify-send "Screenshot saved to clipboard";
    elif [ "$svchoice" = "$svcho2" ]; then
        maim -s -f "$image_format" "$SHOTDIR/$(date +%s).$image_format" && notify-send "Screenshot saved to $SHOTDIR";
    else
        exit 0
    fi
}

#####
## This function does a screenshot of the currently active window, depending
## on what the user chooses on the fsvchos function, the screenshot will be
## saved to the clipboard or to $SHOTDIR
#####
shot_window() {
    fimage_format
    fsvchos
    if [ "$svchoice" = "$svcho1" ]; then
        maim -i "$(xdotool getactivewindow)" | xclip -selection clipboard -t image/"$image_format" && notify-send "Screenshot saved to clipboard";
    elif [ "$svchoice" = "$svcho2" ]; then
        maim -i "$(xdotool getactivewindow)" -f "$image_format" "$SHOTDIR/$(date +%s).$image_format" && notify-send "Screenshot saved to $SHOTDIR";
    else
        exit 0
    fi
}

#####
## This function uses ffmpeg to record the screen, It will check if ffmpeg is
## already running (already recording), to prevent the user from executing
## another ffmpeg instance, if it's not running, it will grab the screen
## resolution with xdpyinfo and use it to record the screen with ffmpeg.
#####
rec_screen() {
    fvideo_format
    if pgrep -x "ffmpeg" > /dev/null; then
        notify-send "A screen recording is already in progress."
    else
        filename="$CASTDIR/$(date +%s).$video_format"
        resolution="$(xdpyinfo | grep dimensions | awk '{print $2}')"
        ffmpeg -f x11grab -s "$resolution" -i :0.0 -c:v libx264 -preset ultrafast -qp 0 "$filename" &
        notify-send "Screen recording started."
    fi
}

#####
## This function will record the screen with a given delay
#####
rec_screen_delay() {
    fvideo_format
    fdel
    case $del in
        "$del1")
            for i in 3 2 1; do
              notify-send -t 1000 "Starting in $i.."
              sleep 1
            done
            ;;
        "$del2")
            for i in 5 4 3 2 1; do
              notify-send -t 1000 "Starting in $i.."
              sleep 1
            done
            ;;
        "$del3")
            for i in 10 9 8 7 6 5 4 3 2 1; do
                notify-send -t 1000 "Starting in $i.."
                sleep 1
            done
            ;;
    esac

    sleep 1
    if pgrep -x "ffmpeg" > /dev/null; then
        notify-send "A screen recording is already in progress."
    else 
        filename="$CASTDIR/$(date +%s).$video_format"
        resolution="$(xdpyinfo | grep dimensions | awk '{print $2}')"
        ffmpeg -f x11grab -s "$resolution" -i :0.0 -c:v libx264 -preset ultrafast -qp 0 "$filename" &
        notify-send "Screen recording started"
    fi
}

## This function allows the user to record a selected area of the screen
rec_area() {
    fvideo_format
    if pgrep -x "ffmpeg" > /dev/null; then
        notify-send "A screen recording is already in progress."
    else 
        filename="$CASTDIR/$(date +%s).$video_format"
        slop_result=$(slop -f "%x %y %w %h")
        if [ -n "$slop_result" ]; then
            IFS=' ' read -r x y w h <<< "$slop_result"
            ffmpeg -f x11grab -s "${w}x${h}" -i :0.0+"$x","$y" -c:v libx264 -preset ultrafast -qp 0 "$filename" &
            notify-send "Screen recording started"
        else
            notify-send "No area selected. Aborting recording."
        fi
    fi
}

## This function allows the user to record the currently active window
rec_window() {
    fvideo_format
    if pgrep -x "ffmpeg" > /dev/null; then
        notify-send "A screen recording is already in progress."
    else 
        filename="$CASTDIR/$(date +%s).$video_format"
        active_window_id=$(xdotool getactivewindow)
        
        if [ -n "$active_window_id" ]; then
            ffmpeg -f x11grab -s "$(xdotool getwindowgeometry --shell $active_window_id | grep WIDTH | cut -d'=' -f2)x$(xdotool getwindowgeometry --shell $active_window_id | grep HEIGHT | cut -d'=' -f2)" -i :0.0+"$(xdotool getwindowgeometry --shell $active_window_id | grep X | cut -d'=' -f2)","$(xdotool getwindowgeometry --shell $active_window_id | grep Y | cut -d'=' -f2)" -c:v libx264 -preset ultrafast -qp 0 "$filename" &
            notify-send "Recording active window started"
        else
            notify-send "No active window found. Aborting recording."
        fi
    fi
}

#####
## This function stops screen recording, it first checks if ffmpeg is already
## running (already recording), to only stop a recording if there's an ffmpeg
## instance running
#####
stop_recording() {
    if pgrep -x "ffmpeg" > /dev/null; then
        pkill -f "ffmpeg -f x11grab"
        notify-send "Screen recording stopped. Video saved to $CASTDIR"
    else
        notify-send "You aren't recording anything at the moment"
    fi
}

## main
check_dependencies;
mchoice=$(echo -e "$mchos" | rofi -dmenu -i -p "   Screen Capture Menu ") # main menu prompt
case $mchoice in
    "$mcho1")
        mschoice=$(echo -e "$mschos" | rofi -dmenu -i -p "   Screenshot Menu ") # screenshot menu prompt
        case $mschoice in
            "$mscho1")
                shot_screen;
                ;;
            "$mscho2")
                shot_screen_delay;
                ;;
            "$mscho3")
                shot_area;
                ;;
            "$mscho4")
                shot_window;
                ;;
            "$mscho5")
                exit 0
        esac
        ;;
    "$mcho2")
        mschoice=$(echo -e "$mschos" | rofi -dmenu -i -p "   Screencast Menu ") # screenshot menu prompt
        case $mschoice in
            "$mscho1")
                rec_screen;
                ;;
            "$mscho2")
                rec_screen_delay;
                ;;
            "$mscho3")
                rec_area;
                ;;
            "$mscho4")
                rec_window;
                ;;
            "$mscho5")
                exit 0
        esac
        ;;
    "$mcho3")
        stop_recording;
        ;;
    "$mcho4")
        exit 0;
        ;;
esac