#!/bin/bash

# Function to display help message
function display_help {
  echo "Usage: $0 -n num_windows -c command"
  echo
  echo "   -n, --number     Number of windows to create"
  echo "   -c, --command    Command to run in each window"
  echo
  exit 1
}

# Parse command line arguments
while getopts "n:c:h" opt; do
  case ${opt} in
  n)
    num_windows=${OPTARG}
    ;;
  c)
    command=${OPTARG}
    ;;
  h)
    display_help
    ;;
  \?)
    echo "Invalid option: -$OPTARG" 1>&2
    display_help
    ;;
  :)
    echo "Option -$OPTARG requires an argument." 1>&2
    display_help
    ;;
  esac
done

# Check if number of windows and command are provided
if [ -z "$num_windows" ] || [ -z "$command" ]; then
  echo "Both number of windows and command are required."
  display_help
fi

# Calculate rows and columns
rows=$(echo "sqrt($num_windows)" | bc)
columns=$(echo "($num_windows + $rows - 1) / $rows" | bc)

# Create a new tmux session
tmux new-session -d -s llm_tester "$command -p 0"

# Create the remaining windows
for ((i = 1; i < $num_windows; i++)); do
  if ((i % $columns == 0)); then
    tmux select-layout -t llm_tester:0 tiled
    tmux select-pane -t 0
    tmux split-window -t llm_tester:0 -v "$command -p $i"
  else
    tmux split-window -t llm_tester:0 -h "$command -p $i"
  fi
done

# Balance the windows
tmux select-layout -t llm_tester:0 tiled

# Attach to the session
tmux attach-session -t llm_tester