#!/bin/sh

show_help() {
  echo "Usage: $0 [options] <total_processes> <config1> [--] <config2> [--] ..."
  echo ""
  echo "Spawn multiple agent processes with different client configurations."
  echo "Processes are distributed evenly across all provided configurations."
  echo ""
  echo "Options:"
  echo "  --help          Show this help message and exit"
  echo ""
  echo "Arguments:"
  echo "  <total_processes>  Total number of agent processes to spawn"
  echo "  <configN>          Client config: [@prefix] <command> [args...]"
  echo "                     Use @prefix to add a name prefix (e.g., @sm for smart)"
  echo ""
  echo "Examples:"
  echo "  # Spawn 26 agents with two different bot programs"
  echo "  $0 26 @v2 ./client_v2 -- @v3 ./client_v3"
  echo "  # Results in 13x 'user-v2-xxx' and 13x 'user-v3-xxx'"
  echo ""
  echo "  # Spawn 10 agents with a single bot (original behavior)"
  echo "  $0 10 ./my_bot"
  echo ""
  echo "  # Spawn 20 agents: 10 with prefix, 10 without"
  echo "  $0 20 @test ./experimental_bot -- ./stable_bot"
  echo ""
  echo "Note: If total_processes doesn't divide evenly, remaining processes"
  echo "      are distributed round-robin starting from the first configuration."
}

if [ "$1" = "--help" ]; then
  show_help
  exit 0
fi

if [ "$#" -lt 2 ]; then
  echo "Error: Missing arguments" >&2
  echo "Usage: $0 <total_processes> <config1> [--] <config2> [--] ..." >&2
  echo "Try '$0 --help' for more information." >&2
  exit 1
fi

NUM="$1"
shift

# Validate NUM is a positive integer
case "$NUM" in
  ''|*[!0-9]*)
    echo "Error: total_processes must be a positive integer" >&2
    exit 1
    ;;
esac

if [ "$NUM" -lt 1 ]; then
  echo "Error: total_processes must be at least 1" >&2
  exit 1
fi

PLAY_SH="$(dirname "$0")/play.sh"

if [ ! -x "$PLAY_SH" ]; then
  echo "Error: $PLAY_SH not found or not executable" >&2
  exit 1
fi

# Create temp file to store configs (format: prefix|command)
CONFIG_FILE="$(mktemp)"
trap "rm -f '$CONFIG_FILE'" EXIT

# Parse commands separated by --
# Each config can optionally start with @prefix
CURRENT_CMD=""
CURRENT_PREFIX=""
CMD_COUNT=0

while [ "$#" -gt 0 ]; do
  if [ "$1" = "--" ]; then
    if [ -n "$CURRENT_CMD" ]; then
      printf '%s|%s\n' "$CURRENT_PREFIX" "$CURRENT_CMD" >> "$CONFIG_FILE"
      CMD_COUNT=$((CMD_COUNT + 1))
      CURRENT_CMD=""
      CURRENT_PREFIX=""
    fi
    shift
  else
    # Check if this is a prefix (@something)
    case "$1" in
      @*)
        # Extract prefix without the @ symbol
        CURRENT_PREFIX="${1#@}"
        shift
        ;;
      *)
        if [ -z "$CURRENT_CMD" ]; then
          CURRENT_CMD="$1"
        else
          CURRENT_CMD="$CURRENT_CMD $1"
        fi
        shift
        ;;
    esac
  fi
done

# Don't forget the last command
if [ -n "$CURRENT_CMD" ]; then
  printf '%s|%s\n' "$CURRENT_PREFIX" "$CURRENT_CMD" >> "$CONFIG_FILE"
  CMD_COUNT=$((CMD_COUNT + 1))
fi

if [ "$CMD_COUNT" -lt 1 ]; then
  echo "Error: At least one command must be provided" >&2
  exit 1
fi

echo "Configuration:"
echo "  Total processes: $NUM"
echo "  Client configurations: $CMD_COUNT"
echo ""

# Calculate distribution
BASE_PER_CMD=$((NUM / CMD_COUNT))
REMAINDER=$((NUM % CMD_COUNT))

echo "Distribution:"
i=1
while IFS='|' read -r prefix cmd; do
  if [ "$i" -le "$REMAINDER" ]; then
    count=$((BASE_PER_CMD + 1))
  else
    count=$BASE_PER_CMD
  fi
  if [ -n "$prefix" ]; then
    echo "  Config $i ($count processes, prefix=$prefix): $cmd"
  else
    echo "  Config $i ($count processes): $cmd"
  fi
  i=$((i + 1))
done < "$CONFIG_FILE"
echo ""

# Download agent once for all workers
AGENT_DIR="$(mktemp -d)"
AGENT_PATH="$AGENT_DIR/agent"
echo "Downloading shared agent to $AGENT_PATH..."
wget -q https://lode.codetribe.cz/assets/agent -O "$AGENT_PATH"
chmod +x "$AGENT_PATH"
echo "Shared agent ready."
echo ""

cleanup() {
  echo ""
  echo "Stopping all workers..."
  kill -- -$$ 2>/dev/null
  echo "Cleaning up shared agent..."
  rm -rf "$AGENT_DIR"
  rm -f "$CONFIG_FILE"
  exit 130
}

trap cleanup INT TERM

# Spawn processes distributed across configurations
echo "Starting $NUM processes..."

spawned=0
cmd_index=1

while IFS='|' read -r prefix cmd; do
  if [ "$cmd_index" -le "$REMAINDER" ]; then
    count=$((BASE_PER_CMD + 1))
  else
    count=$BASE_PER_CMD
  fi
  
  # Build prefix argument if set
  if [ -n "$prefix" ]; then
    PREFIX_ARG="--prefix $prefix"
  else
    PREFIX_ARG=""
  fi
  
  j=1
  while [ "$j" -le "$count" ]; do
    spawned=$((spawned + 1))
    if [ -n "$prefix" ]; then
      echo "  Starting process $spawned (config $cmd_index, prefix=$prefix): $cmd"
    else
      echo "  Starting process $spawned (config $cmd_index): $cmd"
    fi
    (
      eval "exec \"$PLAY_SH\" --agent \"$AGENT_PATH\" $PREFIX_ARG $cmd"
    ) &
    j=$((j + 1))
  done
  
  cmd_index=$((cmd_index + 1))
done < "$CONFIG_FILE"

echo ""
echo "All $NUM processes started. Waiting for completion..."
echo "(Press Ctrl+C to stop all workers)"
echo ""

wait

# Cleanup after all workers finish
echo ""
echo "All workers finished."
echo "Cleaning up shared agent..."
rm -rf "$AGENT_DIR"
rm -f "$CONFIG_FILE"