#!/bin/bash

# Default API URL
API_URL="http://stitcher.local:5000"
OUTPUT_DIR=""
ITERATION_RANGE=""
SAVE_OUTPUT=true

# Parse command line arguments
while [[ "$#" -gt 0 ]]; do
    case $1 in
        -o|--output) OUTPUT_DIR="$2"; shift ;;
        -u|--url) API_URL="$2"; shift ;;
        -i|--iteration) ITERATION_RANGE="$2"; shift ;;
        -n|--no-output) SAVE_OUTPUT=false ;;
        *) echo "Unknown parameter passed: $1"; exit 1 ;;
    esac
    shift
done

# Check if output directory is specified
if [ -z "$OUTPUT_DIR" ]; then
    echo "Error: Output directory not specified. Use -o or --output to specify the directory."
    exit 1
fi

# Create the output directory if it doesn't exist
mkdir -p "$OUTPUT_DIR"

LOG_FILE="${OUTPUT_DIR}/results.txt"
PERFORMANCE_FILE="${OUTPUT_DIR}/performance.csv"

# Add CSV header
echo "Title,Latency(s),Payload" > "$PERFORMANCE_FILE"

exec > >(tee -a "$LOG_FILE") 2>&1

# Define the tests and their corresponding payloads
titles=(
    "01_L19-N21-Small_Square_Image_3x3_tiles"
    "02_H27-K30-Small_Square_Image_4x4_tiles_0.75_scale"
    "03_Q46-W48-Small_Rectangle_image_3x6_tiles"
    "04_K22-N29-Small_Long_Rectangle_Image_8x4_tiles_0.5_scale"
    "05_O30-P31-Small_Square_Cropped_to_LeftOneTile"
    "06_X46-Y47-Small_Square_Cropped_to_MiddleEqually"
    "07_T21-Z27-Medium_7x7_Square_Cropped_to_LeftTop4x3Tile"
    "08_X14-AD20-Medium_7x7_Square_Cropped_to_MiddleEqually"
    "09_A1-AE13-Entire_Left_Panel"
    "10_A14-AE42-Entire_Middle_Panel"
    "11_A1-AE55-Full_Image"
    "12_W48-Q46-ReversedTile_Small_Rectangle_image_3x6_tiles"
)

payloads=(
    '{"canvas_rect":"L19:N21","crop_offset":[0,0],"crop_size":[1,1],"output_scale":1}'
    '{"canvas_rect":"H27:K30","crop_offset":[0,0],"crop_size":[1,1],"output_scale":0.75}'
    '{"canvas_rect":"Q46:W48","crop_offset":[0,0],"crop_size":[1,1],"output_scale":1}'
    '{"canvas_rect":"K22:N29","crop_offset":[0,0],"crop_size":[1,1],"output_scale":0.5}'
    '{"canvas_rect":"O30:P31","crop_offset":[0.075,0.625],"crop_size":[0.4,0.35],"output_scale":1}'
    '{"canvas_rect":"X46:Y47","crop_offset":[0.25,0.25],"crop_size":[0.5,0.5],"output_scale":1}'
    '{"canvas_rect":"T21:Z27","crop_offset":[0.0125,0.0125],"crop_size":[0.55,0.375],"output_scale":1}'
    '{"canvas_rect":"X14:AD20","crop_offset":[0.25,0.25],"crop_size":[0.5,0.5],"output_scale":1}'
    '{"canvas_rect":"A1:AE13","crop_offset":[0,0],"crop_size":[1,1],"output_scale":1}'
    '{"canvas_rect":"A14:AE42","crop_offset":[0,0],"crop_size":[1,1],"output_scale":1}'
    '{"canvas_rect":"A1:AE55","crop_offset":[0,0],"crop_size":[1,1],"output_scale":1}'
    '{"canvas_rect":"W48:Q46","crop_offset":[0,0],"crop_size":[1,1],"output_scale":1}'
)

# Determine the loop range
if [ -z "$ITERATION_RANGE" ]; then
    start=0
    end=$((${#titles[@]} - 1))
elif [[ $ITERATION_RANGE == -* ]]; then
    index=$((${ITERATION_RANGE#-} - 1));
    if [ "$index" -lt 0 ] || [ "$index" -ge "${#titles[@]}" ]; then echo "Error: Invalid iteration number."; exit 1; fi
    start=$index; end=$index
elif [[ $ITERATION_RANGE == *-* ]]; then
    IFS='-' read -r start_range end_range <<< "$ITERATION_RANGE";
    start=$((start_range - 1)); end=$((end_range - 1));
    if [ "$start" -lt 0 ] || [ "$end" -ge "${#titles[@]}" ] || [ "$start" -gt "$end" ]; then echo "Error: Invalid iteration range."; exit 1; fi
else
    index=$(($ITERATION_RANGE - 1));
    if [ "$index" -lt 0 ] || [ "$index" -ge "${#titles[@]}" ]; then echo "Error: Invalid iteration number."; exit 1; fi
    start=$index; end=$index
fi

# Loop through the tests and execute curl
for i in $(seq $start $end); do
    title="${titles[$i]}"
    payload="${payloads[$i]}"
    
    echo "Running test: $title"

    # Set the output target for curl. If --no-output is used, send to /dev/null
    if [ "$SAVE_OUTPUT" = true ]; then
        output_target="$OUTPUT_DIR/${title}.png"
    else
        output_target="/dev/null"
    fi

    time_taken=$(curl -o "$output_target" -w "%{time_total}\n" -X POST "${API_URL}/api/image/generate" -H "Content-Type: application/json" -d "$payload" -s | head -n 1)

    if [ "$SAVE_OUTPUT" = true ]; then
        echo "Image saved to $output_target"
    else
        echo "Request completed (output discarded)."
    fi
    
    echo "Total time taken: ${time_taken}s"
    echo "---------------------------------"

    # Sanitize payload, somehow the python parser can't accept comma "," inside the payload
    payload_sanitized="${payload//,/-}"

    echo "\"$title\",$time_taken,\"$payload_sanitized\"" >> "$PERFORMANCE_FILE"
done

echo "All specified tests completed."


# --- Graph Generation Section ---
PYTHON_SCRIPT_NAME="generate_graph.py"

if [ ! -f "$PYTHON_SCRIPT_NAME" ]; then
    echo "Warning: Python script '$PYTHON_SCRIPT_NAME' not found. Skipping graph generation."
    exit 0
fi
if ! command -v python3 &> /dev/null; then
    echo "Warning: python3 is not installed. Skipping graph generation."
    exit 0
fi

FOLDER_NAME=$(basename "$OUTPUT_DIR")

GRAPH_OUTPUT_PATH="${OUTPUT_DIR}/latency-${FOLDER_NAME}.png"

echo "---------------------------------"
echo "Generating performance graph..."
python3 "$PYTHON_SCRIPT_NAME" --input "$PERFORMANCE_FILE" --output "$GRAPH_OUTPUT_PATH" --user "$FOLDER_NAME"