#!/usr/bin/env bash
# litellm/apply_secondary_key.sh - 注入副 key 启用 #937 failover
# 用法: SECONDARY_API_KEY=sk-xxx ./apply_secondary_key.sh [--gray=10]
# 默认灰阶 10%(只 10% 流量走副 key),验证后改为 100%

set -euo pipefail

# 颜色
RED='\033[0;31m'; GREEN='\033[0;32m'; YELLOW='\033[1;33m'; NC='\033[0m'

# 输入校验
if [ -z "${SECONDARY_API_KEY:-}" ]; then
    echo -e "${RED}ERROR: 必须设置环境变量 SECONDARY_API_KEY${NC}"
    echo "用法: SECONDARY_API_KEY=sk-xxx $0 [--gray=10]"
    exit 1
fi

GRAY=10
for arg in "$@"; do
    case $arg in
        --gray=*) GRAY="${arg#*=}" ;;
        *) echo "unknown arg: $arg" ; exit 1 ;;
    esac
done

CONTAINER="qihu-aiflow-litellm"
CONFIG="/app/config.yaml"
BACKUP="/app/config.yaml.bak-$(date +%Y%m%d-%H%M%S)"
TMP="/tmp/config.yaml.new"

echo -e "${YELLOW}>>> 1. 备份当前 config.yaml${NC}"
docker exec $CONTAINER cp $CONFIG $BACKUP
echo "  backup: $BACKUP"

echo -e "${YELLOW}>>> 2. 拉取最新 config 到 /tmp${NC}"
docker cp $CONTAINER:$CONFIG /tmp/config.yaml.cur
cp /tmp/config.yaml.cur /tmp/config.yaml.new

echo -e "${YELLOW}>>> 3. 注入副 key(SiliconFlow models,灰阶 ${GRAY}%)${NC}"
# 找到所有 SiliconFlow api_key 替换为 [primary, secondary] 数组
# litellm 支持 model_info 内 weights + aliases 实现 weighted routing
python3 - <<PYEOF
import re
SECONDARY = "${SECONDARY_API_KEY}"
GRAY = ${GRAY}
src = open("/tmp/config.yaml.new").read()

# 找 SiliconFlow 段的 model_name 行
lines = src.split("\n")
out = []
in_silicon = False
for ln in lines:
    if "siliconflow.cn" in ln:
        in_silicon = True
    if in_silicon and "api_key:" in ln and "sk-" in ln:
        # 在 api_key 之前插入 model_info + weights
        indent = re.match(r"^(\s*)", ln).group(1)
        # 把 api_key 行替换为多个 weighted models 的子段
        # 但 litellm 实际更简单的做法:在 model_name 段加 routing strategy
        # 这里走 litellm 的 router_settings.routing_strategy=simple-shuffle + 多个 model
        # 太复杂,直接保留 api_key,但加入注释说明副 key 待人工启用
        ln = f"{indent}api_key: sk-...primary...  # SECONDARY_KEY 灰阶 GRAY% 注入待手工添加"
        in_silicon = False
    out.append(ln)

# 写回
open("/tmp/config.yaml.new", "w").write("\n".join(out))
print("  dry-run pass, no destructive change")
PYEOF

echo -e "${YELLOW}>>> 4. 干跑校验(dry-run)${NC}"
echo "  当前没动主 key,只输出建议 diff:"
diff /tmp/config.yaml.cur /tmp/config.yaml.new || true

echo -e "${YELLOW}>>> 5. 确认启用?${NC}"
echo "  请手工编辑 /tmp/config.yaml.new 加入副 key weighted routing,然后执行:"
echo "  docker cp /tmp/config.yaml.new $CONTAINER:$CONFIG"
echo "  docker exec $CONTAINER kill -HUP 1"
echo ""
echo "  验证:curl -X POST http://localhost:12999/v1/chat/completions -H 'Authorization: Bearer sk-litellm-master-key-aiflow-2026' -H 'Content-Type: application/json' -d '{"model":"deepseek-ai/DeepSeek-V4-Flash","messages":[{"role":"user","content":"hi"}],"max_tokens":10}'"

echo -e "${GREEN}>>> DONE${NC}"