#!/bin/bash

# ==============================================================================
# 无人机自然语言控制项目 - vLLM 一键启动脚本（新增，不修改 start_all.sh）
# ==============================================================================
# 功能：启动所有必需的服务（vLLM推理模型、llama-server embedding模型、FastAPI后端）
# 用法：./start_all_vllm.sh [选项]
# ==============================================================================

set -e  # 遇到错误立即退出

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

# 默认配置（可通过环境变量覆盖）
PROJECT_ROOT="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
LLAMA_SERVER_DIR="${LLAMA_SERVER_DIR:-~/llama.cpp/build/bin}"
INFERENCE_MODEL_DIR="${INFERENCE_MODEL_DIR:-/home/a/models/AWQ/Qwen3-4B-Instruct-2507}"
VLLM_CONDA_ENV="${VLLM_CONDA_ENV:-VLLM}"
VLLM_SPECULATIVE_CONFIG="${VLLM_SPECULATIVE_CONFIG:-}"
VLLM_GPU_MEMORY_UTILIZATION="${VLLM_GPU_MEMORY_UTILIZATION:-0.5}"
VLLM_SERVED_MODEL_NAME="${VLLM_SERVED_MODEL_NAME:-local-model}"
VLLM_ENABLE_TOOL_CALLING="${VLLM_ENABLE_TOOL_CALLING:-0}"
VLLM_TOOL_CALL_PARSER="${VLLM_TOOL_CALL_PARSER:-}"
EMBEDDING_MODEL="${EMBEDDING_MODEL:-~/models/gguf/Qwen3/Qwen3-Embedding-4B/Qwen3-Embedding-4B-Q5_K_M.gguf}"
VENV_PATH="${VENV_PATH:-${PROJECT_ROOT}/backend_service/venv}"
LOG_DIR="${LOG_DIR:-${PROJECT_ROOT}/logs}"
PID_FILE="${LOG_DIR}/services_vllm.pid"

# 端口配置
INFERENCE_PORT="${INFERENCE_PORT:-8081}"
EMBEDDING_PORT="${EMBEDDING_PORT:-8090}"
API_PORT="${API_PORT:-8000}"

# 创建日志目录
mkdir -p "${LOG_DIR}"

# ==============================================================================
# 辅助函数
# ==============================================================================

print_info() {
    echo -e "${BLUE}[INFO]${NC} $1"
}

print_success() {
    echo -e "${GREEN}[SUCCESS]${NC} $1"
}

print_warning() {
    echo -e "${YELLOW}[WARNING]${NC} $1"
}

print_error() {
    echo -e "${RED}[ERROR]${NC} $1"
}

# 检查命令是否存在
check_command() {
    if ! command -v "$1" &> /dev/null; then
        print_error "$1 命令未找到，请先安装"
        return 1
    fi
    return 0
}

# 检查端口是否被占用
check_port() {
    local port=$1
    if lsof -Pi :${port} -sTCP:LISTEN -t >/dev/null 2>&1 ; then
        return 0  # 端口被占用
    else
        return 1  # 端口空闲
    fi
}

# 等待服务就绪
wait_for_service() {
    local url=$1
    local service_name=$2
    local max_attempts=30
    local attempt=0
    
    print_info "等待 ${service_name} 启动..."
    while [ $attempt -lt $max_attempts ]; do
        if curl -s "${url}" > /dev/null 2>&1; then
            print_success "${service_name} 已就绪"
            return 0
        fi
        attempt=$((attempt + 1))
        sleep 1
    done
    
    print_error "${service_name} 启动超时"
    return 1
}

# 停止所有服务
stop_services() {
    print_info "正在停止所有服务..."
    
    if [ -f "${PID_FILE}" ]; then
        while read -r pid; do
            if ps -p "$pid" > /dev/null 2>&1; then
                print_info "停止进程 PID: $pid"
                kill "$pid" 2>/dev/null || true
            fi
        done < "${PID_FILE}"
        rm -f "${PID_FILE}"
    fi
    
    # 尝试通过端口停止服务
    for port in ${INFERENCE_PORT} ${EMBEDDING_PORT} ${API_PORT}; do
        if check_port ${port}; then
            local pid
            pid=$(lsof -ti:${port})
            if [ -n "$pid" ]; then
                print_info "停止占用端口 ${port} 的进程 (PID: $pid)"
                kill $pid 2>/dev/null || true
            fi
        fi
    done
    
    print_success "所有服务已停止"
}

# 清理函数（脚本退出时调用）
cleanup() {
    if [ "$?" -ne 0 ]; then
        print_error "启动过程中发生错误，正在清理..."
    fi
    # 注意：这里不自动停止服务，让用户手动控制
}

trap cleanup EXIT

# ==============================================================================
# 主函数
# ==============================================================================

start_services() {
    print_info "=========================================="
    print_info "  无人机自然语言控制项目 - vLLM 服务启动"
    print_info "=========================================="
    echo ""
    
    # 检查必要的命令
    print_info "检查必要的命令..."
    check_command "python3" || exit 1
    check_command "curl" || exit 1
    check_command "lsof" || print_warning "lsof 未安装，将无法检查端口占用"
    check_command "conda" || exit 1
    echo ""
    
    # 检查端口占用
    print_info "检查端口占用..."
    if check_port ${INFERENCE_PORT}; then
        print_warning "端口 ${INFERENCE_PORT} 已被占用，推理模型可能已在运行"
    fi
    if check_port ${EMBEDDING_PORT}; then
        print_warning "端口 ${EMBEDDING_PORT} 已被占用，Embedding模型可能已在运行"
    fi
    if check_port ${API_PORT}; then
        print_error "端口 ${API_PORT} 已被占用，请先停止占用该端口的服务"
        exit 1
    fi
    echo ""
    
    # 检查llama-server（用于 embedding，展开路径中的 ~）
    local llama_server_dir_expanded
    llama_server_dir_expanded=$(eval echo "${LLAMA_SERVER_DIR}")
    local llama_server="${llama_server_dir_expanded}/llama-server"
    if [ ! -f "${llama_server}" ]; then
        print_error "llama-server 未找到: ${llama_server}"
        print_info "请设置 LLAMA_SERVER_DIR 环境变量指向正确的路径"
        print_info "当前路径: ${LLAMA_SERVER_DIR}"
        print_info "展开后路径: ${llama_server_dir_expanded}"
        exit 1
    fi
    print_success "找到 llama-server: ${llama_server}"
    echo ""
    
    # 检查模型路径
    local inference_model_dir_expanded
    inference_model_dir_expanded=$(eval echo "${INFERENCE_MODEL_DIR}")
    local embedding_model_expanded
    embedding_model_expanded=$(eval echo "${EMBEDDING_MODEL}")
    
    if [ ! -e "${inference_model_dir_expanded}" ]; then
        print_error "推理模型目录未找到: ${inference_model_dir_expanded}"
        print_info "请设置 INFERENCE_MODEL_DIR 环境变量指向正确的模型目录"
        exit 1
    fi
    if [ ! -d "${inference_model_dir_expanded}" ]; then
        print_warning "推理模型路径不是目录: ${inference_model_dir_expanded}"
        print_warning "vLLM 通常需要传入模型目录（包含 config/tokenizer/权重等）"
    fi
    print_success "找到推理模型目录: ${inference_model_dir_expanded}"
    
    if [ ! -f "${embedding_model_expanded}" ]; then
        print_error "Embedding模型文件未找到: ${embedding_model_expanded}"
        print_info "请设置 EMBEDDING_MODEL 环境变量指向正确的模型路径"
        exit 1
    fi
    print_success "找到Embedding模型: ${embedding_model_expanded}"
    echo ""
    
    # 检查ROS2环境
    local ros2_setup="${PROJECT_ROOT}/install/setup.bash"
    if [ ! -f "${ros2_setup}" ]; then
        print_warning "ROS2 setup文件未找到: ${ros2_setup}"
        print_warning "如果项目已与ROS2解耦，可以忽略此警告"
    else
        print_success "找到ROS2 setup文件: ${ros2_setup}"
    fi
    echo ""
    
    # 检查venv虚拟环境
    local venv_path_expanded
    venv_path_expanded=$(eval echo "${VENV_PATH}")
    print_info "检查venv虚拟环境: ${venv_path_expanded}"
    if [ ! -d "${venv_path_expanded}" ]; then
        print_error "venv虚拟环境目录不存在: ${venv_path_expanded}"
        print_info "请先创建venv环境: python3 -m venv ${venv_path_expanded}"
        print_info "然后安装依赖: ${venv_path_expanded}/bin/pip install -r backend_service/requirements.txt"
        exit 1
    fi
    if [ ! -f "${venv_path_expanded}/bin/activate" ]; then
        print_error "venv激活脚本不存在: ${venv_path_expanded}/bin/activate"
        print_error "这看起来不是一个有效的venv环境"
        exit 1
    fi
    print_success "venv虚拟环境存在: ${venv_path_expanded}"
    echo ""
    
    # 初始化PID文件
    > "${PID_FILE}"
    
    # ==========================================================================
    # 启动推理模型服务（vLLM）
    # ==========================================================================
    print_info "启动推理模型服务(vLLM) (端口 ${INFERENCE_PORT})..."
    bash -lc "
        set -e
        CONDA_BASE=\$(conda info --base)
        if [ -z \"\${CONDA_BASE}\" ] || [ ! -f \"\${CONDA_BASE}/etc/profile.d/conda.sh\" ]; then
            echo \"[ERROR] conda.sh 未找到：\${CONDA_BASE}/etc/profile.d/conda.sh\" 1>&2
            exit 1
        fi
        source \"\${CONDA_BASE}/etc/profile.d/conda.sh\"
        conda activate \"${VLLM_CONDA_ENV}\"
        export VLLM_USE_V1=1
        echo \"[DEBUG] conda_env=${VLLM_CONDA_ENV}\"
        echo \"[DEBUG] CONDA_PREFIX=\${CONDA_PREFIX}\"
        echo \"[DEBUG] which python: \$(which python)\"
        echo \"[DEBUG] which vllm: \$(which vllm)\"
        vllm --version || true
        python -c \"import sys; print('[DEBUG] python:', sys.version)\" || true
        python -c \"import torch; print('[DEBUG] torch:', torch.__version__, 'cuda:', torch.version.cuda, 'available:', torch.cuda.is_available())\" || true

        SPEC_ARGS=()
        if [ -n \"${VLLM_SPECULATIVE_CONFIG}\" ]; then
            SPEC_ARGS+=(--speculative-config \"${VLLM_SPECULATIVE_CONFIG}\")
        fi
        TOOL_ARGS=()
        if [ \"${VLLM_ENABLE_TOOL_CALLING}\" = \"1\" ]; then
            TOOL_ARGS+=(--enable-auto-tool-choice)
            if [ -n \"${VLLM_TOOL_CALL_PARSER}\" ]; then
                TOOL_ARGS+=(--tool-call-parser \"${VLLM_TOOL_CALL_PARSER}\")
            fi
        fi
        exec vllm serve \"${inference_model_dir_expanded}\" \
          --host=0.0.0.0 \
          --port=${INFERENCE_PORT} \
          --dtype=auto \
          --max-num-seqs=1 \
          --max-model-len=16384 \
          \"\${SPEC_ARGS[@]}\" \
          \"\${TOOL_ARGS[@]}\" \
          --served-model-name \"${VLLM_SERVED_MODEL_NAME}\" \
          --trust-remote-code \
          --gpu-memory-utilization=${VLLM_GPU_MEMORY_UTILIZATION} \
          --uvicorn-log-level=debug \
          --max-parallel-loading-workers 4
    " > "${LOG_DIR}/vllm_inference_model.log" 2>&1 &
    local inference_pid=$!
    echo $inference_pid >> "${PID_FILE}"
    print_success "推理模型服务(vLLM)已启动 (PID: $inference_pid)"
    print_info "日志文件: ${LOG_DIR}/vllm_inference_model.log"
    echo ""
    
    # ==========================================================================
    # 启动Embedding模型服务（llama-server）
    # ==========================================================================
    print_info "启动Embedding模型服务 (端口 ${EMBEDDING_PORT})..."
    cd "${llama_server_dir_expanded}"
    nohup ./llama-server \
        -m "${embedding_model_expanded}" \
        --gpu-layers 36 \
        --port ${EMBEDDING_PORT} \
        --embeddings \
        --pooling last \
        --host 0.0.0.0 \
        > "${LOG_DIR}/embedding_model.log" 2>&1 &
    local embedding_pid=$!
    echo $embedding_pid >> "${PID_FILE}"
    print_success "Embedding模型服务已启动 (PID: $embedding_pid)"
    print_info "日志文件: ${LOG_DIR}/embedding_model.log"
    echo ""
    
    # ==========================================================================
    # 等待模型服务就绪
    # ==========================================================================
    print_info "等待模型服务就绪..."
    sleep 3  # 给服务一些启动时间
    
    # 等待推理模型服务（vLLM）
    if ! wait_for_service "http://localhost:${INFERENCE_PORT}/health" "推理模型服务(vLLM)"; then
        if ! wait_for_service "http://localhost:${INFERENCE_PORT}/v1/models" "推理模型服务(vLLM)"; then
            print_warning "推理模型服务(vLLM)可能未完全就绪，但将继续启动"
        fi
    fi
    
    # 等待Embedding模型服务
    if ! wait_for_service "http://localhost:${EMBEDDING_PORT}/health" "Embedding模型服务"; then
        if ! wait_for_service "http://localhost:${EMBEDDING_PORT}/v1/models" "Embedding模型服务"; then
            print_warning "Embedding模型服务可能未完全就绪，但将继续启动"
        fi
    fi
    echo ""
    
    # ==========================================================================
    # 启动FastAPI后端服务
    # ==========================================================================
    print_info "启动FastAPI后端服务 (端口 ${API_PORT})..."
    cd "${PROJECT_ROOT}"
    
    # 尝试找到 dot 命令 (Graphviz) 路径，以便后端生成图片
    DOT_PATH=$(which dot || true)
    DOT_DIR=""
    if [ -n "$DOT_PATH" ]; then
        DOT_DIR=$(dirname "$DOT_PATH")
        print_info "找到 Graphviz dot 命令: $DOT_PATH"
    else
        print_warning "未找到 dot 命令，可视化图片生成可能会失败"
        # 尝试默认的 miniconda 路径作为备选
        if [ -f "/home/a/miniconda3/bin/dot" ]; then
             DOT_DIR="/home/a/miniconda3/bin"
             print_info "使用备选路径: $DOT_DIR"
        fi
    fi

    bash -c "
        # 激活ROS2环境（如果存在）
        if [ -f '${ros2_setup}' ]; then
            source '${ros2_setup}'
        fi
        
        # 确保 Graphviz 在 PATH 中
        if [ -n '${DOT_DIR}' ]; then
            export PATH=\"${DOT_DIR}:\$PATH\"
        fi

        # 让后端默认使用 vLLM 暴露的模型名，避免请求 local-model 导致 404
        OPENAI_MODEL=\"\${OPENAI_MODEL:-${VLLM_SERVED_MODEL_NAME}}\" \
        CLASSIFIER_MODEL=\"\${CLASSIFIER_MODEL:-${VLLM_SERVED_MODEL_NAME}}\" \
        SIMPLE_MODEL=\"\${SIMPLE_MODEL:-${VLLM_SERVED_MODEL_NAME}}\" \
        COMPLEX_MODEL=\"\${COMPLEX_MODEL:-${VLLM_SERVED_MODEL_NAME}}\" \
        cd '${PROJECT_ROOT}/backend_service' && \
        '${venv_path_expanded}/bin/python' -m uvicorn src.main:app --host 0.0.0.0 --port ${API_PORT}
    " > "${LOG_DIR}/fastapi.log" 2>&1 &
    local api_pid=$!
    echo $api_pid >> "${PID_FILE}"
    print_success "FastAPI服务已启动 (PID: $api_pid)"
    print_info "日志文件: ${LOG_DIR}/fastapi.log"
    echo ""
    
    # 等待FastAPI服务就绪
    sleep 3
    if wait_for_service "http://localhost:${API_PORT}/docs" "FastAPI服务"; then
        print_success "所有服务已成功启动！"
    else
        print_warning "FastAPI服务可能未完全就绪，请检查日志: ${LOG_DIR}/fastapi.log"
    fi
    echo ""
    
    # 显示服务访问信息
    print_info "=========================================="
    print_info "  服务启动完成！"
    print_info "=========================================="
    print_info "推理模型API(vLLM): http://localhost:${INFERENCE_PORT}/v1"
    print_info "Embedding模型API: http://localhost:${EMBEDDING_PORT}/v1"
    print_info "FastAPI后端: http://localhost:${API_PORT}"
    print_info "API文档: http://localhost:${API_PORT}/docs"
    print_info ""
    print_info "日志文件位置:"
    print_info "  - 推理模型(vLLM): ${LOG_DIR}/vllm_inference_model.log"
    print_info "  - Embedding模型: ${LOG_DIR}/embedding_model.log"
    print_info "  - FastAPI服务: ${LOG_DIR}/fastapi.log"
    print_info ""
    print_info "按 Ctrl+C 停止所有服务"
    print_info "=========================================="
    echo ""
    
    # 设置信号处理，确保Ctrl+C时能清理
    trap 'print_info "\n正在停止服务..."; stop_services; exit 0' INT TERM
    
    # 保持脚本运行
    print_info "所有服务正在运行中，查看日志请使用:"
    print_info "  tail -f ${LOG_DIR}/*.log"
    echo ""
    
    wait
}

# ==============================================================================
# 脚本入口
# ==============================================================================

case "${1:-start}" in
    start)
        start_services
        ;;
    stop)
        stop_services
        ;;
    restart)
        stop_services
        sleep 2
        start_services
        ;;
    status)
        print_info "检查服务状态..."
        if [ -f "${PID_FILE}" ]; then
            print_info "已记录的服务进程:"
            while read -r pid; do
                if ps -p "$pid" > /dev/null 2>&1; then
                    print_success "PID $pid: 运行中"
                else
                    print_warning "PID $pid: 已停止"
                fi
            done < "${PID_FILE}"
        else
            print_info "未找到PID文件，服务可能未启动"
        fi
        echo ""
        print_info "端口占用情况:"
        for port in ${INFERENCE_PORT} ${EMBEDDING_PORT} ${API_PORT}; do
            if check_port ${port}; then
                local pid
                pid=$(lsof -ti:${port})
                print_success "端口 ${port}: 被占用 (PID: $pid)"
            else
                print_warning "端口 ${port}: 空闲"
            fi
        done
        ;;
    *)
        echo "用法: $0 {start|stop|restart|status}"
        echo ""
        echo "命令说明:"
        echo "  start   - 启动所有服务（默认）"
        echo "  stop    - 停止所有服务"
        echo "  restart - 重启所有服务"
        echo "  status  - 查看服务状态"
        echo ""
        echo "环境变量配置:"
        echo "  LLAMA_SERVER_DIR     - llama-server所在目录 (默认: ~/llama.cpp/build/bin)"
        echo "  INFERENCE_MODEL_DIR  - 推理模型目录(vLLM) (默认: /home/a/models/safetensors/Qwen3-4B)"
        echo "  VLLM_CONDA_ENV       - conda环境名 (默认: vllm；如果你的是 VLLM，请设置为 VLLM)"
        echo "  VLLM_SPECULATIVE_CONFIG - vLLM speculative JSON（可选；例如: {\"model\":\"/home/a/qwen3-4b-eagle3\",\"num_speculative_tokens\":5,\"method\":\"eagle3\"}）"
        echo "  VLLM_GPU_MEMORY_UTILIZATION - vLLM 显存占用比例 (默认: 0.85)"
        echo "  VLLM_SERVED_MODEL_NAME - vLLM 对外模型名 (默认: local-model；后端也会用它作为默认模型名)"
        echo "  VLLM_ENABLE_TOOL_CALLING - vLLM 工具调用开关(1启用；需版本支持)"
        echo "  VLLM_TOOL_CALL_PARSER - vLLM tool parser (如 hermes/qwen 等，按版本支持)"
        echo "  EMBEDDING_MODEL      - Embedding模型路径 (默认: ~/models/gguf/Qwen3/Qwen3-Embedding-4B/Qwen3-Embedding-4B-Q5_K_M.gguf)"
        echo "  VENV_PATH            - venv虚拟环境路径 (默认: \${PROJECT_ROOT}/backend_service/venv)"
        echo "  LOG_DIR              - 日志目录 (默认: \${PROJECT_ROOT}/logs)"
        echo "  INFERENCE_PORT       - 推理端口 (默认: 8081)"
        echo "  EMBEDDING_PORT       - embedding端口 (默认: 8090)"
        echo "  API_PORT             - 后端端口 (默认: 8000)"
        exit 1
        ;;
esac