#!/bin/bash


reset_after_n_tasks=20

# This script runs a sequence of experiments, resetting the environments and calling `p_run.sh` for each.

# USAGE: 
# Pre-run: Host the captioner with tmux => `./scripts/environments/start_reset_envs.sh host_captioner`
# 1) Define the parameters for the sequence of experiments below
# DEPRECATED: 2) sudo -v => enter your credentials (not needed if `./scripts/environments/start_reset_envs.sh all_vwa` don't require sudo like if homepage is not hosted on port 80)
# 2) Add this line to your /etc/sudoers: <user> ALL=(ALL) NOPASSWD: <full_path>/scripts/environments/start_reset_envs.sh
#   2.1) 
# 3) Run the script on the background: nohup ./scripts/runs/p_run_many.sh > prun.log 2>&1 < /dev/null & 
 
# Notes:
# - Check `prun.log` for the execution progress or any errors.
# - To kill this script running on the background: `pkill -f p_run_many.sh`

#--------------------------------------------------
# Parameters
#--------------------------------------------------

declare -A test_config_base_dir_defaults=(
    ["shopping"]="config_files/vwa_not_vague/test_shopping"
    ["classifieds"]="config_files/vwa_not_vague/test_classifieds"
    ["reddit"]="config_files/vwa_not_vague/test_reddit"
)


# Which domains to evaluate for each agent config, formatted as:
# <agent_config_path>| <domain> | <task_list_path>
# Using a string to preserve the order of runs.

base_tst_dir="evaluation_harness/task_subsets"
agent_test_configs=(
    "config/agent_config_t_100_2p_tri_gemini.yaml | $base_tst_dir/reddit.txt $base_tst_dir/classifieds.txt $base_tst_dir/shopping.txt"
)

#--------------------------------------------------
# Helper functions
#--------------------------------------------------


# Kills the process in case of ctrl+c
cleanup() {
    echo "Ctrl+C pressed or termination signal received. Cleaning up..."
    # kill tmux server
    tmux kill-server 2>/dev/null || true
    # Kill any remaining background jobs started by this script.
    kill $(jobs -p) 2>/dev/null || true
    exit 1
}
trap cleanup SIGINT SIGTERM


# Waits until all tmux panes in the provided session are finished.
wait_for_tmux_session() {
    local session_id="$1"
    echo "Waiting for all jobs in tmux session $session_id to finish..."
    while true; do
        # Check if the session still exists.
        if ! tmux has-session -t "$session_id" 2>/dev/null; then
            echo "Session $session_id not found. Assuming it's closed."
            break
        fi

        alive_count=0
        # List pane statuses; if tmux errors (session closed), ignore.
        while read -r pane_dead; do
            if [[ "$pane_dead" != "1" ]]; then
                alive_count=$((alive_count+1))
            fi
        done < <(tmux list-panes -t "$session_id" -F "#{pane_dead}" 2>/dev/null)
        
        if [[ $alive_count -eq 0 ]]; then
            break
        fi
        echo "[$(date '+%Y-%m-%d %H:%M:%S')] $alive_count pane(s) still running..."
        sleep 60
    done

    echo "All jobs finished in tmux session $session_id."
    # Kill the session if it still exists; ignore error if not.
    tmux kill-session -t "$session_id" 2>/dev/null || true
}


# Runs the p_run script, captures the tmux id, and waits until session completes.
run_and_wait() {
    # Forward all arguments to p_run.
    session_output=$(./scripts/runs/p_run.sh "$@")
    echo "$session_output"
    
    # Extract the tmux session id.
    # This assumes p_run.sh prints a line like:
    # "Experiments running on tmux session $7"
    # The sed command removes a leading "$" if present.
    tmux_session=$(echo "$session_output" | grep -E "tmux session" | tail -n 1 | sed -E 's/.*tmux session\s+\$?([0-9]+).*/\1/')
    
    if [[ -z "$tmux_session" ]]; then
        echo "Failed to capture tmux session id from output. Exiting."
        exit 1
    fi

    echo "Detected tmux session: $tmux_session"
    wait_for_tmux_session "$tmux_session"
}



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
#--------------------------------------------------
# Iterate over each agent config and run the tests for the specified domains.
for agent_test_config in "${agent_test_configs[@]}"; do
    sudo ./scripts/environments/start_reset_envs.sh all_vwa

    IFS='|' read -r agent_config tasks_file <<< "$agent_test_config"
    agent_config=$(echo "$agent_config" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')
    tasks_file=$(echo "$tasks_file" | sed -e 's/^[[:space:]]*//' -e 's/[[:space:]]*$//')
    
    # Split the domains string into an array and iterate over each domain.
    read -ra task_files <<< "$tasks_file"

    
    # Iterate over task files
    for i in "${!task_files[@]}"; do
        task_file="${task_files[$i]}"

        # If custom tasks files were provided, use the corresponding one; otherwise, run all tasks in the test_config_base_dir
        if [[ -n "$tasks_file" ]]; then
            task_list="${task_files[$i]}"
            parsed_output=($(parse_tasks_file "$task_list")) # 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
        else
            test_config_base_dir="${test_config_base_dir_defaults[$task_file]}"
            # get task list from test_config_base_dir
            readarray -t all_tasks < <(ls "$test_config_base_dir" | grep -oP '^\d+' | sort -n)
        fi

        echo "----------------------------------------------------------------------------------------------------------"
        echo "[$(date '+%Y-%m-%d %H:%M:%S')] Starting: ${agent_config}, ${task_list}"
        echo "----------------------------------------------------------------------------------------------------------"
        
        domain=${test_config_base_dir##*/test_}
        echo "Domain: $domain"

        total_tasks=${#all_tasks[@]}
        echo "Total tasks: $total_tasks"
        
        # Divide into chunks of reset_after_n_tasks
        echo "Dividing into chunks of $reset_after_n_tasks tasks"

        chunks=()
        declare -a chunks  # Declare the chunks array
        for ((j=0; j<total_tasks; j+=reset_after_n_tasks)); do
            chunk=("${all_tasks[@]:j:reset_after_n_tasks}")
            chunks+=("${chunk[*]}")  # Add the chunk as a single element
        done
        
        echo "# of chunks: ${#chunks[@]}"
        

        # Run each chunk
        temp_dir="./temp"
        mkdir -p "$temp_dir"
        for chunk_str in "${chunks[@]}"; do
            read -ra chunk <<< "$chunk_str"            

            # Create temporary file with mktemp
            temp_file=$(mktemp "$temp_dir/${domain}_XXXXXX.txt")
            
            # Write config dir and task ids to temp file
            echo "${test_config_base_dir}" > "$temp_file"
            for task_id in "${chunk[@]}"; do
                echo "$task_id" >> "$temp_file"
            done
           
            # # Run the tasks
            echo "Running chunk: ${chunk[@]}"
            run_and_wait -t "$temp_file" -a "$agent_config"

            # Reset envs
            echo "Resetting environment $domain"
            sudo ./scripts/environments/start_reset_envs.sh $domain

            # Remove the temp file
            rm "$temp_file"
        done
    done
done