#!/bin/bash

# 配置参数
FILE="/home/ubuntu/rerank_test/titles.1.8w"
LINES=400
REPEAT=10
URLS=(
    "http://localhost:6007/rerank"
    "http://172.19.27.9:9997/v1/rerank"
)

# 检查必要命令
if ! command -v jq &> /dev/null; then
    echo "错误：jq 未安装，请先安装 jq。"
    exit 1
fi

# 检查文件是否存在
if [ ! -f "$FILE" ]; then
    echo "错误：文件 $FILE 不存在。"
    exit 1
fi

# 读取前 LINES 行，构建 docs 数组的 JSON
echo "正在从 $FILE 读取前 $LINES 行..."
docs_json=$(head -n "$LINES" "$FILE" | jq -R . | jq -s .)
if [ -z "$docs_json" ]; then
    echo "错误：读取文件失败或文件为空。"
    exit 1
fi

# 构建完整请求体 JSON
base_json='{"query": "健身女生T恤短袖", "top_n": 386, "normalize": true}'
full_json=$(echo "$base_json" | jq --argjson docs "$docs_json" '. + {docs: $docs}')
if [ -z "$full_json" ]; then
    echo "错误：构建 JSON 失败。"
    exit 1
fi

# 创建临时文件存放请求体
temp_file=$(mktemp)
echo "$full_json" > "$temp_file"
echo "请求体已保存到临时文件: $temp_file"

# 函数：测试单个 URL
test_url() {
    local url=$1
    local durations=()
    local failed=0

    echo "========================================"
    echo "开始测试 URL: $url"
    echo "重复次数: $REPEAT"
    echo "----------------------------------------"

    for i in $(seq 1 $REPEAT); do
        # 开始计时（纳秒）
        start=$(date +%s%N)

        # 发送请求，-s 静默，-o /dev/null 丢弃响应体，--data-binary 直接读取文件
        http_code=$(curl -s -o /dev/null -w "%{http_code}" -X POST "$url" \
            -H "Content-Type: application/json" \
            --data-binary @"$temp_file")

        # 结束计时
        end=$(date +%s%N)
        duration=$(( (end - start) / 1000000 ))  # 毫秒

        # 判断是否成功（HTTP 2xx 视为成功）
        if [[ "$http_code" -ge 200 && "$http_code" -lt 300 ]]; then
            durations+=($duration)
            echo "第 $i 次: 耗时 ${duration} ms (HTTP $http_code)"
        else
            echo "第 $i 次: 失败 (HTTP $http_code) 耗时 ${duration} ms"
            failed=$((failed + 1))
        fi

        # 可选：每次请求间隔一点时间，避免过载（取消注释以启用）
        # sleep 0.1
    done

    # 统计成功请求的耗时
    local count=${#durations[@]}
    echo "----------------------------------------"
    if [ $count -eq 0 ]; then
        echo "所有请求均失败，无法统计耗时。"
        return
    fi

    # 计算总和、平均值、最小值、最大值
    local sum=0
    local min=${durations[0]}
    local max=${durations[0]}
    for d in "${durations[@]}"; do
        sum=$((sum + d))
        (( d < min )) && min=$d
        (( d > max )) && max=$d
    done
    local avg=$((sum / count))

    echo "成功请求数: $count / $REPEAT"
    echo "失败请求数: $failed"
    echo "平均耗时: $avg ms"
    echo "最小耗时: $min ms"
    echo "最大耗时: $max ms"
    echo "========================================"
}

# 依次测试每个 URL
for url in "${URLS[@]}"; do
    test_url "$url"
done

# 清理临时文件
rm -f "$temp_file"
echo "临时文件已删除。"