#!/usr/bin/env bash
# 启动 Label Studio 的 ML backend（预标注 + 触发重训）。
#
#   bash scripts/start_ml_backend.sh
#
# 会先停掉旧进程再起新的。所有参数都可以用环境变量覆盖，例如换权重：
#
#   YOLO_WEIGHTS=/path/to/best.pt bash scripts/start_ml_backend.sh
#
# 权重只在进程启动时加载一次，重训完必须重跑本脚本才会生效。
set -euo pipefail

REPO=$(cd "$(dirname "${BASH_SOURCE[0]}")/.." && pwd)

# 默认用 runs/detect/current.pt 这个软链接：重训脚本训完会把它指向最新权重，
# 这样换模型不用改脚本。链接不存在时退回第一版训练的结果。
DEFAULT_WEIGHTS="$REPO/runs/detect/current.pt"
[ -e "$DEFAULT_WEIGHTS" ] || DEFAULT_WEIGHTS="$REPO/runs/detect/film/weights/best.pt"

export YOLO_WEIGHTS="${YOLO_WEIGHTS:-$DEFAULT_WEIGHTS}"
export YOLO_DEVICE="${YOLO_DEVICE:-cuda}"
export YOLO_CONF="${YOLO_CONF:-0.25}"
export YOLO_IMGSZ="${YOLO_IMGSZ:-640}"

# 必须与运行 Label Studio 时的值一致，否则按 ?d=... 还原不出图片路径
export LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT="${LABEL_STUDIO_LOCAL_FILES_DOCUMENT_ROOT:-/home/deep/labeling}"

# 点 Start Training 无条件触发；攒够 RETRAIN_AFTER 条新标注也会自动触发（0=只手动）
export RETRAIN_AFTER="${RETRAIN_AFTER:-0}"
export RETRAIN_COMMAND="${RETRAIN_COMMAND:-$REPO/scripts/retrain_from_label_studio.sh}"
export RETRAIN_LOG="${RETRAIN_LOG:-/tmp/retrain.log}"

PORT="${LS_ML_PORT:-9090}"
LOG="${ML_BACKEND_LOG:-/tmp/mlb.log}"
PYTHON="${ML_BACKEND_PYTHON:-$REPO/.venv-ml/bin/python}"

if [ ! -x "$PYTHON" ]; then
    echo "找不到 $PYTHON。先建独立环境（见 ml_backend/README.md）：" >&2
    echo "  uv venv .venv-ml --python 3.12" >&2
    exit 1
fi
if [ ! -e "$YOLO_WEIGHTS" ]; then
    echo "找不到权重 $YOLO_WEIGHTS" >&2
    exit 1
fi

# 模式里不要出现 _wsgi.py 的完整写法：ssh 单行执行时会匹配到命令自身，把会话一起杀掉
pkill -f "wsgi" 2>/dev/null || true
sleep 1

cd "$REPO/ml_backend"
setsid nohup "$PYTHON" ./_wsgi.py -p "$PORT" > "$LOG" 2>&1 < /dev/null &
echo "ML backend 已启动 pid=$!"
echo "  权重     : $YOLO_WEIGHTS"
echo "  重训命令 : $RETRAIN_COMMAND (RETRAIN_AFTER=$RETRAIN_AFTER)"
echo "  服务日志 : $LOG"
echo "  重训日志 : $RETRAIN_LOG"
echo "  网页端 URL: http://localhost:$PORT"