#!/bin/bash

###############################################################################
# CN-CLIP 图像编码服务启动脚本（增强版）
#
# 用途：
#   启动 CN-CLIP 模型推理服务，用于图像和文本编码
#
# 使用方法：
#   ./scripts/start_cnclip_service.sh [选项]
#
# 选项：
#   --port PORT           服务端口（默认：51000）
#   --device DEVICE       设备类型：cuda 或 cpu（默认：自动检测）
#   --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 cpu
#   ./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="auto"  # auto, cuda, cpu
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  # 副本数

# 项目路径
PROJECT_ROOT="/data/tw/SearchEngine"
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（默认：自动检测）"
    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 cpu                # 使用 CPU 模式，端口 52000"
    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

# 检查 conda 环境
if [ -z "${CONDA_DEFAULT_ENV}" ] || [ "${CONDA_DEFAULT_ENV}" != "clip_service" ]; then
    echo -e "${YELLOW}警告: 当前未激活 clip_service 环境${NC}"
    echo -e "${YELLOW}正在激活环境...${NC}"

    if [ -f "/home/tw/miniconda3/etc/profile.d/conda.sh" ]; then
        source "/home/tw/miniconda3/etc/profile.d/conda.sh"
        conda activate clip_service
        echo -e "${GREEN}✓ 环境已激活${NC}"
    else
        echo -e "${RED}错误: 无法找到 conda 初始化脚本${NC}"
        exit 1
    fi
else
    echo -e "${GREEN}✓ Conda 环境: ${CONDA_DEFAULT_ENV}${NC}"
fi

# 检查 Python 依赖
echo -e "${BLUE}检查 Python 依赖...${NC}"
python -c "import cn_clip" 2>/dev/null || {
    echo -e "${RED}错误: cn_clip 未安装${NC}"
    echo -e "${YELLOW}请运行: pip install cn-clip${NC}"
    exit 1
}

python -c "from clip_client import Client" 2>/dev/null || {
    echo -e "${RED}错误: clip_client 未安装${NC}"
    echo -e "${YELLOW}请运行: pip install clip-client${NC}"
    exit 1
}

echo -e "${GREEN}✓ 所有依赖已安装${NC}"
echo ""

# 自动检测设备
if [ "${DEVICE}" == "auto" ]; then
    if command -v nvidia-smi &> /dev/null && nvidia-smi &> /dev/null; then
        DEVICE="cuda"
        echo -e "${GREEN}✓ 检测到 NVIDIA GPU，使用 CUDA${NC}"
    else
        DEVICE="cpu"
        echo -e "${YELLOW}未检测到 GPU，使用 CPU${NC}"
    fi
else
    echo -e "${GREEN}✓ 设备: ${DEVICE}${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

# 生成新的配置文件（使用官方默认配置，只指定模型名称）
cat > "${TEMP_FLOW_FILE}" << EOF
jtype: Flow
version: '1'
with:
  port: ${PORT}
executors:
  - name: clip_t
    uses:
      jtype: CLIPEncoder
      with:
        name: '${MODEL_NAME}'
      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