#!/bin/bash

# 定义变量
DB_FILE="$1"
DB_NAME="$(basename "$DB_FILE")"
DEFAULT_DB="$HOME/.local/share/milvus-lite/default.db"
SYSTEM_DB="/usr/share/milvus-lite/default.db"
ADDRESS="$2"
BIN_PATH="/opt/milvus-lite/bin"
LIB_PATH="/opt/milvus-lite/lib"
MILVUS_BIN="$BIN_PATH/milvus"
WORK_DIR="$(dirname "$DB_FILE")"
LOCK_FILE="${TMPDIR:-/tmp}/milvus-lite-$(basename "$DB_FILE")-$(id -u).lock"
UDS_PATH="${TMPDIR:-/tmp}/milvus-lite-$(id -u).sock"
LOG_LEVEL="${LOG_LEVEL:-ERROR}"

echo "Starting Milvus service..."
echo "DB_FILE: $DB_FILE"
if [ -n "$ADDRESS" ]; then
    echo "ADDRESS: $ADDRESS"
else
    echo "UDS_PATH: $UDS_PATH"
fi
echo "LOCK_FILE: $LOCK_FILE"
echo "LOG_LEVEL: $LOG_LEVEL"

# 检查DB_FILE是否符合要求
if ! [[ "$DB_NAME" =~ ^[a-zA-Z0-9.\-_]+$ ]]; then
    echo "Unsupport db name $(basename "$DB_FILE"), the name must match ^[a-zA-Z0-9.\-_]+$"
    exit 1
fi

if [ ${#DB_NAME} -gt 36 ]; then
    echo "Db name $(basename "$DB_FILE") is too long, should be less than 36"
    exit 1
fi

# 检查BIN_PATH是否存在
if [ ! -d "$BIN_PATH" ]; then
    echo "Bin path $BIN_PATH not exists"
    exit 1
fi

# 检查WORK_DIR是否存在
if [ ! -d "$WORK_DIR" ]; then
    echo "Dir $WORK_DIR not exist, creating it..."
    mkdir -p "$WORK_DIR"
    if [ $? -eq 0 ]; then
        echo "Dir $WORK_DIR created successfully."
    else
        echo "Failed to create dir $WORK_DIR."
    fi
fi

if [ ! -f "$DEFAULT_DB" ]; then
    echo "Default DB $DEFAULT_DB not found, copying from $SYSTEM_DB..."
    mkdir -p "$(dirname "$DEFAULT_DB")"
    cp "$SYSTEM_DB" "$DEFAULT_DB"
    if [ $? -eq 0 ]; then
        echo "Default DB copied successfully to $DEFAULT_DB."
    else
        echo "Failed to copy default DB from $SYSTEM_DB to $DEFAULT_DB."
        exit 1
    fi
fi

# 尝试获取锁
if ! (set -o noclobber && echo "locked" > "$LOCK_FILE") 2> /dev/null; then
    echo "Open $DB_FILE failed, the file has been opened by another program"
    exit 1
fi

# 设置环境变量
if [[ ":$LD_LIBRARY_PATH:" != *":$LIB_PATH:"* ]]; then
    export LD_LIBRARY_PATH=$LIB_PATH:$LD_LIBRARY_PATH
fi

# 启动Milvus服务
if [ -n "$ADDRESS" ]; then
    "$MILVUS_BIN" "$DB_FILE" "$ADDRESS" "$LOG_LEVEL" &
else
    "$MILVUS_BIN" "$DB_FILE" "unix:$UDS_PATH" "$LOG_LEVEL" &
fi

# 保存进程ID
PID=$!

# 注册信号处理函数
function cleanup {
    echo "Stopping Milvus service..."
    kill -TERM "$PID" 2>/dev/null
    wait "$PID" 2>/dev/null
    rm -f "$UDS_PATH" "$LOCK_FILE"
    echo "Milvus service stopped."
    exit 0
}

trap cleanup EXIT

# 等待服务停止
wait "$PID"
