#!/usr/bin/env bash
# =============================================================================
# 视频字幕提取系统 — 服务器一键环境配置脚本
#
# 安装/配置：gcc 编译工具链、ffmpeg(含 AV1/dav1d 解码)、uv、Python 虚拟环境、
#            项目依赖，并处理常见环境坑(opencv libGL、UTF-8 locale、HF/pip 镜像、
#            OMP 线程、文件句柄上限)。
#
# 用法(在项目根目录):
#   bash setup_server.sh
# 跑完后:
#   source activate.sh      # 激活环境(设置 venv + 各种环境变量)
#   python main.py --source-list ./text2 --max-workers 12
#
# 幂等：可重复运行，已装的会跳过。
# =============================================================================
set -euo pipefail

# ----------------------------- 可调参数 --------------------------------------
PYTHON_VERSION="${PYTHON_VERSION:-3.11}"
PIP_MIRROR="${PIP_MIRROR:-https://repo.huaweicloud.com/repository/pypi/simple}"
HF_ENDPOINT_URL="${HF_ENDPOINT_URL:-https://hf-mirror.com}"
FFMPEG_STATIC_URL="${FFMPEG_STATIC_URL:-https://johnvansickle.com/ffmpeg/releases/ffmpeg-release-amd64-static.tar.xz}"
PROJECT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)"
VENV_DIR="$PROJECT_DIR/.venv"

cd "$PROJECT_DIR"

# 让本脚本(及自检)优先用项目 bin/ 里的 ffmpeg(我们装的静态版)
export PATH="$PROJECT_DIR/bin:$PATH"

# ----------------------------- 工具函数 --------------------------------------
log()  { echo -e "\n\033[1;32m==> $*\033[0m"; }
warn() { echo -e "\033[1;33m[警告] $*\033[0m"; }
err()  { echo -e "\033[1;31m[错误] $*\033[0m" >&2; }
have() { command -v "$1" >/dev/null 2>&1; }

# root 则不用 sudo，否则用 sudo
if [ "$(id -u)" -eq 0 ]; then SUDO=""; else SUDO="sudo"; fi

# 检测包管理器
PKG=""
if have dnf;     then PKG="dnf"
elif have yum;   then PKG="yum"
elif have apt-get; then PKG="apt"
else warn "未识别的包管理器，跳过系统包安装(请确保已手动装好 gcc/ffmpeg 等)"; fi

# 解压后 coscli/bin 里的二进制常丢失 +x，先无条件加执行权限（放最前，确保一定执行）
for b in coscli-linux-amd64 coscli bin/ffmpeg bin/ffprobe; do
  [ -f "$PROJECT_DIR/$b" ] && chmod +x "$PROJECT_DIR/$b" 2>/dev/null || true
done
log "已给 coscli / bin 二进制加执行权限"

# =============================================================================
log "[1/7] 安装系统依赖(gcc 工具链 + opencv/运行库)"
# =============================================================================
# 关键环境坑：
#  - opencv-python 在无图形界面的服务器上 import 会报 libGL.so.1 缺失 → 装 mesa-libGL/libgl1
#  - 还需要 glib2(libgthread-2.0) 否则 opencv 也可能报错
#  - gcc/gcc-c++/make：少数包没有 wheel 时要本地编译
case "$PKG" in
  dnf|yum)
    # openblas-devel：万一 scipy/numpy 走源码编译时需要(优先还是用 wheel)
    $SUDO $PKG install -y gcc gcc-c++ gcc-gfortran make git tar xz which \
                          mesa-libGL glib2 openblas-devel || warn "部分系统包安装失败，继续"
    ;;
  apt)
    $SUDO apt-get update -y || true
    $SUDO apt-get install -y build-essential gfortran git tar xz-utils curl \
                             libgl1 libglib2.0-0 libopenblas-dev || warn "部分系统包安装失败，继续"
    ;;
esac

# =============================================================================
log "[2/7] 安装 ffmpeg(含 AV1/dav1d 解码) → 项目目录 bin/"
# =============================================================================
# 直接下载静态构建(johnvansickle，自带 dav1d)到项目的 bin/，代码会自动优先用它
# (config_utils 启动时把 bin/ 提到 PATH 最前)。不依赖系统 ffmpeg，也不用 sudo。
# 注意：CPU 版(无 NVDEC)；要 GPU 硬解需另装 cuda 版 ffmpeg。
FFMPEG_DIR="$PROJECT_DIR/bin"
LOCAL_FFMPEG="$FFMPEG_DIR/ffmpeg"
need_ffmpeg=1
if [ -x "$LOCAL_FFMPEG" ] && "$LOCAL_FFMPEG" -hide_banner -decoders 2>/dev/null | grep -qiE 'dav1d|av1'; then
  log "项目自带 ffmpeg 已存在且支持 AV1: $LOCAL_FFMPEG，跳过下载"
  need_ffmpeg=0
fi
if [ "$need_ffmpeg" -eq 1 ]; then
  if have curl || have wget; then
    mkdir -p "$FFMPEG_DIR"
    tmp="$(mktemp -d)"
    log "下载 ffmpeg 静态包到 $FFMPEG_DIR ..."
    if have curl; then curl -L "$FFMPEG_STATIC_URL" -o "$tmp/ff.tar.xz"
    else wget -O "$tmp/ff.tar.xz" "$FFMPEG_STATIC_URL"; fi
    tar -xf "$tmp/ff.tar.xz" -C "$tmp"
    ffdir="$(find "$tmp" -maxdepth 1 -type d -name 'ffmpeg-*' | head -1)"
    if [ -n "$ffdir" ]; then
      cp "$ffdir/ffmpeg" "$ffdir/ffprobe" "$FFMPEG_DIR/"
      chmod +x "$FFMPEG_DIR/ffmpeg" "$FFMPEG_DIR/ffprobe"
      log "ffmpeg/ffprobe 已装到 $FFMPEG_DIR"
    else
      err "解压后未找到 ffmpeg 目录，请手动安装 ffmpeg"
    fi
    rm -rf "$tmp"
  else
    err "没有 curl/wget，无法下载 ffmpeg，请手动安装"
  fi
fi

# =============================================================================
log "[3/7] 安装 uv"
# =============================================================================
if ! have uv; then
  curl -LsSf https://astral.sh/uv/install.sh | sh
  # uv 默认装到 ~/.local/bin 或 ~/.cargo/bin
  export PATH="$HOME/.local/bin:$HOME/.cargo/bin:$PATH"
fi
have uv || { err "uv 安装失败，请检查网络"; exit 1; }
log "uv 版本: $(uv --version)"

# =============================================================================
log "[4/7] 用 uv 安装 Python ${PYTHON_VERSION}"
# =============================================================================
uv python install "$PYTHON_VERSION" || warn "uv python install 失败(可能已装)"

# =============================================================================
log "[5/7] 创建虚拟环境并安装依赖"
# =============================================================================
# 优先用 pyproject.toml + uv.lock(规范、可复现)；没有则回退 requirements.txt。
# pyproject 的 [tool.uv] 已配置：禁止 av/numpy/sklearn 本地编译 + 清华镜像。
# 用 uv venv + uv pip install(现代解析器)。--only-binary 让它自动挑有 wheel 的版本，
# 避免某些受限镜像把新版本只给源码包导致本地编译/报错(不像 uv sync 的 no-build 会硬失败)。
[ -d "$VENV_DIR" ] || uv venv --python "$PYTHON_VERSION" "$VENV_DIR"
# shellcheck disable=SC1091
source "$VENV_DIR/bin/activate"
if [ -f requirements.txt ]; then
  uv pip install -r requirements.txt --index-url "$PIP_MIRROR" \
    --only-binary numpy,scipy,scikit-learn,opencv-python-headless,Pillow,av \
  || uv pip install -r requirements.txt --index-url "$PIP_MIRROR"
else
  warn "未找到 requirements.txt，跳过依赖安装"
fi

# 激活 venv 供后续自检使用
# shellcheck disable=SC1091
[ -f "$VENV_DIR/bin/activate" ] && source "$VENV_DIR/bin/activate"
log "Python: $(python --version 2>&1)  位置: $(command -v python)"

# =============================================================================
log "[6/7] 写出环境变量脚本 activate.sh"
# =============================================================================
# 常见环境坑统一在这里处理，跑项目前 source 一下即可
# 用单引号 heredoc 让脚本体保持字面(不在生成时展开)，实现"自动定位所在目录"，
# 这样整个项目可以搬到任意路径，activate.sh 仍能找到自己的 .venv / bin。
cat > "$PROJECT_DIR/activate.sh" <<'EOF'
#!/usr/bin/env bash
# 跑项目前先: source activate.sh —— 自动定位项目目录，项目可放任意路径
HERE="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

# Python：优先项目内自包含的 ./python(make_portable.sh 产物，可随便放)；否则用 .venv
if [ -x "$HERE/python/bin/python3" ]; then
  export PATH="$HERE/bin:$HERE/python/bin:$PATH"
else
  export VIRTUAL_ENV="$HERE/.venv"
  export PATH="$HERE/bin:$VIRTUAL_ENV/bin:$HOME/.local/bin:$HOME/.cargo/bin:$PATH"
fi

# UTF-8：项目有中文文件名(数据.jsonl)。优先选系统存在的 UTF-8 locale，避免 setlocale 警告；
# 并强制 Python 用 UTF-8(即使系统没有 UTF-8 locale 也能正确处理中文文件名)。
if locale -a 2>/dev/null | grep -qiE '^C\.(utf-?8|UTF-?8)$'; then
  export LANG=C.UTF-8 LC_ALL=C.UTF-8
elif locale -a 2>/dev/null | grep -qiE 'en_US\.(utf-?8|UTF-?8)'; then
  export LANG=en_US.UTF-8 LC_ALL=en_US.UTF-8
fi
export PYTHONUTF8=1
export PYTHONIOENCODING=utf-8

# HuggingFace 镜像（联网下模型时用）
export HF_ENDPOINT=__HF_ENDPOINT__
# 模型缓存：项目内有 ./models 就用它(自包含包)，否则用默认 ~/.cache/huggingface
[ -d "$HERE/models" ] && export HF_HOME="$HERE/models"
# 只在模型【确实已缓存】时才强制离线，否则允许联网下载(避免"没模型又禁网"的死锁)
if ls "${HF_HOME:-$HOME/.cache/huggingface}"/hub/models--*faster-whisper* >/dev/null 2>&1; then
  export HF_HUB_OFFLINE=1
  export TRANSFORMERS_OFFLINE=1
fi

# 若项目自带 CUDA 库(WHISPER_GPU=1 打包)，加进 LD_LIBRARY_PATH，让 ctranslate2 找到(whisper GPU)
_nvlibs="$(ls -d "$HERE"/python/lib/python*/site-packages/nvidia/*/lib 2>/dev/null | tr '\n' ':')"
[ -n "$_nvlibs" ] && export LD_LIBRARY_PATH="${_nvlibs}${LD_LIBRARY_PATH:-}"

# 限制数值库线程：防止 whisper/numpy/openblas 每个都抢满核(多视频并发时互相踩)
export OMP_NUM_THREADS=2
export OPENBLAS_NUM_THREADS=2
export MKL_NUM_THREADS=2

echo "环境已激活: $(python --version 2>&1) | $HERE"
EOF
sed -i "s|__HF_ENDPOINT__|$HF_ENDPOINT_URL|" "$PROJECT_DIR/activate.sh"
chmod +x "$PROJECT_DIR/activate.sh"

# 提高文件句柄上限(高并发 + 海量小帧文件)，当前会话生效；永久需改 /etc/security/limits.conf
ulimit -n 65535 2>/dev/null || warn "ulimit 提升失败(非致命)，如需永久生效请改 /etc/security/limits.conf"


# =============================================================================
log "[7/7] 环境自检"
# =============================================================================
ok=1
check() { if eval "$2" >/dev/null 2>&1; then echo "  ✓ $1"; else echo "  ✗ $1"; ok=0; fi; }

check "gcc"            "gcc --version"
check "ffmpeg"         "ffmpeg -version"
check "ffprobe"        "ffprobe -version"
check "ffmpeg 支持 AV1" "ffmpeg -hide_banner -decoders | grep -iE 'dav1d|av1'"
check "uv"             "uv --version"
check "python"         "python --version"
check "import cv2"     "python -c 'import cv2'"
check "import numpy"   "python -c 'import numpy'"
check "import sklearn" "python -c 'import sklearn'"
check "import av(PyAV)" "python -c 'import av'"
check "import faster_whisper" "python -c 'import faster_whisper'"
check "import qcloud_cos"     "python -c 'import qcloud_cos'"

# GPU 检测(可选，有则抽帧能硬解加速)
if have nvidia-smi && nvidia-smi >/dev/null 2>&1; then
  echo "  ✓ 检测到 NVIDIA GPU(抽帧可硬解加速)"
else
  echo "  · 未检测到 GPU(抽帧用 CPU，正常)"
fi

echo
if [ "$ok" -eq 1 ]; then
  log "全部就绪 ✅"
else
  warn "有项目未通过(见上面 ✗)，请按提示排查"
fi

cat <<EOF

下一步：
  source activate.sh
  python main.py --source-list ./text2 --max-workers 12

EOF