#!/usr/bin/env bash
# 分区重分配封装脚本：生成 → 执行 → 验证 → 取消限速
# 用法：
#   ./reassign_helper.sh generate "topic1,topic2" "1,2,3,4"     # 生成方案
#   ./reassign_helper.sh execute reassign.json 50000000          # 执行（限速 50MB/s）
#   ./reassign_helper.sh verify reassign.json                    # 验证
#   ./reassign_helper.sh cancel-throttle reassign.json           # 单独取消限速
set -euo pipefail

BS=${BS:-localhost:9092}
KAFKA_HOME=${KAFKA_HOME:-/opt/kafka}
SCRIPT="$KAFKA_HOME/bin/kafka-reassign-partitions.sh"

usage() {
    cat <<EOF
Usage:
  $0 generate <topic_csv> <broker_csv>
       生成 reassign.json（基于 broker-list 平均分布）
  $0 execute  <reassign.json> [throttle_bytes_per_sec]
       执行（默认限速 50MB/s）
  $0 verify   <reassign.json>
       验证当前进度
  $0 cancel-throttle <reassign.json>
       手动取消限速（verify 完成后会自动取消，但失败时可手动）
EOF
    exit 1
}

cmd=${1:-}
case "$cmd" in
    generate)
        topics=$2; brokers=$3
        # 生成 topics-to-move JSON
        TJSON=$(mktemp)
        echo -n '{"topics":[' > "$TJSON"
        IFS=',' read -ra arr <<< "$topics"
        for i in "${!arr[@]}"; do
            [[ $i -gt 0 ]] && echo -n ',' >> "$TJSON"
            echo -n "{\"topic\":\"${arr[$i]}\"}" >> "$TJSON"
        done
        echo ']  ,"version":1}' >> "$TJSON"
        echo "==> 生成方案 (broker-list=$brokers)"
        "$SCRIPT" --bootstrap-server "$BS" \
            --topics-to-move-json-file "$TJSON" \
            --broker-list "$brokers" \
            --generate
        echo ""
        echo "✅ 把上面 'Proposed partition reassignment configuration' 段落保存为 reassign.json"
        ;;

    execute)
        plan=$2; throttle=${3:-50000000}
        echo "==> 执行重分配，限速 = $throttle bytes/s ($((throttle/1024/1024)) MB/s)"
        "$SCRIPT" --bootstrap-server "$BS" \
            --reassignment-json-file "$plan" \
            --execute --throttle "$throttle"
        echo "✅ 已提交，下一步用 verify 查看进度"
        ;;

    verify)
        plan=$2
        echo "==> 验证进度"
        "$SCRIPT" --bootstrap-server "$BS" \
            --reassignment-json-file "$plan" \
            --verify
        ;;

    cancel-throttle)
        plan=$2
        echo "==> 取消限速"
        "$SCRIPT" --bootstrap-server "$BS" \
            --reassignment-json-file "$plan" \
            --verify
        echo "（verify 会自动清理；如仍残留请用 kafka-configs.sh 手动 delete）"
        ;;

    *) usage ;;
esac