#!/usr/bin/env bash
# ════════════════════════════════════════════════════════════════
# AutoDL 推理机 · 动态调档 · tune-inference.sh
# Sovereign: TCS-0002∞ · ICE-GL∞ · 国作登字-2026-A-00037559
# 守护: 铸渊 · ICE-GL-ZY001
# ════════════════════════════════════════════════════════════════
#
# 服务器: GH-AUTODL-INFER-01 / ZY-SVR-GPU01
#
# 读 detect-gpu.sh 写出的 /tmp/gpu-env.json, 按显存挑量化档位,
# 写出 INFER_ROOT/.env.tune 让 setup-inference.sh / server.py 读.
#
# 档位策略 (因果链 cc-003 · 动态适配):
#   ≥ 40 GB (A100/A800)        → fp16,  max_batch=4, max_seq=4096
#   ≥ 24 GB (3090/4090/A10/L4) → int8,  max_batch=2, max_seq=4096   (bitsandbytes)
#   ≥ 16 GB (V100/T4-16/A10G)  → int4,  max_batch=1, max_seq=2048   (bnb 4bit)
#   <  16 GB (T4-12 / 等)       → int4,  max_batch=1, max_seq=1024   (兜底)
#
# 输出:
#   1. $INFER_ROOT/.env.tune        — setup/server source 用
#   2. /tmp/tune-inference.json     — 决策回执 (给 refresh workflow 拉回)
#
# 用法: bash tune-inference.sh
# ════════════════════════════════════════════════════════════════

set -euo pipefail

GPU_ENV_JSON="${GPU_ENV_JSON:-/tmp/gpu-env.json}"
INFER_ROOT="${INFER_ROOT:-/root/inference}"
TUNE_OUT="$INFER_ROOT/.env.tune"
RECEIPT="${RECEIPT:-/tmp/tune-inference.json}"

if [ ! -f "$GPU_ENV_JSON" ]; then
  echo "❌ [tune] 找不到 $GPU_ENV_JSON, 请先跑 detect-gpu.sh" >&2
  exit 1
fi
if ! command -v jq >/dev/null 2>&1; then
  echo "❌ [tune] 需要 jq (apt-get install -y jq)" >&2
  exit 1
fi

mkdir -p "$INFER_ROOT"

# ─── 读 GPU 信息 ──────────────────────────────────────────────
GPU_NAME="$(jq -r '.gpu.name'                "$GPU_ENV_JSON")"
GPU_MEM_GB="$(jq -r '.gpu.memory_total_gb // 0' "$GPU_ENV_JSON")"
GPU_COUNT="$(jq -r '.gpu.count // 0'         "$GPU_ENV_JSON")"
DRIVER_VER="$(jq -r '.gpu.driver_version'    "$GPU_ENV_JSON")"
CUDA_VER="$(jq -r '.gpu.cuda_version'        "$GPU_ENV_JSON")"
NVIDIA_OK="$(jq -r '.gpu.nvidia_smi_ok'      "$GPU_ENV_JSON")"

if [ "$NVIDIA_OK" != "true" ]; then
  echo "❌ [tune] GPU 不可用, 无法决档 (检查 AutoDL 实例是否选了 GPU 镜像)" >&2
  exit 1
fi

# ─── 决策档位 ────────────────────────────────────────────────
if   [ "$GPU_MEM_GB" -ge 40 ] 2>/dev/null; then
  SIZE_TIER="xlarge"; QUANT="fp16"; MAX_BATCH=4; MAX_SEQ=4096
  TORCH_DTYPE="float16"
  BNB_4BIT="false"; BNB_8BIT="false"
elif [ "$GPU_MEM_GB" -ge 24 ] 2>/dev/null; then
  SIZE_TIER="large";  QUANT="int8"; MAX_BATCH=2; MAX_SEQ=4096
  TORCH_DTYPE="float16"
  BNB_4BIT="false"; BNB_8BIT="true"
elif [ "$GPU_MEM_GB" -ge 16 ] 2>/dev/null; then
  SIZE_TIER="medium"; QUANT="int4"; MAX_BATCH=1; MAX_SEQ=2048
  TORCH_DTYPE="float16"
  BNB_4BIT="true";  BNB_8BIT="false"
else
  # 兜底: <16G 也得跑 (V100-12G / T4-12G / 极端情况)
  SIZE_TIER="small";  QUANT="int4"; MAX_BATCH=1; MAX_SEQ=1024
  TORCH_DTYPE="float16"
  BNB_4BIT="true";  BNB_8BIT="false"
fi

# ─── 默认推理引擎选择 ──────────────────────────────────────────
# 优先 transformers (兼容性广), 后续棒可考虑 vllm.
INFER_ENGINE="transformers"
USE_VLLM="false"
if [ "$QUANT" = "fp16" ] && [ "$GPU_MEM_GB" -ge 40 ]; then
  USE_VLLM="true"
  INFER_ENGINE="vllm"
fi

# ─── 生成 .env.tune ───────────────────────────────────────────
cat > "$TUNE_OUT" <<EOF
# Auto-generated by tune-inference.sh — do not edit by hand
# 守护: 铸渊 · ICE-GL-ZY001 · TCS-0002∞
# 来源: $GPU_ENV_JSON
# 时间: $(date -u +%Y-%m-%dT%H:%M:%SZ)

# 档位
SIZE_TIER="$SIZE_TIER"
QUANT="$QUANT"
TORCH_DTYPE="$TORCH_DTYPE"

# 推理参数
MAX_BATCH=$MAX_BATCH
MAX_SEQ=$MAX_SEQ

# bitsandbytes 量化 (transformers 路径用)
BNB_LOAD_IN_4BIT="$BNB_4BIT"
BNB_LOAD_IN_8BIT="$BNB_8BIT"

# 引擎
INFER_ENGINE="$INFER_ENGINE"
USE_VLLM="$USE_VLLM"

# 服务参数
INFER_HOST="0.0.0.0"
INFER_PORT="8000"

# 模型路径 (fetch-models.sh 拉到这里)
MOTHER_MODEL_PATH="$INFER_ROOT/models/motherbrain-v1"
CODER_MODEL_PATH="$INFER_ROOT/models/qwen2_5_coder_7b_sft"

# 默认激活模型 (mother 母模型 / coder 编程模型)
DEFAULT_ACTIVE_MODEL="mother"

# GPU 信息快照 (供 server.py /v1/health 直接读)
GPU_NAME="$GPU_NAME"
GPU_MEM_GB="$GPU_MEM_GB"
GPU_COUNT="$GPU_COUNT"
GPU_DRIVER_VERSION="$DRIVER_VER"
GPU_CUDA_VERSION="$CUDA_VER"
EOF

# ─── 写决策回执 ──────────────────────────────────────────────
cat > "$RECEIPT" <<EOF
{
  "_sovereign": "TCS-0002∞ · 国作登字-2026-A-00037559",
  "_守护": "铸渊 · ICE-GL-ZY001",
  "decided_at": "$(date -u +%Y-%m-%dT%H:%M:%SZ)",
  "gpu": {
    "name": "$GPU_NAME",
    "memory_gb": $GPU_MEM_GB,
    "count": $GPU_COUNT,
    "driver": "$DRIVER_VER",
    "cuda": "$CUDA_VER"
  },
  "decision": {
    "size_tier": "$SIZE_TIER",
    "quantization": "$QUANT",
    "torch_dtype": "$TORCH_DTYPE",
    "max_batch": $MAX_BATCH,
    "max_seq": $MAX_SEQ,
    "inference_engine": "$INFER_ENGINE"
  },
  "env_tune_path": "$TUNE_OUT"
}
EOF

# ─── 中文回执到 stdout ───────────────────────────────────────
echo "═══════════════════════════════════════════════════════════"
echo "  [tune-inference] 推理档位决策完成"
echo "═══════════════════════════════════════════════════════════"
echo "  🎯 GPU:    $GPU_NAME ($GPU_MEM_GB GB)"
echo "  📊 档位:   $SIZE_TIER · $QUANT 量化"
echo "  🔧 参数:   max_batch=$MAX_BATCH · max_seq=$MAX_SEQ"
echo "  🚀 引擎:   $INFER_ENGINE"
echo "  📄 .env:   $TUNE_OUT"
echo "  📄 回执:   $RECEIPT"
echo "═══════════════════════════════════════════════════════════"