#!/bin/bash
# =====================================================================
# 换发型服务 - 本机统一启动脚本
# 依赖顺序：webui(57860) → photo_service(32678) → hair_service_sd(8801)
#
# 用法：
#   bash start_all.sh start    # 启动全部服务
#   bash start_all.sh stop     # 停止全部服务
#   bash start_all.sh restart  # 重启
#   bash start_all.sh status   # 查看状态
# =====================================================================
set -u
BASE="/home/xsl/change_hair"
PROJ="$BASE/project"
LOGDIR="$PROJ/logs"
PIDD="$LOGDIR/pids"
mkdir -p "$LOGDIR" "$PIDD"

# conda 环境的 python 路径
PY_HAIR="/home/xsl/miniconda3/envs/my_hair/bin/python"
PY_SD="/home/xsl/miniconda3/envs/sdwebui/bin/python"
PY_PHOTO="/home/xsl/miniconda3/envs/py310/bin/python"

# 公共环境变量
export CRYPTOGRAPHY_OPENSSL_NO_LEGACY=1   # 旧版 cryptography 兼容
export CUDA_VISIBLE_DEVICES=0              # 单卡固定 0 号
export APP_WORKER_ID=1                     # hair_service 多 worker 识别

start_webui() {
  if pgrep -f "webui.py.*57860" >/dev/null; then echo "webui 已在运行，跳过"; return; fi
  echo "启动 webui (端口 57860)..."
  # HF 离线模式：使用本地缓存，避免访问 huggingface.co
  cd "$PROJ/onediff/stable-diffusion-webui"
  HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 \
  nohup "$PY_SD" webui.py --api --listen --xformers --port 57860 \
    > "$LOGDIR/webui.log" 2>&1 &
  echo $! > "$PIDD/webui.pid"
  echo "  webui PID=$(cat $PIDD/webui.pid)"
}

start_photo() {
  if pgrep -f "lora_train_service_1.py" >/dev/null; then echo "photo_service 已在运行，跳过"; return; fi
  echo "启动 photo_service (端口 32678)..."
  cd "$PROJ/photo_service"
  nohup "$PY_PHOTO" -u lora_train_service_1.py \
    > "$LOGDIR/photo_service.log" 2>&1 &
  echo $! > "$PIDD/photo_service.pid"
  echo "  photo_service PID=$(cat $PIDD/photo_service.pid)"
}

start_hair() {
  if pgrep -f "run_copy_cost_colorb64" >/dev/null; then echo "hair_service_sd 已在运行，跳过"; return; fi
  echo "启动 hair_service_sd (端口 8801)..."
  cd "$PROJ/hair_service_sd"
  nohup "$PY_HAIR" run_copy_cost_colorb64.py \
    > "$LOGDIR/hair_service.log" 2>&1 &
  echo $! > "$PIDD/hair_service.pid"
  echo "  hair_service_sd PID=$(cat $PIDD/hair_service.pid)"
}

start_all() {
  echo "=== 按依赖顺序启动 $(date '+%Y-%m-%d %H:%M:%S') ==="
  start_webui
  echo "  等待 webui 加载模型(约1-2分钟)..."
  start_photo
  start_hair
  echo "=== 启动命令已发送 ==="
  echo "提示：模型加载需要时间，约2分钟后用 'bash start_all.sh status' 确认就绪"
}

stop_one() {
  local label="$1" pattern="$2"
  pkill -f "$pattern" 2>/dev/null
  echo "  已停止 $label"
}

stop_all() {
  echo "=== 停止全部服务 ==="
  stop_one "hair_service_sd" "run_copy_cost_colorb64"
  stop_one "photo_service"   "lora_train_service_1.py"
  stop_one "webui"           "webui.py.*57860"
  sleep 2
  echo "=== 已停止 ==="
}

status_all() {
  echo "=== 服务状态 $(date '+%Y-%m-%d %H:%M:%S') ==="
  check() {
    local label="$1" pattern="$2" port="$3"
    if pgrep -f "$pattern" >/dev/null 2>&1; then
      local pid=$(pgrep -f "$pattern" | head -1)
      if ss -tln 2>/dev/null | grep -q ":$port "; then
        echo "  ✓ $label: 运行中 (PID=$pid, 端口 $port 就绪)"
      else
        echo "  ⏳ $label: 运行中 (PID=$pid, 端口 $port 未就绪，可能仍在加载)"
      fi
    else
      echo "  ✗ $label: 未运行"
    fi
  }
  check "webui"            "webui.py.*57860"         "57860"
  check "photo_service"    "lora_train_service_1.py" "32678"
  check "hair_service_sd"  "run_copy_cost_colorb64"  "8801"
}

case "${1:-status}" in
  start)   start_all ;;
  stop)    stop_all ;;
  restart) stop_all; sleep 3; start_all ;;
  status)  status_all ;;
  *)
    echo "用法: bash start_all.sh {start|stop|restart|status}"
    echo ""
    echo "服务端口："
    echo "  webui           : 57860 (SD模型 + ControlNet)"
    echo "  photo_service   : 32678 (LoRA加载 + webui调度)"
    echo "  hair_service_sd : 8801  (主服务：换发色/换发型)"
    echo ""
    echo "接口："
    echo "  POST http://127.0.0.1:8801/hairColor/v2     换发色"
    echo "  POST http://127.0.0.1:8801/api/swapHair/v1  换发型"
    echo "  POST http://127.0.0.1:8801/api/uploadHair/v1 新增发型(需Docker,暂不可用)"
    exit 1 ;;
esac