#!/bin/bash

###############################################################################
# CN-CLIP 图像编码服务启动脚本（增强版）
#
# 用途：
#   启动 CN-CLIP 模型推理服务，用于图像和文本编码
#
# 使用方法：
#   ./scripts/start_cnclip_service.sh [选项]
#
# 选项：
#   --port PORT           服务端口（默认：51000）
#   --device DEVICE       设备类型：cuda 或 cpu（默认：cuda）
#   --model-name NAME     模型名称（默认读取 embeddings/config.py）
#   --replicas NUM        副本数（默认：1）
#   --help                显示帮助信息
#
# 示例：
#   ./scripts/start_cnclip_service.sh
#   ./scripts/start_cnclip_service.sh --port 52000 --device cuda
#
###############################################################################

set -euo pipefail

# 颜色定义
RED='\033[0;31m'
GREEN='\033[0;32m'
YELLOW='\033[1;33m'
BLUE='\033[0;34m'
NC='\033[0m' # No Color

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

resolve_default_model_name() {
    local python_bin
    local resolved_model_name
    for python_bin in python3 python; do
        if command -v "${python_bin}" >/dev/null 2>&1; then
            if resolved_model_name="$(PYTHONPATH="${PROJECT_ROOT}${PYTHONPATH:+:${PYTHONPATH}}" "${python_bin}" -c "from embeddings.config import CONFIG; print(CONFIG.CLIP_AS_SERVICE_MODEL_NAME)" 2>/dev/null)"; then
                if [ -n "${resolved_model_name}" ]; then
                    echo "${resolved_model_name}"
                    return 0
                fi
            fi
        fi
    done
    echo "CN-CLIP/ViT-L-14"
}

# 默认配置
DEFAULT_PORT=51000
DEFAULT_DEVICE="cuda"
DEFAULT_MODEL_NAME="$(resolve_default_model_name)"
DEFAULT_REPLICAS=1  # 副本数

CLIP_SERVER_DIR="${PROJECT_ROOT}/third-party/clip-as-service/server"
LOG_DIR="${PROJECT_ROOT}/logs"
PID_FILE="${LOG_DIR}/cnclip.pid"
LOG_LINK="${LOG_DIR}/cnclip.log"
LOG_FILE="${LOG_DIR}/cnclip-$(date +%F).log"
LOG_ROUTER_SCRIPT="${PROJECT_ROOT}/scripts/daily_log_router.sh"

# 帮助信息
show_help() {
    echo "CN-CLIP 图像编码服务启动脚本（增强版）"
    echo ""
    echo "用法: $0 [选项]"
    echo ""
    echo "选项:"
    echo "  --port PORT           服务端口（默认：${CNCLIP_PORT:-${DEFAULT_PORT}}）"
    echo "  --device DEVICE       设备类型：cuda 或 cpu（默认：cuda）"
    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 --model-name CN-CLIP/ViT-L-14            # 临时覆盖模型"
    echo "  $0 --replicas 2                            # 启动2个副本（需8-10GB显存）"
    echo ""
    echo "说明:"
    echo "  - 默认模型取自 embeddings/config.py 的 CLIP_AS_SERVICE_MODEL_NAME"
    echo "  - 也可通过环境变量 CNCLIP_MODEL_NAME 覆盖，再由 --model-name 最终覆盖"
    echo ""
    echo "支持的模型:"
    local supported_models=(
        "CN-CLIP/ViT-B-16|基础版本，速度快"
        "CN-CLIP/ViT-L-14|平衡版本"
        "CN-CLIP/ViT-L-14-336|高分辨率版本"
        "CN-CLIP/ViT-H-14|大型版本，精度高"
        "CN-CLIP/RN50|ResNet-50 版本"
    )
    local item model desc suffix
    for item in "${supported_models[@]}"; do
        model="${item%%|*}"
        desc="${item#*|}"
        suffix=""
        if [ "${model}" = "${DEFAULT_MODEL_NAME}" ]; then
            suffix="（当前默认）"
        fi
        echo "  - ${model}  ${desc}${suffix}"
    done
}

# 解析命令行参数
PORT="${CNCLIP_PORT:-${DEFAULT_PORT}}"
DEVICE=${DEFAULT_DEVICE}
MODEL_NAME="${CNCLIP_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
            ;;
        --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 [ ! -x "${LOG_ROUTER_SCRIPT}" ]; then
    echo -e "${RED}错误: 日志路由脚本不存在或不可执行: ${LOG_ROUTER_SCRIPT}${NC}"
    exit 1
fi
ln -sfn "$(basename "${LOG_FILE}")" "${LOG_LINK}"

# 检查是否已经有服务在运行
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:+:${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_LINK}"
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:+:${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 bash -lc "exec python -m clip_server \"${TEMP_FLOW_FILE}\" > >(\"${LOG_ROUTER_SCRIPT}\" \"cnclip\" \"${LOG_DIR}\" \"${LOG_RETENTION_DAYS:-30}\") 2>&1" >/dev/null 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_LINK}"
    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_LINK}${NC}"

else
    echo -e "${RED}========================================${NC}"
    echo -e "${RED}✗ 服务启动失败！${NC}"
    echo -e "${RED}========================================${NC}"
    echo ""
    echo -e "请查看日志获取详细错误信息:"
    echo -e "  tail -F ${LOG_LINK}"
    echo ""
    rm -f "${PID_FILE}"
    exit 1
fi