grabit/main.sh

#!/bin/bash

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

source "$SCRIPT_DIR/helpers/variables.sh"
source "$SCRIPT_DIR/helpers/installer.sh"
source "$SCRIPT_DIR/helpers/config.sh"

validate_config

if [[ "$NEEDS_SETUP" -eq 1 ]]; then setup_config; fi
if [[ $(get_value "SHOW_CAPTURE_PREVIEW") == "true" ]]; then setup_venv; fi

arguments=()
while [[ $# -gt 0 ]]; do
    arg="$1"
    case "$arg" in
        --help)
            echo "Usage: $0 [options]"
            echo "--help       Show this help message"
            echo "--version    Show version information"
            echo "--config     Allows you to set up the config file"
            echo "-c           Copies the screenshot to clipboard"
            echo "-u           Uploads the screenshot"
            echo "-f <file>    Uploads the given file instead of taking a screenshot"
            echo "-d           Debug mode"
            echo "--silent     Suppresses Image preview and notifications"
            echo
            echo "Domain Specific Options:"
            for service in "${valid_services[@]}"; do
                echo "--$service"
            done
            exit 0
            ;;
        --version)
            echo "Version: $VERSION"
            exit 0
            ;;
        --config)
            setup_config
            exit 0
            ;;
        -c)
            arguments+=("-c")
            shift
            ;;
        -u)
            arguments+=("-u")
            shift
            ;;
        -f)
            shift
            if [[ -n "$1" && "$1" != -* ]]; then
                arguments+=("-f" "$1")
                shift
            else
                echo -e "${RED}Error: No valid file specified for -f option.${RESET}"
                exit 1
            fi
            ;;
        -d)
            DEBUG=true
            shift
            ;;
        --silent)
            SHOW_NOTIFICATIONS=false
            SHOW_CAPTURE_PREVIEW=false
            shift
            ;;
        --*)
            service="${arg:2}"
            if [[ " ${valid_services[*]} " =~ " ${service} " ]]; then
                SHOULD_UPLOAD=true
                SHOULD_COPY=false
                SERVICE="${service}"
                shift
            else
                echo -e "${RED}Error: Invalid service: $service${RESET}"
                exit 1
            fi
            ;;
        *)
            echo -e "${RED}Error: Invalid argument: $arg${RESET}"
            echo "Use --help for usage information."
            exit 1
            ;;
    esac
done

SHOULD_UPLOAD=$(get_value "DEFAULT_OPTION")
[[ "$SHOULD_UPLOAD" == "upload" ]] && SHOULD_UPLOAD=true || SHOULD_UPLOAD=false

SHOW_NOTIFICATIONS=$(get_value "SHOW_NOTIFICATIONS")
[[ "$SHOW_NOTIFICATIONS" == "true" ]] && SHOW_NOTIFICATIONS=true || SHOW_NOTIFICATIONS=false

SHOW_CAPTURE_PREVIEW=$(get_value "SHOW_CAPTURE_PREVIEW")
[[ "$SHOW_CAPTURE_PREVIEW" == "true" ]] && SHOW_CAPTURE_PREVIEW=true || SHOW_CAPTURE_PREVIEW=false

REQUIRED_COMMANDS=("notify-send" "curl" "jq")

if [[ -n "$WAYLAND_DISPLAY" ]]; then
    if [[ "$XDG_CURRENT_DESKTOP" == "KDE" ]]; then
        SCREENSHOT_CMD=("spectacle" "-rbn" "-o")
        REQUIRED_COMMANDS+=("xclip" "spectacle")
    else
        SCREENSHOT_CMD=("grim" "-g")
        SELECT_REGION="slurp"
        REQUIRED_COMMANDS+=("wl-copy" "grim" "slurp")
    fi
elif [[ -n "$DISPLAY" ]]; then
    SCREENSHOT_CMD=("flameshot" "gui" "-p")
    REQUIRED_COMMANDS+=("xclip" "flameshot")
else
    echo "Error: Unable to detect display server (Wayland/X11)."
    exit 1
fi

MISSING_COMMANDS=()
for cmd in "${REQUIRED_COMMANDS[@]}"; do
    command -v "$cmd" &>/dev/null || MISSING_COMMANDS+=("$cmd")
done

# Check for xdg-utils for xdg-open
if ! command -v xdg-open &>/dev/null; then
    MISSING_COMMANDS+=("xdg-utils")
fi

if [[ ${#MISSING_COMMANDS[@]} -gt 0 ]]; then
    echo -e "${RED}Error: The following commands are not installed: ${MISSING_COMMANDS[*]}${RESET}"
    echo -e "${RED}Please install them"
    exit 1
fi

send_notification() {
    [[ "$SHOW_NOTIFICATIONS" == "false" ]] && return

    local title="$1"
    local message="$2"
    local image="$3"

    if [[ -z "$title" && -z "$message" ]]; then
        return
    fi

    local notify_cmd=("notify-send")

    [[ -n "$image" ]] && notify_cmd+=("-i" "$image")
    [[ -n "$title" ]] && notify_cmd+=("$title")
    [[ -n "$message" ]] && notify_cmd+=("$message")

    "${notify_cmd[@]}"
}

SHOULD_UPLOAD=false
SHOULD_SAVE="$(get_value "SAVE_IMAGES")"
SHOULD_COPY=false

case "$(get_value "DEFAULT_OPTION")" in
    upload) SHOULD_UPLOAD=true ;;
    copy)   SHOULD_COPY=true ;;
esac

for arg in "${arguments[@]}"; do
    case "$arg" in
        -c)
            SHOULD_UPLOAD=false
            SHOULD_COPY=true
            ;;
        -u|-f)
            SHOULD_UPLOAD=true
            SHOULD_COPY=false
            ;;
    esac
done

CURRENT_DATE=$(date +%Y-%m-%d-%H-%M-%S)
RANDOM_FILE_NAME="${CURRENT_DATE}_$(date +%s%N | cut -b1-13)"

if [[ "$SHOULD_COPY" == "true" ]]; then    
    if [[ "$SHOULD_SAVE" == "true" ]]; then
        SAVE_DIR=$(get_value "SAVE_DIR")
        if [[ ! -d "$SAVE_DIR" ]]; then
            mkdir -p "$SAVE_DIR"
        fi
        FILE="$SAVE_DIR/$CURRENT_DATE.png"
    else
        FILE="/tmp/$RANDOM_FILE_NAME.png"
    fi

    if [[ -n "$WAYLAND_DISPLAY" ]]; then
        if [[ "$XDG_CURRENT_DESKTOP" == "KDE" && -n "$WAYLAND_DISPLAY" ]]; then
            "${SCREENSHOT_CMD[@]}" "$FILE"
        else
            "${SCREENSHOT_CMD[@]}" "$("$SELECT_REGION")" "$FILE"
        fi
    else
        "${SCREENSHOT_CMD[@]}" "$FILE"
    fi

    if [[ $? -ne 0 ]]; then
        echo "Error: Screenshot capture failed." >&2
        exit 1
    fi

    if [[ ! -f "$FILE" ]]; then
        echo "Error: File not found: $FILE" >&2
        exit 1
    fi

    if [[ -n "$WAYLAND_DISPLAY" ]]; then
        echo "Using wl-copy for clipboard."
        cat "$FILE" | wl-copy
    elif [[ -n "$DISPLAY" ]]; then
        echo "Using xclip for clipboard."
        xclip -selection clipboard -t image/png -i "$FILE"
    else
        echo "Error: Unable to detect display server (Wayland/X11)."
        exit 1
    fi

    if [[ "$SHOW_CAPTURE_PREVIEW" == "true" ]]; then
        "$PYTHON_EXEC" "$SCRIPT_DIR/helpers/show_image.py" "$FILE" "Copied to your clipboard"
    fi

    if [[ "$SHOULD_SAVE" != "true" ]]; then
        [[ -n "$FILE" ]] && rm -f "$FILE"
    fi

    exit 0
fi

if [[ -z "$SERVICE" ]]; then
    SERVICE=$(get_value "SERVICE")
    if [[ -z "$SERVICE" ]]; then
        echo -e "${RED}Error: No service specified and no default service set.${RESET}"
        exit 1
    fi
fi

if [[ "$SERVICE" == "zipline" ]]; then
    DOMAIN=$(get_value "DOMAIN")
    if [[ -z "$DOMAIN" ]]; then
        echo -e "${RED}Error: Domain not set for zipline.${RESET}"
        setup_domain
    fi
else
    DOMAIN="${UPLOADERS[$SERVICE]}"
fi

AUTH=$(get_value "${SERVICE}_auth")
if [[ -z "$AUTH" ]]; then
    echo -e "${RED}Error: Auth not set for $SERVICE.${RESET}"
    setup_domain

    AUTH=$(get_value "${SERVICE}_auth")
    if [[ -z "$AUTH" ]]; then
        echo -e "${RED}Error: No auth key provided for $SERVICE.${RESET}"
        exit 1
    fi
fi

USER_FILE=""
if [[ "$SHOULD_UPLOAD" == "true" ]]; then
    for ((i = 0; i < ${#arguments[@]}; i++)); do
        if [[ "${arguments[i]}" == "-f" ]]; then
            USER_FILE="${arguments[i+1]}"
            break
        fi
    done
fi

if [[ -z "$USER_FILE" ]]; then
    if [[ "$SHOULD_SAVE" == "true" ]]; then
        SAVE_DIR=$(get_value "SAVE_DIR")
        if [[ ! -d "$SAVE_DIR" ]]; then
            mkdir -p "$SAVE_DIR"
        fi
        FILE="$SAVE_DIR/$CURRENT_DATE.png"
    else
        FILE="/tmp/$RANDOM_FILE_NAME.png"
    fi

    if [[ -n "$WAYLAND_DISPLAY" ]]; then
        if [[ "$XDG_CURRENT_DESKTOP" == "KDE" && -n "$WAYLAND_DISPLAY" ]]; then
            "${SCREENSHOT_CMD[@]}" "$FILE"
        else
            "${SCREENSHOT_CMD[@]}" "$("$SELECT_REGION")" "$FILE"
        fi
    else
        "${SCREENSHOT_CMD[@]}" "$FILE"
    fi

    if [[ $? -ne 0 ]]; then
        echo "Error: Screenshot capture failed." >&2
        exit 1
    fi

    if [[ ! -f "$FILE" ]]; then
        echo "Error: File not found: $FILE" >&2
        exit 1
    fi
else
    FILE="$USER_FILE"
fi

if [[ ! -f "$FILE" ]]; then
    echo "Error: File not found: $FILE" >&2
    exit 1
fi

HEADERS=()
if [[ "$SERVICE" == "zipline" ]]; then
    declare -A ZIPLINE_HEADERS
    ZIPLINE_HEADERS=$(get_starting_with "x-zipline")
    for key in "${!ZIPLINE_HEADERS[@]}"; do
        header_name="${key#ZIPLINE__}" | tr '[:upper:]' '[:lower:]' | tr '_' '-'
        header_value="${ZIPLINE_HEADERS[$key]}"
        headers+=("-H" "$header_name: $header_value")
    done
fi

headers+=("-H" "${UPLOAD_HEADERS[$SERVICE]}: $AUTH")

echo headers "${headers[@]}"

FILE_URL=""
RESPONSE=$(curl -s -w "%{http_code}" -o /tmp/upload_response.json \
    -X POST -F "file=@$FILE" "${headers[@]}" "$DOMAIN")

HTTP_CODE="${RESPONSE:(-3)}"
BODY=$(cat /tmp/upload_response.json)
rm /tmp/upload_response.json

if [[ "$HTTP_CODE" -ne 200 ]]; then
    echo -e "${RED}Error: Upload failed with status code $HTTP_CODE.${RESET}"
    echo "$BODY"
    exit 1
fi

if ! echo "$BODY" | jq . &>/dev/null; then
    echo -e "${RED}Error: Invalid JSON response from server.${RESET}"
    echo "$BODY"
    exit 1
fi

FILE_URL=$(echo "$BODY" | jq -r ".${UPLOAD_JSON_KEYS[$SERVICE]}")
if [[ "$FILE_URL" == "null" || -z "$FILE_URL" ]]; then
    echo -e "${RED}Error: File URL not found in response.${RESET}"
    echo "$BODY"
    exit 1
fi

if [[ -n "$WAYLAND_DISPLAY" ]]; then
    echo "Using wl-copy for clipboard."
    echo "$FILE_URL" | wl-copy
elif [[ -n "$DISPLAY" ]]; then
    echo "Using xclip for clipboard."
    echo "$FILE_URL" | xclip -selection clipboard
else
    echo "Error: Unable to detect display server (Wayland/X11)."
    exit 1
fi

MIME_TYPE=$(file --mime-type -b "$FILE")
if [[ "$MIME_TYPE" == image/* ]]; then
    if [[ "$SHOW_NOTIFICATIONS" == "true" ]]; then
        send_notification "Uploaded to $SERVICE" "$FILE_URL" "$FILE"
    fi

    if [[ "$SHOW_CAPTURE_PREVIEW" == "true" ]]; then
        "$PYTHON_EXEC" "$SCRIPT_DIR/helpers/show_image.py" "$FILE" "Uploaded to $SERVICE"
    fi
else
    if [[ "$SHOW_NOTIFICATIONS" == "true" ]]; then
        send_notification "Uploaded to $SERVICE" "$FILE_URL"
    fi
fi

# delete if should not save and not a user file
if [[ "$SHOULD_SAVE" != "true" && "$USER_FILE" == "" ]]; then
    rm -f "$FILE"
fi