#!/bin/bash
#
# Copyright 2026 KylinSoft  Co., Ltd.
#
# Licensed under the Apache License, Version 2.0 (the "License");
# you may not use this file except in compliance with the License.
# You may obtain a copy of the License at
#
#     http://www.apache.org/licenses/LICENSE-2.0
#
# Unless required by applicable law or agreed to in writing, software
# distributed under the License is distributed on an "AS IS" BASIS,
# WITHOUT WARRANTIES OR CONDITIONS OF ANY KIND, either express or implied.
# See the License for the specific language governing permissions and
# limitations under the License.
#
#
# XSched WRR (Weighted Round Robin) fairness test script
# Defaults to local llama.cpp server, can be overridden by env vars:
#   API_URL="http://host:port" ./this_script.sh
#
# Prerequisites:
#   Start the llama.cpp server with XSCHED_POLICY=WRR before running.
#   e.g. export XSCHED_POLICY=WRR && ./build/bin/llama-server -m <model> ...
#

set -e

# Support environment variable override, default compatible with llama.cpp standard port
API_URL=${API_URL:-"http://127.0.0.1:8080"}

# Test task parameters (identical tasks for fair comparison)
PROMPT="Write a very long story about artificial intelligence."
N_PREDICT=500

# Execute a single task and return elapsed time (seconds, 2 decimal places)
run_task() {
    local TASK_NAME="$1"
    local LOG_FILE="$2"
    local START END COST
    START=$(date +%s.%N)
    curl -s -X POST "${API_URL}/completion" \
        -H "Content-Type: application/json" \
        -d "{
            \"prompt\": \"$PROMPT\",
            \"n_predict\": $N_PREDICT,
            \"stream\": false,
            \"temperature\": 0.1
        }" -o /dev/null
    END=$(date +%s.%N)
    COST=$(echo "scale=2; ($END - $START)" | bc)
    echo "$COST" > "$LOG_FILE"
}

echo "============================================================="
echo "          XSched WRR Round-Robin Fairness Test"
echo "============================================================="

# Baseline: single task running alone
echo -e "\n[1] Task A Running Alone (Baseline)"
LOG_A_ALONE="./a_alone.tmp"
run_task "A" "$LOG_A_ALONE"
t_alone=$(cat "$LOG_A_ALONE")
echo "    Task A (alone): ${t_alone}s"

# Two tasks running concurrently under WRR
echo -e "\n[2] Task A and Task B Running Concurrently (WRR Fair Sharing)"
LOG_A="./a_wrr.tmp"
LOG_B="./b_wrr.tmp"

# Launch Task B in background, then Task A shortly after
run_task "B" "$LOG_B" &
sleep 0.3
run_task "A" "$LOG_A"

wait

t_A=$(cat "$LOG_A")
t_B=$(cat "$LOG_B")

echo "    Task A: ${t_A}s"
echo "    Task B: ${t_B}s"

# Fairness ratio: under WRR, both tasks should have similar times
# We expect |T_A - T_B| / max(T_A, T_B) to be small (e.g., < 30%)
FAIR_DELTA=$(echo "scale=2; ($t_A - $t_B)" | bc)
FAIR_DELTA=${FAIR_DELTA#-}  # absolute value
FAIR_MAX=$(echo "if ($t_A > $t_B) $t_A else $t_B" | bc)
FAIR_RATIO=$(echo "scale=2; ($FAIR_DELTA * 100 / $FAIR_MAX)" | bc)

echo ""
echo "-------------------------------------------------------------"
echo "                   WRR Fairness Summary"
echo "-------------------------------------------------------------"
echo "Single task baseline:     ${t_alone}s"
echo "Two tasks concurrent:     Task A: ${t_A}s     Task B: ${t_B}s"
echo "Time difference ratio:    ${FAIR_RATIO}%"

if (( $(echo "$FAIR_RATIO < 30" | bc) )); then
    echo "Fairness verdict:         PASS (both tasks completed closely)"
else
    echo "Fairness verdict:         CHECK (difference > 30%, may need investigation)"
fi

echo "============================================================="

# Cleanup
rm -f "$LOG_A_ALONE" "$LOG_A" "$LOG_B"