#!/bin/bash
# run_job.sh

# —- Config ——
METAS_FILE="genius_metas.jsonl"      # JSONL file path
CHUNK_SIZE=5000                       # Number of lines per worker
MAX_JOBS=12                           # Maximum number of concurrent processes
DOCKER_IMAGE="my-python-jupyter-suno"
SCRIPT_CMD="python3 -u process_audio_features.py"  # Added -u for unbuffered output
# ——————

# 1) Calculate the total number of lines in the file
TOTAL_LINES=$(wc -l < "$METAS_FILE")
echo "Total lines in $METAS_FILE: $TOTAL_LINES"

# 2) Calculate the total number of chunks (rounded up)
NUM_CHUNKS=$(( (TOTAL_LINES + CHUNK_SIZE - 1) / CHUNK_SIZE ))
echo "Will split into $NUM_CHUNKS chunks (chunk size = $CHUNK_SIZE)."

# Function to count running containers
count_running_containers() {
    sudo docker ps --filter "name=audio_proc_batch_" -q | wc -l
}

# 3) Loop through each chunk and submit it to the background
for (( batch_idx=0; batch_idx < 8; batch_idx++ )); do
  echo "[`date +'%H:%M:%S'`] Launch batch #${batch_idx} (${batch_idx}/${NUM_CHUNKS-1})"
  
  CONTAINER_NAME="audio_proc_batch_${batch_idx}"
  
  sudo docker run -d \
    --name "$CONTAINER_NAME" \
    --rm \
    -v "$(pwd)":/app \
    -w /app \
    "$DOCKER_IMAGE" \
    $SCRIPT_CMD --batch_idx="$batch_idx"

  # If the number of running containers ≥ MAX_JOBS, sleep until there is a free slot
  while (( $(count_running_containers) >= MAX_JOBS )); do
    sleep 1
  done
done

echo "All $NUM_CHUNKS batches have been launched. Waiting for them to finish…"

# 4) Wait for all containers to finish
while [ $(count_running_containers) -gt 0 ]; do
    echo "[`date +'%H:%M:%S'`] Still running: $(count_running_containers) containers"
    sleep 10
done

echo ">>> All batches processed. <<<"