#!/bin/bash

# Configuration
SESSION_NAME="jetski_swarm"
CLI_PATH="/google/bin/releases/jetski-devs/tools/cli"
SCRIPT_DIR="$( cd "$( dirname "${BASH_SOURCE[0]}" )" &> /dev/null && pwd )"
ORCHESTRATOR_PATH="$SCRIPT_DIR/orchestrator.py"

# Number of agents
N=${1:-4}

cleanup() {
    echo "Enforcing smooth start (cleaning up previous state)..."
    tmux kill-session -t $SESSION_NAME 2>/dev/null
    rm -f "$SCRIPT_DIR/comms/requests/"*.json
    echo "Cleanup complete."
}

# Smooth Start: Ensure clean state
cleanup

echo "Bootstrapping swarm in background..."

# Create session with Master Agent window
mkdir -p "$SCRIPT_DIR/agents/master"
tmux new-session -d -s $SESSION_NAME -n master -c "$SCRIPT_DIR/agents/master"
tmux send-keys -t $SESSION_NAME:master "$CLI_PATH -cli=true" C-m

# Wait a few seconds for Master Agent to initialize, then inject skill instruction
sleep 5
tmux send-keys -t $SESSION_NAME:master "Please read the file 'SKILL.md' in your current directory to understand how to control the swarm." C-m

# Create Orchestrator window
tmux new-window -t $SESSION_NAME -n orchestrator -c "$SCRIPT_DIR"
# Run orchestrator with unbuffered output and redirection
tmux send-keys -t $SESSION_NAME:orchestrator "python3 -u $ORCHESTRATOR_PATH > orchestrator.log 2>&1" C-m

# Create Grid window for sub-agents
tmux new-window -t $SESSION_NAME -n grid -c "$SCRIPT_DIR"

# Add sub-agents to Grid window as panes
for i in $(seq 1 $N); do
    mkdir -p "$SCRIPT_DIR/agents/agent_$i"
    if [ $i -eq 1 ]; then
        tmux send-keys -t $SESSION_NAME:grid "$CLI_PATH -cli=true" C-m
    else
        tmux split-window -t $SESSION_NAME:grid -c "$SCRIPT_DIR/agents/agent_$i"
        tmux send-keys -t $SESSION_NAME:grid "$CLI_PATH -cli=true" C-m
    fi
done

# Arrange Grid window in tiles
tmux select-layout -t $SESSION_NAME:grid tiled

# Wait for agents to settle and handle any prompts
echo "Aligning agents..."
python3 "$SCRIPT_DIR/../tmp/swarm/scripts/swarm_auto_aligner.py"

tmux select-window -t $SESSION_NAME:master
echo "Setup complete."
echo "To use the Master Agent (Default Window), run:"
echo "  tmux attach -t $SESSION_NAME"
echo "Or:"
echo "  tmux attach -t $SESSION_NAME:master"
echo ""
echo "To use the Grid of Sub-Agents, run:"
echo "  tmux attach -t $SESSION_NAME:grid"
echo ""
echo "To see the Orchestrator logs, attach to its window:"
echo "  tmux attach -t $SESSION_NAME:orchestrator"
echo "Or check the file: $SCRIPT_DIR/orchestrator.log"
echo ""
echo "To view all windows:"
echo "  tmux attach -t $SESSION_NAME"
echo ""
echo "To stop the swarm and cleanup:"
echo "  tmux kill-session -t $SESSION_NAME"