#!/bin/bash
# PlantAI-dev 9000 服务管理脚本
# 用法: bash manage_plantai.sh {start|stop|restart|status} [--port PORT]
# 默认端口: 9000

set -euo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
if [ -f "$SCRIPT_DIR/.env" ]; then
    set -a
    . "$SCRIPT_DIR/.env"
    set +a
fi

# 解析命令行参数：第一个参数是命令，后面可选 --port / -p
CMD=""
WEB_PORT="${PLANTAI_WEB_PORT:-9000}"
for arg in "$@"; do
    case "$arg" in
        start|stop|restart|status)
            CMD="$arg"
            ;;
        --port|-p)
            WEB_PORT=""
            ;;
        *)
            if [ -z "$WEB_PORT" ] && [[ "$arg" =~ ^[0-9]+$ ]]; then
                WEB_PORT="$arg"
            elif [ -z "$CMD" ]; then
                echo "错误: 未知命令 '$arg'"
                echo "用法: bash manage_plantai.sh {start|stop|restart|status} [--port PORT]"
                exit 1
            fi
            ;;
    esac
done
CMD="${CMD:-status}"

# Conda 环境：启动脚本内部也依赖这个环境
CONDA_BASE="${BBAV5_CONDA_BASE:-/home/xia-lab/miniconda3}"
DEFAULT_ENV="${BBAV5_DEFAULT_ENV:-test4py311}"
CONDA_ENV="${CONDA_ENV:-$CONDA_BASE/envs/$DEFAULT_ENV}"
export PATH="$CONDA_ENV/bin:$PATH"

# 项目配置（与 .env 和启动脚本保持一致）
PROJECT_DIR="$SCRIPT_DIR"
QUEUE_PREFIX="${BBAV5_QUEUE_PREFIX:-plantai_4090_9000}"
REDIS_URL="${BBAV5_REDIS_URL:-redis://localhost:6379/10}"
WEB_LOG="/tmp/plantai_${WEB_PORT}.log"
CELERY_LOG="/tmp/celery_${WEB_PORT}.log"

red()  { echo -e "\033[31m$*\033[0m"; }
green(){ echo -e "\033[32m$*\033[0m"; }
cyan() { echo -e "\033[36m$*\033[0m"; }

web_pid_for_port() {
    ss -tlnp 2>/dev/null \
        | grep ":$WEB_PORT " \
        | grep -oP 'pid=\K\d+' \
        | head -n 1 \
        || true
}

web_process_cwd() {
    readlink -f "/proc/$1/cwd" 2>/dev/null || true
}

web_pid_belongs_to_project() {
    [ "$(web_process_cwd "$1")" = "$PROJECT_DIR" ]
}

status() {
    echo ""
    cyan "=== PlantAI-dev ${WEB_PORT} 服务状态 ==="
    echo ""

    # Web 服务
    WEB_PID="$(web_pid_for_port)"
    if [ -n "${WEB_PID:-}" ]; then
        WEB_CWD="$(web_process_cwd "$WEB_PID")"
        if web_pid_belongs_to_project "$WEB_PID"; then
            green "[✓] Web  服务运行中  (端口 $WEB_PORT, PID $WEB_PID)"
        else
            red "[!] 端口 $WEB_PORT 被其他项目占用  (PID $WEB_PID, cwd ${WEB_CWD:-未知})"
        fi
    else
        red "[✗] Web  服务未启动  (端口 $WEB_PORT)"
    fi

    # Redis
    if redis-cli -u "$REDIS_URL" ping > /dev/null 2>&1; then
        green "[✓] Redis 运行中  ($REDIS_URL)"
    else
        red "[✗] Redis 未启动"
    fi

    # Celery workers
    CELERY_COUNT="$(pgrep -fc "$QUEUE_PREFIX" || true)"
    if [ "$CELERY_COUNT" -gt 0 ]; then
        green "[✓] Celery workers 运行中  (${CELERY_COUNT} 进程, 队列前缀: $QUEUE_PREFIX)"
        echo "    Worker 列表:"
        ps aux | grep "$QUEUE_PREFIX" | grep -v grep | grep -oP "queues=\K\S+" | sort | while read q; do
            echo "      - $q"
        done
    else
        red "[✗] Celery workers 未启动  (队列前缀: $QUEUE_PREFIX)"
    fi
    echo ""
}

stop() {
    echo ""
    cyan "=== 停止 PlantAI-dev ${WEB_PORT} 服务 ==="
    echo ""

    # 停止 Web 服务
    WEB_PID="$(web_pid_for_port)"
    if [ -n "${WEB_PID:-}" ]; then
        WEB_CWD="$(web_process_cwd "$WEB_PID")"
        if ! web_pid_belongs_to_project "$WEB_PID"; then
            red "[!] 拒绝停止：端口 $WEB_PORT 属于其他项目  (PID $WEB_PID, cwd ${WEB_CWD:-未知})"
            return 1
        fi
        echo "正在停止 Web 服务 (PID $WEB_PID)..."
        kill "$WEB_PID" 2>/dev/null || true
        sleep 1
        # 仅强制停止仍属于本项目的原进程，避免 PID 复用造成误杀。
        if kill -0 "$WEB_PID" 2>/dev/null && web_pid_belongs_to_project "$WEB_PID"; then
            kill -9 "$WEB_PID" 2>/dev/null || true
        fi
        green "[✓] Web 服务已停止"
    else
        echo "[ ] Web 服务未运行"
    fi

    # 停止 Celery workers
    CELERY_PIDS=$(pgrep -f "$QUEUE_PREFIX" 2>/dev/null) || true
    if [ -n "${CELERY_PIDS:-}" ]; then
        echo "正在停止 Celery workers..."
        echo "$CELERY_PIDS" | xargs -r kill 2>/dev/null || true
        sleep 2
        # 清理残留
        REMAIN=$(pgrep -f "$QUEUE_PREFIX" 2>/dev/null) || true
        if [ -n "${REMAIN:-}" ]; then
            echo "  强制清理残留进程..."
            echo "$REMAIN" | xargs -r kill -9 2>/dev/null || true
        fi
        green "[✓] Celery workers 已停止"
    else
        echo "[ ] Celery workers 未运行"
    fi
    echo ""
}

start() {
    echo ""
    cyan "=== 启动 PlantAI-dev ${WEB_PORT} 服务 ==="
    echo ""

    # 检查是否已在运行
    WEB_PID="$(web_pid_for_port)"
    if [ -n "${WEB_PID:-}" ]; then
        WEB_CWD="$(web_process_cwd "$WEB_PID")"
        if web_pid_belongs_to_project "$WEB_PID"; then
            red "[!] 本项目 Web 已在端口 $WEB_PORT 运行  (PID $WEB_PID)"
        else
            red "[!] 端口 $WEB_PORT 被其他项目占用  (PID $WEB_PID, cwd ${WEB_CWD:-未知})"
        fi
        exit 1
    fi

    if ! redis-cli -u "$REDIS_URL" ping > /dev/null 2>&1; then
        red "[!] Redis 未运行，请先启动 Redis"
        exit 1
    fi

    cd "$PROJECT_DIR"

    # 启动 Celery workers
    echo "正在启动 Celery workers..."
    nohup bash start_celery_9000.sh > "$CELERY_LOG" 2>&1 &
    sleep 3
    CELERY_COUNT="$(pgrep -fc "$QUEUE_PREFIX" || true)"
    if [ "$CELERY_COUNT" -gt 0 ]; then
        green "  [✓] Celery workers 已启动 (${CELERY_COUNT} 进程)"
    else
        red "  [✗] Celery workers 启动失败，查看 $CELERY_LOG"
    fi

    # 启动 Web 服务
    echo "正在启动 Web 服务..."
    nohup bash start_plantai_9000.sh --port "$WEB_PORT" > "$WEB_LOG" 2>&1 &
    sleep 3
    if ss -tlnp 2>/dev/null | grep -q ":$WEB_PORT"; then
        green "  [✓] Web 服务已启动 (http://0.0.0.0:$WEB_PORT)"
    else
        red "  [✗] Web 服务启动失败，查看 $WEB_LOG"
    fi

    # 最终状态
    echo ""
    echo "日志文件:"
    echo "  Web:    $WEB_LOG"
    echo "  Celery: $CELERY_LOG"
    echo "  Worker: $PROJECT_DIR/logs/celery_${QUEUE_PREFIX}/"
    echo ""
}

restart() {
    stop
    sleep 1
    start
}

case "$CMD" in
    start)   start   ;;
    stop)    stop    ;;
    restart) restart ;;
    status)  status  ;;
    *)
        echo "用法: bash manage_plantai.sh {start|stop|restart|status} [--port PORT]"
        exit 1
        ;;
esac