#!/bin/bash
# ============================================
# Dragon Q6A 环境自动配置脚本
# 适用：Python 3.12 + PyQt5 + Open3D + PyOpenGL + GPU 硬件加速
# ============================================

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

echo "============================================"
echo "  Dragon Q6A 环境配置开始"
echo "============================================"

# ============================================
# 1. 系统更新
# ============================================
echo ""
echo "[1/7] 更新系统..."
sudo rsetup

# ============================================
# 2. 安装系统依赖
# ============================================
echo ""
echo "[2/7] 安装系统依赖..."
sudo apt install -y \
    build-essential git wget curl \
    python3-pyqt5 pyqt5-dev-tools \
    python3-pyqt5.qtsvg python3-pyqt5.qtwebengine \
    mesa-utils mesa-vulkan-drivers \
    libgl1-mesa-dri libglx-mesa0 \
    libegl1-mesa-dev libgles2-mesa-dev \
    libgbm-dev libdrm-dev libvulkan-dev \
    vulkan-tools libglu1-mesa-dev libx11-xcb-dev \
    python3-venv python3-pip \
    v4l-utils ffmpeg

# ============================================
# 3. 验证 GPU
# ============================================
echo ""
echo "[3/7] 验证 GPU 驱动..."
echo "OpenGL 渲染器信息:"
glxinfo | grep "OpenGL renderer" | head -1
echo ""
echo "Vulkan 驱动信息:"
vulkaninfo | grep -E "driverName|deviceName" | head -4 || echo "⚠️  请检查 Vulkan 驱动"

# ============================================
# 4. 设置虚拟环境
# ============================================
echo ""
echo "[4/7] 设置 Python 虚拟环境..."

# 定义项目目录
PROJECT_DIR1="$HOME/orisys_gripper/ctag2f-90d-sdk"
PROJECT_DIR2="$HOME/orisys_gripper/orisys-sdk"

configure_project() {
    local PROJECT_DIR="$1"
    if [ -d "$PROJECT_DIR" ]; then
        echo "  配置项目: $PROJECT_DIR"
        cd "$PROJECT_DIR"
        
        # 删除旧的虚拟环境（如果存在）
        if [ -d ".venv" ]; then
            echo "    移除旧的虚拟环境..."
            rm -rf .venv
        fi
        
        # 创建新虚拟环境
        echo "    创建虚拟环境..."
        python3 -m venv .venv --system-site-packages
        
        # 激活并安装依赖（使用 . 代替 source）
        echo "    安装 Python 依赖..."
        . .venv/bin/activate
        
        # 升级 pip
        pip install --upgrade pip
        
        # 安装 requirements.txt（如果存在）
        if [ -f "requirements.txt" ]; then
            pip install -r requirements.txt
        fi
        #pip install pyqtgraph
        
        # 安装基础包
        #pip install PyOpenGL
        #pip install open3d-unofficial-arm
        #pip install numpy scipy matplotlib pyyaml
        
        # 验证安装
        echo "    验证安装..."
        python -c "import PyQt5.QtCore; print('      PyQt5 OK')"
        python -c "from OpenGL.GL import *; print('      OpenGL OK')"
        python -c "import open3d as o3d; print('      Open3D OK: v' + o3d.__version__)"
        
        # 退出虚拟环境
        deactivate
        
        cd - > /dev/null
        echo "    ✅ $PROJECT_DIR 配置完成"
        return 0
    else
        echo "  ⚠️  目录不存在: $PROJECT_DIR (跳过)"
        return 1
    fi
}

# 配置项目
configure_project "$PROJECT_DIR1"
configure_project "$PROJECT_DIR2"

# ============================================
# 6. 创建启动脚本
# ============================================
echo ""
echo "[6/7] 创建启动脚本..."

cat > ~/run_visualizer.sh << 'EOF'
#!/bin/bash
# Dragon Q6A 可视化程序启动脚本

# 查找项目目录
if [ -d "$HOME/orisys_gripper/ctag2f-90d-sdk" ]; then
    PROJECT_DIR="$HOME/orisys_gripper/ctag2f-90d-sdk"
elif [ -d "$HOME/orisys_gripper/orisys-sdk" ]; then
    PROJECT_DIR="$HOME/orisys_gripper/orisys-sdk"
else
    echo "错误: 找不到项目目录"
    exit 1
fi

echo "进入项目目录: $PROJECT_DIR"
cd "$PROJECT_DIR"

# 激活虚拟环境
if [ -f ".venv/bin/activate" ]; then
    . .venv/bin/activate
else
    echo "错误: 找不到虚拟环境"
    exit 1
fi

# 运行程序
python -m gripper.visualizer "$@"
EOF

chmod +x ~/run_visualizer.sh
echo "✅ 启动脚本已创建: ~/run_visualizer.sh"

# ============================================
# 7. 最终验证
# ============================================
echo ""
echo "[7/7] 最终验证..."

# 查找存在的项目目录
if [ -d "$PROJECT_DIR1" ]; then
    cd "$PROJECT_DIR1"
elif [ -d "$PROJECT_DIR2" ]; then
    cd "$PROJECT_DIR2"
else
    echo "⚠️  未找到项目目录，跳过验证"
    cd "$HOME"
fi

if [ -f ".venv/bin/activate" ]; then
    . .venv/bin/activate
    echo "✅ 虚拟环境激活成功"
    
    echo "验证 Python 包:"
    python -c "import PyQt5.QtCore; print('  ✅ PyQt5: OK')" 2>/dev/null || echo "  ❌ PyQt5: 失败"
    python -c "from OpenGL.GL import *; print('  ✅ PyOpenGL: OK')" 2>/dev/null || echo "  ❌ PyOpenGL: 失败"
    python -c "import open3d as o3d; print('  ✅ Open3D: v' + o3d.__version__)" 2>/dev/null || echo "  ❌ Open3D: 失败"
    
    deactivate
fi

# ============================================
echo ""
echo "============================================"
echo "  Dragon Q6A 环境配置完成！"
echo "============================================"
echo ""
echo "GPU 状态: 硬件加速已启用 (FD643/Turnip)"
echo ""
echo "下一步："
echo "  1. 重启终端或执行: source ~/.bashrc"
echo "  2. 运行程序: ~/run_visualizer.sh"
echo ""
echo "如果遇到图形问题，可以设置软件渲染模式："
echo "  export LIBGL_ALWAYS_SOFTWARE=1"
echo "============================================"