#!/bin/bash

# i3-screenshot-countdown.sh - Tool for taking screenshots with a countdown
# This script provides a countdown before taking the screenshot to avoid capturing menus

# Default save location - user's Pictures directory
SAVE_DIR="$HOME/Pictures/Screenshots"
TIMESTAMP=$(date +%Y-%m-%d_%H-%M-%S)
SCREENSHOT_PATH="$SAVE_DIR/screenshot_$TIMESTAMP.png"

# Create the save directory if it doesn't exist
mkdir -p "$SAVE_DIR"

# Function to display notification
notify() {
    notify-send "Screenshot Tool" "$1" -i camera-photo
}

# Function to get a list of monitors
get_monitors() {
    xrandr --listactivemonitors | grep -v "Monitors" | awk '{print $4}'
}

# Function to get monitor geometry
get_monitor_geometry() {
    local monitor=$1
    xrandr | grep "$monitor" | grep -oP '\d+x\d+\+\d+\+\d+'
}

# Function to do a countdown before taking the screenshot
countdown_screenshot() {
    # Countdown from 3
    for i in 3 2 1; do
        notify "Taking screenshot in $i..."
        sleep 1
    done
    
    # Take the screenshot
    if [ -n "$1" ]; then
        maim -g "$1" "$SCREENSHOT_PATH"
    else
        maim "$SCREENSHOT_PATH"
    fi
    
    # Check if screenshot was successful
    if [ $? -ne 0 ]; then
        notify "Screenshot failed"
        exit 1
    fi
}

# Check if rofi is installed
if ! command -v rofi &> /dev/null; then
    notify "Rofi is not installed. Using direct selection mode."
    
    # Use slop to select an area and maim to take the screenshot
    notify "Select an area to capture"
    SELECTION=$(slop -f "%x,%y,%w,%h")
    if [ $? -ne 0 ]; then
        notify "Screenshot canceled"
        exit 1
    fi
    
    # Extract coordinates and dimensions
    IFS=',' read -r X Y W H <<< "$SELECTION"
    
    # Take the screenshot of the selected area with countdown
    countdown_screenshot "${W}x${H}+${X}+${Y}"
else
    # Get list of monitors
    MONITORS=$(get_monitors)
    
    # Add "Select Area" and "Full Screenshot" options
    OPTIONS="Select Area\nFull Screenshot"
    
    # Add each monitor to the options
    for MONITOR in $MONITORS; do
        OPTIONS="$OPTIONS\n$MONITOR"
    done
    
    # Show selection menu with rofi
    CHOICE=$(echo -e "$OPTIONS" | rofi -dmenu -i -p "Screenshot Mode")
    
    case "$CHOICE" in
        "Select Area")
            # Use slop to select an area
            notify "Select an area to capture"
            SELECTION=$(slop -f "%x,%y,%w,%h")
            if [ $? -ne 0 ]; then
                notify "Screenshot canceled"
                exit 1
            fi
            
            # Extract coordinates and dimensions
            IFS=',' read -r X Y W H <<< "$SELECTION"
            
            # Take the screenshot of the selected area with countdown
            countdown_screenshot "${W}x${H}+${X}+${Y}"
            ;;
            
        "Full Screenshot")
            # Take screenshot of all monitors with countdown
            countdown_screenshot
            ;;
            
        *)
            # Take screenshot of selected monitor
            if [ -n "$CHOICE" ]; then
                GEOMETRY=$(get_monitor_geometry "$CHOICE")
                if [ -n "$GEOMETRY" ]; then
                    notify "Taking screenshot of monitor $CHOICE..."
                    countdown_screenshot "$GEOMETRY"
                else
                    notify "Could not determine monitor geometry"
                    exit 1
                fi
            else
                notify "Screenshot canceled"
                exit 1
            fi
            ;;
    esac
fi

# Copy to clipboard
xclip -selection clipboard -t image/png -i "$SCREENSHOT_PATH"

# Notify user
notify "Screenshot saved to $SCREENSHOT_PATH and copied to clipboard"

echo "Screenshot saved to: $SCREENSHOT_PATH"