#!/usr/bin/env bash

# Ensure exactly 10 spaces exist on the main display (display 1)
# These can then be moved to other displays as needed

DESIRED_SPACES=10
MAIN_DISPLAY=1

# Get current number of spaces on main display
CURRENT_SPACE_COUNT=$(yabai -m query --spaces --display "$MAIN_DISPLAY" | jq 'length')

echo "Current spaces on display $MAIN_DISPLAY: $CURRENT_SPACE_COUNT"
echo "Desired spaces: $DESIRED_SPACES"

# Focus main display first
yabai -m display --focus "$MAIN_DISPLAY"

if [ "$CURRENT_SPACE_COUNT" -lt "$DESIRED_SPACES" ]; then
    MISSING_SPACES=$((DESIRED_SPACES - CURRENT_SPACE_COUNT))
    echo "Creating $MISSING_SPACES spaces on display $MAIN_DISPLAY..."

    for i in $(seq 1 $MISSING_SPACES); do
        yabai -m space --create
        echo "Created space $((CURRENT_SPACE_COUNT + i))"
    done
elif [ "$CURRENT_SPACE_COUNT" -gt "$DESIRED_SPACES" ]; then
    EXTRA_SPACES=$((CURRENT_SPACE_COUNT - DESIRED_SPACES))
    echo "Removing $EXTRA_SPACES extra spaces from display $MAIN_DISPLAY..."

    # Get the last space on main display and destroy it
    for i in $(seq 1 $EXTRA_SPACES); do
        LAST_SPACE=$(yabai -m query --spaces --display "$MAIN_DISPLAY" | jq 'map(select(."is-native-fullscreen" == false))[-1].index')
        yabai -m space --destroy "$LAST_SPACE"
        echo "Destroyed space $LAST_SPACE"
    done
fi

echo "Total spaces on display $MAIN_DISPLAY: $(yabai -m query --spaces --display $MAIN_DISPLAY | jq 'length')"

sketchybar --trigger space_change --trigger windows_on_spaces