#!/bin/bash
# =============================================================================
# AI Agent 启动与管理脚本
# 用法: ./start.sh [check|backend|frontend|both|docker-up|docker-down]
# =============================================================================

set -e

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

# 项目根目录
PROJECT_DIR="/home/huang/Study/AIProject/Agent1"

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

# =============================================================================
# 配置检查函数
# =============================================================================
check_config() {
    echo -e "${BLUE}📋 开始环境配置检查...${NC}"
    echo ""
    
    PASS=0
    FAIL=0
    WARN=0
    
    # 辅助函数
    check_pass() {
        echo -e "${GREEN}✓${NC} $1"
        ((PASS++))
    }
    
    check_fail() {
        echo -e "${RED}✗${NC} $1"
        ((FAIL++))
    }
    
    check_warn() {
        echo -e "${YELLOW}⚠${NC} $1"
        ((WARN++))
    }
    
    # 1. 检查 .env 文件
    echo "🔍 检查配置文件..."
    if [ -f "$PROJECT_DIR/.env" ]; then
        check_pass ".env 文件存在"
        
        # 检查文件格式
        if grep -q "^EOF" .env 2>/dev/null; then
            check_fail ".env 文件格式错误：发现多余的 EOF 标记"
        else
            check_pass ".env 文件格式正确"
        fi
    else
        check_fail ".env 文件不存在"
        echo "   提示: 请创建 .env 文件并配置环境变量"
        return 1
    fi
    
    # 2. 检查必需的环境变量
    echo ""
    echo "🔑 检查环境变量..."
    
    # 检查 ZHIPUAI_API_KEY
    if grep -q "^ZHIPUAI_API_KEY=" "$PROJECT_DIR/.env" 2>/dev/null; then
        API_KEY=$(grep "^ZHIPUAI_API_KEY=" "$PROJECT_DIR/.env" | head -1 | cut -d'=' -f2- | tr -d '[:space:]')
        if [ ${#API_KEY} -gt 10 ]; then
            check_pass "ZHIPUAI_API_KEY 已配置（长度: ${#API_KEY}）"
        else
            check_fail "ZHIPUAI_API_KEY 配置可能无效（过短）"
        fi
    else
        check_fail "ZHIPUAI_API_KEY 未配置或格式错误"
    fi
    
    # 检查 VLLM_LOCAL_KEY
    if grep -q "^VLLM_LOCAL_KEY=" "$PROJECT_DIR/.env" 2>/dev/null; then
        check_pass "VLLM_LOCAL_KEY 已配置"
    else
        check_warn "VLLM_LOCAL_KEY 未配置（如不使用本地模型可忽略）"
    fi
    
    # 检查 DB_URI (远程服务器)
    if grep -q "^DB_URI=.*115.190.121.151" "$PROJECT_DIR/.env" 2>/dev/null; then
        check_pass "DB_URI 已配置（远程服务器）"
    elif grep -q "^DB_URI=" "$PROJECT_DIR/.env" 2>/dev/null; then
        check_warn "DB_URI 已配置（非远程服务器地址）"
    else
        check_fail "DB_URI 未配置"
    fi
    
    # 检查 QDRANT_URL (远程服务器)
    if grep -q "^QDRANT_URL=.*115.190.121.151" "$PROJECT_DIR/.env" 2>/dev/null; then
        check_pass "QDRANT_URL 已配置（远程服务器）"
    elif grep -q "^QDRANT_URL=" "$PROJECT_DIR/.env" 2>/dev/null; then
        check_warn "QDRANT_URL 已配置（非远程服务器地址）"
    else
        check_fail "QDRANT_URL 未配置"
    fi
    
    # 3. 检查 Docker 环境
    echo ""
    echo "🐳 检查 Docker 环境..."
    
    if command -v docker &> /dev/null; then
        check_pass "Docker 已安装"
        
        if docker info &> /dev/null; then
            check_pass "Docker 守护进程正在运行"
        else
            check_fail "Docker 守护进程未运行"
            echo "   提示: sudo systemctl start docker"
        fi
    else
        check_fail "Docker 未安装"
    fi
    
    if command -v docker compose version &> /dev/null || command -v docker-compose &> /dev/null; then
        check_pass "Docker Compose 已安装"
    else
        check_fail "Docker Compose 未安装"
    fi
    
    # 4. 检查端口占用
    echo ""
    echo "🔌 检查端口占用..."
    
    for port in 8001 8501; do
        if lsof -i :$port &> /dev/null; then
            check_warn "端口 $port 已被占用"
        else
            check_pass "端口 $port 可用"
        fi
    done
    
    # 5. 检查远程服务连接
    echo ""
    echo "🌐 检查远程服务连接..."
    
    # 测试 PostgreSQL 连接
    if command -v psql &> /dev/null; then
        # 注意：这里假设密码为 mysecretpassword，如果不同需调整或从 env 读取
        if PGPASSWORD=mysecretpassword psql -h 115.190.121.151 -U postgres -d langgraph_db -c "SELECT 1;" &> /dev/null; then
            check_pass "PostgreSQL 远程连接正常 (115.190.121.151:5432)"
        else
            check_fail "PostgreSQL 远程连接失败"
            echo "   提示: 检查网络连接和防火墙设置"
        fi
    else
        check_warn "psql 客户端未安装，跳过 PostgreSQL 连接测试"
    fi
    
    # 测试 Qdrant 连接
    if curl -s http://115.190.121.151:6333/collections &> /dev/null; then
        check_pass "Qdrant 远程连接正常 (115.190.121.151:6333)"
    else
        check_fail "Qdrant 远程连接失败"
        echo "   提示: 检查网络连接和防火墙设置"
    fi
    
    # 总结
    echo ""
    echo "=========================================="
    echo "  检查结果汇总"
    echo "=========================================="
    echo -e "${GREEN}通过: $PASS${NC}"
    echo -e "${RED}失败: $FAIL${NC}"
    echo -e "${YELLOW}警告: $WARN${NC}"
    echo ""
    
    if [ $FAIL -eq 0 ]; then
        echo -e "${GREEN}✅ 配置检查通过！${NC}"
        return 0
    else
        echo -e "${RED}❌ 发现 $FAIL 个错误，请修复后重试${NC}"
        return 1
    fi
}

# =============================================================================
# Docker 容器检查函数（仅检查 llama.cpp 服务）
# =============================================================================
check_llamacpp() {
    echo -e "${BLUE}🔍 检查 llama.cpp LLM 容器...${NC}"
    if ! docker ps --format '{{.Names}}' | grep -q "^gemma4-llamacpp-server$"; then
        echo -e "${YELLOW}⚠️  llama.cpp LLM 容器未运行${NC}"
        return 1
    else
        echo -e "${GREEN}✓ llama.cpp LLM 容器正在运行 (gemma4-llamacpp-server)${NC}"
        return 0
    fi
}

check_embedding() {
    echo -e "${BLUE}🔍 检查 llama.cpp Embedding 容器...${NC}"
    if ! docker ps --format '{{.Names}}' | grep -q "^embedding-server$"; then
        echo -e "${YELLOW}⚠️  llama.cpp Embedding 容器未运行${NC}"
        return 1
    else
        echo -e "${GREEN}✓ llama.cpp Embedding 容器正在运行 (embedding-server)${NC}"
        return 0
    fi
}

# =============================================================================
# 启动 Docker 依赖服务（llama.cpp）
# =============================================================================
start_llamacpp() {
    echo -e "${BLUE}🚀 启动 llama.cpp LLM 容器...${NC}"
    
    # 检查模型文件
    if [ ! -f "/home/huang/Study/AIModel/GGUF/Gemma-4-E2B-Uncensored-HauhauCS-Aggressive-Q6_K_P.gguf" ]; then
        echo -e "${RED}✗ 错误：LLM 模型文件不存在${NC}"
        exit 1
    fi
    
    if [ ! -f "/home/huang/Study/AIModel/GGUF/mmproj-Gemma-4-E2B-Uncensored-HauhauCS-Aggressive-f16.gguf" ]; then
        echo -e "${RED}✗ 错误：多模态投影文件不存在${NC}"
        exit 1
    fi
    
    docker run -d \
        --name gemma4-llamacpp-server \
        --restart=unless-stopped \
        --group-add=video \
        --device=/dev/kfd \
        --device=/dev/dri \
        -v /home/huang/Study/AIModel/GGUF:/models \
        -p 8081:8080 \
        ghcr.io/ggml-org/llama.cpp:server-rocm \
        -m /models/Gemma-4-E2B-Uncensored-HauhauCS-Aggressive-Q6_K_P.gguf \
        --mmproj /models/mmproj-Gemma-4-E2B-Uncensored-HauhauCS-Aggressive-f16.gguf \
        --host 0.0.0.0 \
        --port 8080 \
        -ngl 99
    
    echo -e "${GREEN}✓ llama.cpp LLM 容器已启动 (端口 8081)${NC}"
    echo -e "${YELLOW}⏳ 等待模型加载（可能需要几分钟）...${NC}"
    sleep 15
}

start_embedding() {
    echo -e "${BLUE}🚀 启动 llama.cpp Embedding 容器...${NC}"
    
    # 检查模型文件
    if [ ! -f "/home/huang/Study/AIModel/GGUF/embeddinggemma-300M-Q8_0.gguf" ]; then
        echo -e "${RED}✗ 错误：Embedding 模型文件不存在${NC}"
        exit 1
    fi
    
    docker run -d \
        --name embedding-server \
        --restart=unless-stopped \
        --group-add=video \
        --device=/dev/kfd \
        --device=/dev/dri \
        -v /home/huang/Study/AIModel/GGUF:/models \
        -p 8082:8080 \
        ghcr.io/ggml-org/llama.cpp:server-rocm \
        -m /models/embeddinggemma-300M-Q8_0.gguf \
        --host 0.0.0.0 \
        --port 8080 \
        -ngl 99 \
        --embeddings \
        -c 512
    
    echo -e "${GREEN}✓ llama.cpp Embedding 容器已启动 (端口 8082)${NC}"
    sleep 5
}

# =============================================================================
# 启动 Python 服务
# =============================================================================
start_backend() {
    echo -e "\n${BLUE}🚀 启动后端服务 (端口 8003)...${NC}"
    cd "$PROJECT_DIR"
    
    # 加载 .env 文件中的环境变量
    set -a
    source .env 2>/dev/null || true
    set +a
    
    export PYTHONPATH="$PROJECT_DIR"
    export BACKEND_PORT=8003
    python app/backend.py &
    BACKEND_PID=$!
    echo -e "${GREEN}✓ 后端服务已启动 (PID: $BACKEND_PID)${NC}"
    sleep 2
}

start_frontend() {
    echo -e "\n${BLUE}🎨 启动前端界面 (端口 8501)...${NC}"
    cd "$PROJECT_DIR"
    
    # 加载 .env 文件中的环境变量
    set -a
    source .env 2>/dev/null || true
    set +a
    
    export PYTHONPATH="$PROJECT_DIR"
    streamlit run frontend/frontend.py &
    FRONTEND_PID=$!
    echo -e "${GREEN}✓ 前端服务已启动 (PID: $FRONTEND_PID)${NC}"
    echo -e "${GREEN}✓ 访问地址:${NC}"
    echo -e "   本地开发: http://localhost:8501"
}

# =============================================================================
# Docker Compose 管理
# =============================================================================
docker_up() {
    echo -e "${BLUE}🐳 使用 Docker Compose 启动所有服务...${NC}"
    cd "$PROJECT_DIR/docker"
    
    # 检查 .env 文件
    if [ ! -f "../.env" ]; then
        echo -e "${RED}✗ 错误：.env 文件不存在${NC}"
        echo "   请先复制配置文件:"
        echo "   cp .env.docker .env  # 服务器部署"
        exit 1
    fi
    
    docker compose up -d --build
    
    echo -e "\n${GREEN}✓ Docker Compose 服务已启动${NC}"
    echo -e "${BLUE}📊 查看服务状态:${NC} docker compose ps"
    echo -e "${BLUE}📝 查看日志:${NC} docker compose logs -f"
    echo -e "${BLUE}🌐 访问应用:${NC} http://localhost:8501"
}

docker_down() {
    echo -e "${BLUE}🛑 停止 Docker Compose 服务...${NC}"
    cd "$PROJECT_DIR/docker"
    docker compose down
    echo -e "${GREEN}✓ 服务已停止${NC}"
}

# =============================================================================
# 清理函数
# =============================================================================
cleanup() {
    echo -e "\n${RED}🛑 正在停止 Python 服务...${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 "   停止 llama.cpp LLM: docker stop gemma4-llamacpp-server"
    echo -e "   停止 llama.cpp Embedding: docker stop embedding-server"
    echo -e "   或使用: $0 docker-down"
    exit 0
}

# 捕获 Ctrl+C
trap cleanup SIGINT SIGTERM

# =============================================================================
# 主逻辑
# =============================================================================
case "${1:-help}" in
    check)
        check_config
        ;;
    
    backend)
        check_config || exit 1
        check_llamacpp || start_llamacpp
        check_embedding || start_embedding
        start_backend
        echo -e "\n${GREEN}后端服务正在运行，按 Ctrl+C 停止${NC}"
        wait $BACKEND_PID
        ;;
    
    frontend)
        check_config || exit 1
        start_frontend
        echo -e "\n${GREEN}前端服务正在运行，按 Ctrl+C 停止${NC}"
        wait $FRONTEND_PID
        ;;
    
    both)
        check_config || exit 1
        check_llamacpp || start_llamacpp
        check_embedding || start_embedding
        start_backend
        start_frontend
        echo -e "\n${GREEN}所有服务正在运行，按 Ctrl+C 停止 Python 服务${NC}"
        echo -e "${YELLOW}注意：Docker 容器会在后台继续运行${NC}"
        wait
        ;;
    
    docker-up)
        check_config || exit 1
        docker_up
        ;;
    
    docker-down)
        docker_down
        ;;
    
    help|*)
        echo -e "${BLUE}用法:${NC} $0 [command]"
        echo ""
        echo -e "${BLUE}命令:${NC}"
        echo "  check       检查环境配置"
        echo "  backend     仅启动后端服务"
        echo "  frontend    仅启动前端服务"
        echo "  both        启动前后端服务（默认）"
        echo "  docker-up   使用 Docker Compose 启动所有服务"
        echo "  docker-down 停止 Docker Compose 服务"
        echo "  help        显示此帮助信息"
        echo ""
        echo -e "${BLUE}示例:${NC}"
        echo "  $0 check              # 检查配置"
        echo "  $0 both               # 启动本地开发环境"
        echo "  $0 docker-up          # 启动 Docker 部署环境"
        ;;
esac