#!/bin/bash

# AI Agent 启动脚本
# 用法: ./start.sh [backend|frontend|both]

set -e

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

echo -e "${BLUE}========================================${NC}"
echo -e "${BLUE}   AI Agent - 个人生活助手启动脚本${NC}"
echo -e "${BLUE}========================================${NC}"
echo ""

# 检查 vLLM 容器是否运行
check_vllm() {
    if ! docker ps --format '{{.Names}}' | grep -q "^gemma4-server$"; then
        echo -e "${YELLOW}⚠️  vLLM 容器未运行！${NC}"
        echo "正在启动 vLLM 容器（Gemma-4 模型）..."
        
        # 检查模型文件是否存在
        if [ ! -d "/home/huang/Study/AIModel/gemma-4-E2B-it" ]; then
            echo -e "${RED}✗ 错误：模型目录不存在: /home/huang/Study/AIModel/gemma-4-E2B-it${NC}"
            echo "请先下载模型或修改模型路径"
            exit 1
        fi
        
        docker run -d \
            --name gemma4-server \
            --group-add=video \
            --cap-add=SYS_PTRACE \
            --security-opt seccomp=unconfined \
            --device=/dev/kfd \
            --device=/dev/dri \
            -v /home/huang/Study/AIModel/gemma-4-E2B-it:/models/gemma-4-E2B-it \
            -e VLLM_ROCM_USE_AITER=0 \
            -e HF_TOKEN="${HF_TOKEN}" \
            -p 8000:8000 \
            --ipc=host \
            --entrypoint vllm \
            my-vllm-gemma4:working \
            serve /models/gemma-4-E2B-it \
            --served-model-name gemma-4-E2B-it \
            --dtype auto \
            --api-key token-abc123 \
            --trust-remote-code \
            --port 8000 \
            --gpu-memory-utilization 0.85 \
            --max-model-len 8192
        
        echo -e "${GREEN}✓ vLLM 容器已启动${NC}"
        echo -e "${YELLOW}⏳ 等待模型加载（可能需要几分钟）...${NC}"
        sleep 10
    else
        echo -e "${GREEN}✓ vLLM 容器正在运行${NC}"
    fi
}

# 检查 PostgreSQL 容器是否运行
check_postgres() {
    if ! docker ps | grep -q postgres-langgraph; then
        echo -e "${YELLOW}⚠️  PostgreSQL 容器未运行！${NC}"
        echo "正在启动 PostgreSQL 容器..."
        docker run -d \
            --name postgres-langgraph \
            -e POSTGRES_PASSWORD=mysecretpassword \
            -e POSTGRES_DB=langgraph_db \
            -p 5432:5432 \
            -v ~/docker_volumes/postgres_data:/var/lib/postgresql/data \
            postgres:16
        
        echo -e "${GREEN}✓ PostgreSQL 容器已启动${NC}"
        sleep 3
    else
        echo -e "${GREEN}✓ PostgreSQL 容器正在运行${NC}"
    fi
}

# 启动后端
start_backend() {
    echo -e "\n${BLUE}🚀 启动后端服务 (端口 8001)...${NC}"
    python backend.py &
    BACKEND_PID=$!
    echo -e "${GREEN}✓ 后端服务已启动 (PID: $BACKEND_PID)${NC}"
    sleep 2
}

# 启动前端
start_frontend() {
    echo -e "\n${BLUE}🎨 启动前端界面...${NC}"
    streamlit run frontend.py &
    FRONTEND_PID=$!
    echo -e "${GREEN}✓ 前端服务已启动 (PID: $FRONTEND_PID)${NC}"
    echo -e "${GREEN}✓ 请在浏览器中打开: http://localhost:8501${NC}"
}

# 清理函数
cleanup() {
    echo -e "\n${RED}🛑 正在停止所有服务...${NC}"
    if [ ! -z "$BACKEND_PID" ]; then
        kill $BACKEND_PID 2>/dev/null || true
        echo -e "${GREEN}✓ 后端服务已停止${NC}"
    fi
    if [ ! -z "$FRONTEND_PID" ]; then
        kill $FRONTEND_PID 2>/dev/null || true
        echo -e "${GREEN}✓ 前端服务已停止${NC}"
    fi
    echo -e "${YELLOW}💡 提示：Docker 容器需要手动停止${NC}"
    echo -e "   停止 vLLM: docker stop gemma4-server"
    echo -e "   停止 PostgreSQL: docker stop postgres-langgraph"
    exit 0
}

# 捕获 Ctrl+C
trap cleanup SIGINT SIGTERM

# 主逻辑
case "${1:-both}" in
    backend)
        check_vllm
        check_postgres
        start_backend
        echo -e "\n${GREEN}后端服务正在运行，按 Ctrl+C 停止${NC}"
        wait $BACKEND_PID
        ;;
    frontend)
        start_frontend
        echo -e "\n${GREEN}前端服务正在运行，按 Ctrl+C 停止${NC}"
        wait $FRONTEND_PID
        ;;
    both|*)
        check_vllm
        check_postgres
        start_backend
        start_frontend
        echo -e "\n${GREEN}所有服务正在运行，按 Ctrl+C 停止 Python 服务${NC}"
        echo -e "${YELLOW}注意：Docker 容器会在后台继续运行${NC}"
        wait
        ;;
esac