#!/usr/bin/env bash
# Capture all windows and output JSON mapping window metadata to screenshot paths

set -euo pipefail

# Check required commands
for cmd in niri jq logger; do
    if ! command -v "$cmd" >/dev/null 2>&1; then
        echo "Error: Required command not found: $cmd" >&2
        exit 1
    fi
done

LOG_TAG="niri-capture"
OUTPUT_DIR="${1:-/tmp/niri-window-captures}"
SCREENSHOT_DIR="$HOME/Pictures/Screenshots"

mkdir -p "$OUTPUT_DIR"
logger -t "$LOG_TAG" "Starting capture of all windows to $OUTPUT_DIR"

# Ensure screenshot directory exists (niri writes here first)
if [[ ! -d "$SCREENSHOT_DIR" ]]; then
    mkdir -p "$SCREENSHOT_DIR"
fi

echo "Capturing all windows to $OUTPUT_DIR..." >&2

# Get all window IDs
if ! WINDOWS_JSON=$(niri msg --json windows 2>&1); then
    logger -t "$LOG_TAG" "ERROR: Failed to query niri: $WINDOWS_JSON"
    echo "Error: Failed to query niri. Is niri running?" >&2
    exit 1
fi

WINDOW_IDS=$(echo "$WINDOWS_JSON" | jq -r '.[].id')
WINDOW_COUNT=$(echo "$WINDOW_IDS" | wc -w)

if [[ -z "$WINDOW_IDS" ]]; then
    logger -t "$LOG_TAG" "ERROR: No windows found"
    echo "Error: No windows found" >&2
    exit 1
fi

logger -t "$LOG_TAG" "Found $WINDOW_COUNT windows to capture"

# Capture each window
FAILED=0
for id in $WINDOW_IDS; do
    # Get window metadata from cached query
    METADATA=$(echo "$WINDOWS_JSON" | jq --arg id "$id" '.[] | select(.id == ($id | tonumber))')
    TITLE=$(echo "$METADATA" | jq -r '.title')
    APP_ID=$(echo "$METADATA" | jq -r '.app_id')
    WORKSPACE=$(echo "$METADATA" | jq -r '.workspace_id')

    # Sanitize title for filename
    SAFE_TITLE=$(echo "$TITLE" | tr '/' '-' | tr ' ' '_' | cut -c1-50)
    OUTPUT_PATH="$OUTPUT_DIR/window-${id}-${SAFE_TITLE}.png"

    # Capture window
    if ! niri msg action screenshot-window --id "$id" --write-to-disk true >/dev/null 2>&1; then
        logger -t "$LOG_TAG" "WARNING: Failed to capture window $id: $TITLE"
        echo "Warning: Failed to capture window $id: $TITLE" >&2
        FAILED=$((FAILED + 1))
        continue
    fi
    sleep 0.1

    # Move from Screenshots to our output dir
    if ! LATEST=$(ls -t "$SCREENSHOT_DIR"/*.png 2>/dev/null | head -1) || [[ -z "$LATEST" ]]; then
        logger -t "$LOG_TAG" "WARNING: Screenshot not found for window $id"
        echo "Warning: Screenshot not found for window $id" >&2
        FAILED=$((FAILED + 1))
        continue
    fi
    mv "$LATEST" "$OUTPUT_PATH"
    logger -t "$LOG_TAG" "Captured window $id ($APP_ID): $OUTPUT_PATH"

    # Output JSON for this window
    echo "$METADATA" | jq --arg path "$OUTPUT_PATH" '. + {screenshot_path: $path}'
done | jq -s '.'

if [[ $FAILED -gt 0 ]]; then
    logger -t "$LOG_TAG" "Completed with $FAILED failed captures out of $WINDOW_COUNT"
    echo "Completed with $FAILED failed captures" >&2
else
    logger -t "$LOG_TAG" "Successfully captured all $WINDOW_COUNT windows"
fi
