#!/bin/bash
  
  ###############################################################################
  # CN-CLIP 图像编码服务启动脚本（增强版）
  #
  # 用途：
  #   启动 CN-CLIP 模型推理服务，用于图像和文本编码
  #
  # 使用方法：
  #   ./scripts/start_cnclip_service.sh [选项]
  #
  # 选项：
  #   --port PORT           服务端口（默认：51000）

  #   --device DEVICE       设备类型：cuda 或 cpu（默认：cuda）

  #   --batch-size SIZE     批处理大小（默认：32）
  #   --num-workers NUM     预处理线程数（默认：4）
  #   --dtype TYPE          数据类型：float16 或 float32（默认：float16）
  #   --model-name NAME     模型名称（默认：CN-CLIP/ViT-H-14）
  #   --replicas NUM        副本数（默认：1）
  #   --help                显示帮助信息
  #
  # 示例：
  #   ./scripts/start_cnclip_service.sh

  #   ./scripts/start_cnclip_service.sh --port 52000 --device cuda

  #   ./scripts/start_cnclip_service.sh --batch-size 16 --dtype float32
  #
  ###############################################################################
  
  set -e  # 遇到错误立即退出
  
  # 颜色定义
  RED='\033[0;31m'
  GREEN='\033[0;32m'
  YELLOW='\033[1;33m'
  BLUE='\033[0;34m'
  NC='\033[0m' # No Color
  
  # 默认配置
  DEFAULT_PORT=51000

  DEFAULT_DEVICE="cuda"

  DEFAULT_BATCH_SIZE=32
  DEFAULT_NUM_WORKERS=4
  DEFAULT_DTYPE="float16"
  DEFAULT_MODEL_NAME="CN-CLIP/ViT-H-14"

  # DEFAULT_MODEL_NAME="CN-CLIP/ViT-L-14-336"

  DEFAULT_REPLICAS=1  # 副本数
  

  # 项目路径（以仓库实际路径为准，避免写死 /data/tw/...）
  PROJECT_ROOT="$(cd "$(dirname "$0")/.." && pwd)"

  CLIP_SERVER_DIR="${PROJECT_ROOT}/third-party/clip-as-service/server"
  LOG_DIR="${PROJECT_ROOT}/logs"
  PID_FILE="${LOG_DIR}/cnclip_service.pid"
  LOG_FILE="${LOG_DIR}/cnclip_service.log"
  
  # 帮助信息
  show_help() {
      echo "CN-CLIP 图像编码服务启动脚本（增强版）"
      echo ""
      echo "用法: $0 [选项]"
      echo ""
      echo "选项:"
      echo "  --port PORT           服务端口（默认：${DEFAULT_PORT}）"

      echo "  --device DEVICE       设备类型：cuda 或 cpu（默认：cuda）"

      echo "  --batch-size SIZE     批处理大小（默认：${DEFAULT_BATCH_SIZE}）"
      echo "  --num-workers NUM     预处理线程数（默认：${DEFAULT_NUM_WORKERS}）"
      echo "  --dtype TYPE          数据类型：float16 或 float32（默认：${DEFAULT_DTYPE}）"
      echo "  --model-name NAME     模型名称（默认：${DEFAULT_MODEL_NAME}）"
      echo "  --replicas NUM        副本数（默认：${DEFAULT_REPLICAS}）"
      echo "  --help                显示此帮助信息"
      echo ""
      echo "示例:"
      echo "  $0                                          # 使用默认配置启动"

      echo "  $0 --port 52000 --device cuda               # 指定 CUDA 模式，端口 52000"
      echo "  $0 --port 52000 --device cpu                # 显式使用 CPU 模式"

      echo "  $0 --batch-size 16 --dtype float32          # 小批处理，float32 精度"
      echo "  $0 --replicas 2                            # 启动2个副本（需8-10GB显存）"
      echo ""
      echo "支持的模型:"
      echo "  - CN-CLIP/ViT-B-16      基础版本，速度快"
      echo "  - CN-CLIP/ViT-L-14      平衡版本"
      echo "  - CN-CLIP/ViT-L-14-336  高分辨率版本（默认）"
      echo "  - CN-CLIP/ViT-H-14      大型版本，精度高"
      echo "  - CN-CLIP/RN50          ResNet-50 版本"
  }
  
  # 解析命令行参数
  PORT=${DEFAULT_PORT}
  DEVICE=${DEFAULT_DEVICE}
  BATCH_SIZE=${DEFAULT_BATCH_SIZE}
  NUM_WORKERS=${DEFAULT_NUM_WORKERS}
  DTYPE=${DEFAULT_DTYPE}
  MODEL_NAME=${DEFAULT_MODEL_NAME}
  REPLICAS=${DEFAULT_REPLICAS}
  
  while [[ $# -gt 0 ]]; do
      case $1 in
          --port)
              PORT="$2"
              shift 2
              ;;
          --device)
              DEVICE="$2"
              shift 2
              ;;
          --batch-size)
              BATCH_SIZE="$2"
              shift 2
              ;;
          --num-workers)
              NUM_WORKERS="$2"
              shift 2
              ;;
          --dtype)
              DTYPE="$2"
              shift 2
              ;;
          --model-name)
              MODEL_NAME="$2"
              shift 2
              ;;
          --replicas)
              REPLICAS="$2"
              shift 2
              ;;
          --help)
              show_help
              exit 0
              ;;
          *)
              echo -e "${RED}错误: 未知参数 $1${NC}"
              show_help
              exit 1
              ;;
      esac
  done
  
  # 检查环境
  echo -e "${BLUE}========================================${NC}"
  echo -e "${BLUE}CN-CLIP 服务启动脚本${NC}"
  echo -e "${BLUE}========================================${NC}"
  echo ""
  
  # 检查项目目录
  if [ ! -d "${PROJECT_ROOT}" ]; then
      echo -e "${RED}错误: 项目目录不存在: ${PROJECT_ROOT}${NC}"
      exit 1
  fi
  
  # 检查 CLIP 服务目录
  if [ ! -d "${CLIP_SERVER_DIR}" ]; then
      echo -e "${RED}错误: CLIP 服务目录不存在: ${CLIP_SERVER_DIR}${NC}"
      exit 1
  fi
  
  # 创建日志目录
  mkdir -p "${LOG_DIR}"
  
  # 检查是否已经有服务在运行
  if [ -f "${PID_FILE}" ]; then
      OLD_PID=$(cat "${PID_FILE}")
      if ps -p ${OLD_PID} > /dev/null 2>&1; then
          echo -e "${YELLOW}警告: 服务已经在运行 (PID: ${OLD_PID})${NC}"
          echo -e "${YELLOW}请先运行 ./scripts/stop_cnclip_service.sh 停止服务${NC}"
          exit 1
      else
          echo -e "${YELLOW}清理旧的 PID 文件${NC}"
          rm -f "${PID_FILE}"
      fi
  fi
  
  # 检查端口是否被占用
  if lsof -Pi :${PORT} -sTCP:LISTEN -t >/dev/null 2>&1; then
      echo -e "${RED}错误: 端口 ${PORT} 已被占用${NC}"
      echo -e "${YELLOW}请检查是否有其他服务正在使用该端口${NC}"
      echo -e "${YELLOW}可以使用: lsof -i :${PORT} 查看占用情况${NC}"
      exit 1
  fi
  

  # 使用 CN-CLIP 专用环境（避免与主项目依赖冲突）

  CNCLIP_VENV="${PROJECT_ROOT}/.venv-cnclip"
  if [ -x "${CNCLIP_VENV}/bin/python" ]; then
      export PATH="${CNCLIP_VENV}/bin:${PATH}"
      export VIRTUAL_ENV="${CNCLIP_VENV}"
      echo -e "${GREEN}✓ 使用 CN-CLIP 专用环境: .venv-cnclip${NC}"

  else

      echo -e "${RED}错误: 未找到 CN-CLIP 专用环境 ${CNCLIP_VENV}${NC}"
      echo -e "${YELLOW}请先执行: ./scripts/setup_cnclip_venv.sh${NC}"
      exit 1

  fi
  

  # 检查 Python 依赖（CN-CLIP 服务端需要 cn_clip 与 clip_server）

  echo -e "${BLUE}检查 Python 依赖...${NC}"
  python -c "import cn_clip" 2>/dev/null || {
      echo -e "${RED}错误: cn_clip 未安装${NC}"

      echo -e "${YELLOW}请重建专用环境: ./scripts/setup_cnclip_venv.sh${NC}"

      exit 1
  }
  

  # clip_server 通过 PYTHONPATH 加载（见下方启动命令），此处仅做可导入性检查
  export PYTHONPATH="${CLIP_SERVER_DIR}:${PYTHONPATH}"
  python -c "import clip_server" 2>/dev/null || {
      echo -e "${RED}错误: clip_server 不可用${NC}"

      echo -e "${YELLOW}请重建专用环境: ./scripts/setup_cnclip_venv.sh${NC}"

      exit 1
  }
  

  echo -e "${GREEN}✓ 所有依赖已就绪${NC}"

  echo ""
  

  # 自动检测设备（可通过环境变量 CNCLIP_DEVICE 指定，供 service_ctl/restart 使用）

  if [ -n "${CNCLIP_DEVICE:-}" ]; then
      DEVICE="${CNCLIP_DEVICE}"
  fi
  
  DEVICE="$(echo "${DEVICE}" | tr '[:upper:]' '[:lower:]')"
  if [ "${DEVICE}" != "cuda" ] && [ "${DEVICE}" != "cpu" ]; then
      echo -e "${RED}错误: 不支持的 device=${DEVICE}，仅支持 cuda/cpu${NC}"
      exit 1
  fi
  if [ "${DEVICE}" == "cuda" ]; then
      if ! command -v nvidia-smi &> /dev/null || ! nvidia-smi &> /dev/null; then
          echo -e "${RED}错误: 已配置 --device cuda，但未检测到可用 NVIDIA GPU；禁止自动降级到 CPU${NC}"
          exit 1

      fi

      echo -e "${GREEN}✓ 设备: cuda（严格 GPU 模式，失败不降级）${NC}"

  else

      echo -e "${YELLOW}✓ 设备: cpu（显式配置）${NC}"

  fi
  
  # 显示配置信息
  echo -e "${BLUE}服务配置:${NC}"
  echo "  模型名称:     ${MODEL_NAME}"
  echo "  服务端口:     ${PORT}"
  echo "  协议:         gRPC (默认，官方推荐)"
  echo "  其他参数:     使用官方默认值"
  echo "  副本数:       ${REPLICAS}"
  echo "  日志文件:     ${LOG_FILE}"
  echo ""
  
  # 副本数显存警告
  if [ "${DEVICE}" == "cuda" ] && [ ${REPLICAS} -gt 1 ]; then
      ESTIMATED_MEMORY=$((REPLICAS * 5))
      echo -e "${YELLOW}⚠ 预计显存占用: ~${ESTIMATED_MEMORY}GB (${REPLICAS} 副本 × ~5GB/副本)${NC}"
      echo -e "${YELLOW}  请确保有足够的显存！${NC}"
      echo ""
  fi
  
  # 直接启动，不需要确认
  
  # 构建启动命令
  cd "${CLIP_SERVER_DIR}"
  
  # 设置环境变量
  export PYTHONPATH="${CLIP_SERVER_DIR}:${PYTHONPATH}"
  export NO_VERSION_CHECK=1  # 跳过版本检查
  
  # 启动服务
  echo -e "${BLUE}正在启动服务...${NC}"
  
  # 创建动态配置文件
  FLOW_FILE="${CLIP_SERVER_DIR}/torch-flow.yml"
  TEMP_FLOW_FILE="${CLIP_SERVER_DIR}/torch-flow-temp.yml"
  
  # 备份原配置文件
  if [ -f "${FLOW_FILE}" ] && [ ! -f "${FLOW_FILE}.original" ]; then
      cp "${FLOW_FILE}" "${FLOW_FILE}.original"
      echo -e "${YELLOW}已备份原配置文件: ${FLOW_FILE}.original${NC}"
  fi
  

  # 生成新的配置文件（使用官方默认配置，显式传入 device）

  cat > "${TEMP_FLOW_FILE}" << EOF
  jtype: Flow
  version: '1'
  with:
    port: ${PORT}
  executors:
    - name: clip_t
      uses:
        jtype: CLIPEncoder
        with:
          name: '${MODEL_NAME}'

          device: '${DEVICE}'

        metas:
          py_modules:
            - clip_server.executors.clip_torch
      timeout_ready: 3000000
      replicas: ${REPLICAS}
  EOF
  
  echo -e "${GREEN}✓ 已生成配置文件: ${TEMP_FLOW_FILE}${NC}"
  
  # 使用 nohup 在后台启动服务
  cd "${CLIP_SERVER_DIR}"
  nohup python -m clip_server "${TEMP_FLOW_FILE}" > "${LOG_FILE}" 2>&1 &
  
  # 保存 PID
  SERVICE_PID=$!
  echo ${SERVICE_PID} > "${PID_FILE}"
  
  # 等待服务启动
  echo -e "${YELLOW}等待服务启动...${NC}"
  sleep 5
  
  # 检查服务是否启动成功
  if ps -p ${SERVICE_PID} > /dev/null 2>&1; then
      echo -e "${GREEN}========================================${NC}"
      echo -e "${GREEN}✓ CN-CLIP 服务启动成功！${NC}"
      echo -e "${GREEN}========================================${NC}"
      echo ""
      echo -e "服务信息:"
      echo -e "  PID:          ${SERVICE_PID}"
      echo -e "  端口:         ${PORT}"
      echo -e "  模型:         ${MODEL_NAME}"
      echo -e "  设备:         ${DEVICE}"
      echo ""
      echo -e "测试服务 (使用 Python 客户端):"
      echo -e "  from clip_client import Client"
      echo -e "  c = Client('grpc://localhost:${PORT}')"
      echo -e "  r = c.encode(['测试文本'])"
      echo ""
      echo -e "查看日志:"
      echo -e "  tail -f ${LOG_FILE}"
      echo ""
      echo -e "停止服务:"
      echo -e "  ./scripts/stop_cnclip_service.sh"
      echo ""
  
      # 等待服务完全就绪（gRPC 协议无法用 curl 检查，等待固定时间）
      echo -e "${YELLOW}等待模型加载完成（约30-60秒）...${NC}"
      sleep 30
      echo -e "${GREEN}✓ 服务已启动，请查看日志确认模型是否加载完成${NC}"
      echo -e "${YELLOW}查看日志: tail -f ${LOG_FILE}${NC}"
  
  else
      echo -e "${RED}========================================${NC}"
      echo -e "${RED}✗ 服务启动失败！${NC}"
      echo -e "${RED}========================================${NC}"
      echo ""
      echo -e "请查看日志获取详细错误信息:"
      echo -e "  tail -f ${LOG_FILE}"
      echo ""
      rm -f "${PID_FILE}"
      exit 1
  fi