#!/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)..."
    
    # Kill all sessions in the group or containing the session name
    tmux list-sessions -F '#{session_name} #{session_group}' 2>/dev/null | while read -r s g; do
        if [ "$g" = "$SESSION_NAME" ] || [[ "$s" == *"$SESSION_NAME"* ]]; then
            tmux kill-session -t "$s" 2>/dev/null
        fi
    done
    
    # Kill all processes with CWD under this directory (leaked subagents)
    # Only kill language_server (jetski) and server.par (codemind) to avoid killing ourselves!
    find /proc -maxdepth 2 -name cwd -exec ls -l {} + 2>/dev/null | grep "$SCRIPT_DIR" | while read -r line; do
        pid=$(echo "$line" | awk -F/ '{print $3}')
        cmd=$(ps -p "$pid" -o comm= 2>/dev/null)
        if [[ "$cmd" == *"language_server"* ]] || [[ "$cmd" == *"server.par"* ]]; then
            kill -9 "$pid" 2>/dev/null
        fi
    done
    
    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 Physical Dispatcher window
tmux new-window -t $SESSION_NAME -n dispatcher -c "$SCRIPT_DIR"
tmux send-keys -t $SESSION_NAME:dispatcher "python3 -u $ORCHESTRATOR_PATH > orchestrator.log 2>&1" C-m

# Create Grid window for all agents (Orchestrator + Sub-agents)
mkdir -p "$SCRIPT_DIR/agents/orchestrator"
tmux new-window -t $SESSION_NAME -n grid -c "$SCRIPT_DIR/agents/orchestrator"

# Enable pane borders and titles
tmux set-window-option -t $SESSION_NAME:grid pane-border-status top
tmux set-window-option -t $SESSION_NAME:grid pane-border-format " #P: #{pane_title} "

# Pane 0: Orchestrator Brain
tmux select-pane -t $SESSION_NAME:grid.0 -T "Orchestrator Brain"
tmux send-keys -t $SESSION_NAME:grid.0 "$CLI_PATH -cli=true" C-m

# Add sub-agents to Grid window as panes
for i in $(seq 1 $N); do
    mkdir -p "$SCRIPT_DIR/agents/agent_$i"
    tmux split-window -t $SESSION_NAME:grid -c "$SCRIPT_DIR/agents/agent_$i"
    tmux select-layout -t $SESSION_NAME:grid tiled
    tmux select-pane -t $SESSION_NAME:grid.$i -T "Sub Agent $i"
    tmux send-keys -t $SESSION_NAME:grid.$i "$CLI_PATH -cli=true" C-m
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, it is Pane 0 in the Grid window."
echo "Or check the log file: $SCRIPT_DIR/orchestrator.log"

echo ""
echo "To view Master and Grid independently in two terminals:"
echo "  Run this in the second terminal to open the grid automatically:"
echo "  tmux new-session -t $SESSION_NAME -s ${SESSION_NAME}_view2 \\; select-window -t grid"
echo ""
echo "To view all windows in one linked session:"
echo "  tmux attach -t $SESSION_NAME"
echo ""
echo "To stop the swarm and cleanup:"
echo "  tmux kill-session -t $SESSION_NAME"