#!/bin/bash
CONDA_ENV_NAME='vwebarena'          # the name of your conda environment
TOTAL_PROCESSES=4                # the number of processes to run in parallel   
test_config_base_dir=''
TASKS_FILE=''               
SKIP_COOKIES=false

# Parse optional arguments (-t for tasks file, -s to skip cookies, -c for config dir, -n for number of processes, -a for agent config)

while getopts ":t:c:n:a:s" opt; do
    case $opt in
        t) TASKS_FILE="$OPTARG" ;;
        s) SKIP_COOKIES=true ;;
        c) test_config_base_dir="$OPTARG" ;;
        n) TOTAL_PROCESSES="$OPTARG" ;;
        a) agent_config="$OPTARG" ;;
        \?) echo "Invalid option -$OPTARG" >&2 ;;
    esac
done
# Example usage:
# ./p_run.sh -c config_files/vwa_not_vague/test_shopping -n 4 -a config/agent_config.yaml
# ./p_run.sh -t evaluation_harness/task_subsets/shopping.txt -n 4 -a config/agent_config.yaml



# Tmux commands for creating panes
tmux_commands=(
    'tmux split-window -h'  # 2 panes
    'tmux select-pane -t 0; tmux split-window -v' # 3 panes
    'tmux select-pane -t 0; tmux split-window -v' # 4 panes
    'tmux select-pane -t 3; tmux split-window -v' # 5 panes
    'tmux select-pane -t 3; tmux split-window -v' # 6 panes
)

# Fill with NUM_TOTAL_PROCESSES strings if want results to be saved in specific directories
results_dirs=(
    '' '' '' '' ''
)

# Fill with at least NUM_TOTAL_PROCESSES API keys
# api_keys=(
# )

api_keys=(

)
# Tells which device the captioner model is running on
agent_captioning_model_devices=(
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
    'server-cuda'
)
#===============================================================================
# Helpers
#===============================================================================

# Function to run a single job in a tmux pane.
run_job() {
    pane_index=$1; start_id=$2; end_id=$3; index=$4; task_file=$5;

    if [[ -z "$task_file" ]]; then
        task_file="''"
    fi

    tmux select-pane -t $pane_index

    if [[ -n "${results_dirs[$index]}" ]]; then
        cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -d ${results_dirs[$index]} -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r; exit"
        # cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -d ${results_dirs[$index]} -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r"
    else
        cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r; exit"
        # cmd="conda activate ${CONDA_ENV_NAME}; ./scripts/runs/run.sh -s $start_id -e $end_id -c ${test_config_base_dir} -k \"${api_keys[$index]}\" -t $task_file -m ${agent_captioning_model_devices[$index]} -a \"${agent_config}\" -r"
    fi
    tmux send-keys "$cmd" ENTER
    sleep 3
}
# Function to run batches of jobs
run_batch() {
    if [[ -n "$TASKS_FILE" ]]; then
        num_jobs=$TOTAL_PROCESSES
        for ((i=0; i<num_jobs; i++)); do
            pane_index=$i
            task_file="tasks$((i+1)).txt"
            run_job $pane_index 0 1 $i $task_file
        done
    else
        args=("$@")  # all arguments as an array
        num_jobs=$(( ${#args[@]} / 2 ))
        for ((i=0; i<num_jobs; i++)); do
            pane_index=$i
            start_id=${args[$((i*2))]}
            end_id=${args[$((i*2+1))]}
            run_job $pane_index $start_id $end_id $i ""
        done
    fi
}


parse_tasks_file() {
    local tasks_file_path="$1"
    local config_dir=""
    
    # Read the first line as config base directory
    config_dir=$(head -n 1 "$tasks_file_path")
    
    # Read the rest of the lines as task IDs and filter integers
    # Using mapfile (readarray) to read directly into an array
    mapfile -t task_ids_local < <(tail -n +2 "$tasks_file_path" | grep -E '^[0-9]+$')
    
    # Return the config_dir and the array elements via stdout
    echo "$config_dir"
    printf "%s\n" "${task_ids_local[@]}"
}

#===============================================================================
# MAIN
#===============================================================================
# --- Copy api_keys.json to api_keys_copy.json if it exists -----
if [[ -f "api_keys.json" ]]; then
    cp api_keys.json api_keys_copy.json
fi

# --- Create autologin cookies if not skipped -----
# Obs.: the `run.sh` calls don't create cookies; it is only created here.
if [[ "$SKIP_COOKIES" == false ]]; then
    echo "Creating autologin cookies"
    ./scripts/environments/autologin_cookies.sh 'local_vwebarena'
fi

# Check if host_captioner.py is running; if not, host the captioner in a new tmux session
if ! pgrep -f "host_captioner.py" > /dev/null; then
    echo "No captioner found. Starting  'host_captioner.py' in tmux 'captioner' session"
    tmux new-session -d -s "captioner" "conda activate ${CONDA_ENV_NAME}; python utils/host_captioner.py"
    sleep 20
fi

# --- Create a new tmux session and the required panes -----
TMUX_SESSION=$(tmux new-session -d -P -F "#{session_id}")
for ((i=0; i<$TOTAL_PROCESSES-1; i++)); do
    eval ${tmux_commands[$i]}
done

# --- Distribute and run tasks -----
if [[ -n "$TASKS_FILE" ]]; then
    # If task file provided, get tasks from it
    parsed_output=($(parse_tasks_file "$TASKS_FILE")) # Capture ALL output of parse_tasks_file
    test_config_base_dir="${parsed_output[0]}"       # First line is config_dir
    all_tasks=("${parsed_output[@]:1}")              # Remaining lines are task_ids

    total_tasks=${#all_tasks[@]}
    
    echo "Test config base dir: $test_config_base_dir"
    echo "Total tasks: $total_tasks"
    
    tasks_per_process=$(( (total_tasks + TOTAL_PROCESSES - 1) / TOTAL_PROCESSES ))

    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        start_index=$(( i * tasks_per_process ))
        end_index=$(( (i + 1) * tasks_per_process - 1 ))
        (( end_index >= total_tasks )) && end_index=$(( total_tasks - 1 ))

        task_file="tasks$((i+1)).txt"
        > "$task_file"
        for ((j=start_index; j<=end_index && j<total_tasks; j++)); do
            echo "${all_tasks[j]}" >> "$task_file"
        done
    done

    run_batch
    sleep 1
    # Remove temporary task files
    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        rm -f "tasks$((i+1)).txt"
    done
else
    # Determine the start and end IDs based on files in the config base directory
    start_id=$(ls "$test_config_base_dir" | grep -oP '^\d+' | sort -n | head -n 1)
    end_id=$(ls "$test_config_base_dir" | grep -oP '^\d+' | sort -n | tail -n 1)

    # If no task file provided, run all tasks in `test_config_base_dir`
    total_jobs=$((end_id - start_id + 1))
    batch_size=$((total_jobs / TOTAL_PROCESSES))
    remaining_jobs=$((total_jobs % TOTAL_PROCESSES))
    
    job_args=()
    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        batch_start=$(( start_id + i * batch_size + (i < remaining_jobs ? i : remaining_jobs) ))
        batch_end=$(( batch_start + batch_size - 1 + (i < remaining_jobs ? 1 : 0) ))
        job_args+=($batch_start $batch_end)
    done

    for ((i=0; i<TOTAL_PROCESSES; i++)); do
        echo "Process $i: Start ID = ${job_args[$((i*2))]}, End ID = ${job_args[$((i*2+1))]}"
    done

    run_batch "${job_args[@]}"
fi

echo "Experiments running on tmux session $TMUX_SESSION"