@echo off
chcp 65001 >nul
setlocal enabledelayedexpansion

:: ============================================================
:: JRXML RAG 项目 - Windows 一键部署脚本
:: ============================================================

title JRXML RAG 环境部署

echo.
echo ============================================================
echo   JRXML RAG 项目 - 环境部署脚本
echo ============================================================
echo.

:: 检查 Python
echo [1/5] 检查 Python 环境...
python --version >nul 2>&1
if %errorlevel% neq 0 (
    echo [错误] 未找到 Python，请先安装 Python 3.11+
    pause
    exit /b 1
)
for /f "tokens=2" %%v in ('python --version 2^>^&1') do echo    Python 版本: %%v

:: 检查 CUDA
echo.
echo [2/5] 检查 CUDA 环境...
python -c "import torch; print(f'   PyTorch: {torch.__version__}'); print(f'   CUDA 可用: {torch.cuda.is_available()}'); print(f'   GPU: {torch.cuda.get_device_name(0) if torch.cuda.is_available() else \"N/A\"}')" 2>nul
if %errorlevel% neq 0 (
    echo    PyTorch 未安装，将在后续步骤安装
)

:: 初始化 .env 配置
echo.
echo [3/5] 初始化环境配置...
if not exist ".env" (
    if exist ".env.example" (
        copy ".env.example" ".env" >nul
        echo    已从 .env.example 创建 .env 配置文件
        echo    请根据需要编辑 .env 文件修改配置
    ) else (
        echo    [警告] 未找到 .env.example，跳过配置初始化
    )
) else (
    echo    .env 配置文件已存在，跳过
)

:: 创建必要目录
echo.
echo [4/5] 创建项目目录...
set "DIRS=jrxml_source jrxml_chunker_output models embeddings chroma_db"
for %%d in (%DIRS%) do (
    if not exist "%%d" (
        mkdir "%%d" >nul 2>&1
        echo    创建目录: %%d
    )
)
echo    目录结构已就绪

:: 安装依赖
echo.
echo [5/5] 安装 Python 依赖...
echo.

:: 检测是否有 NVIDIA GPU
python -c "import subprocess; exit(0 if subprocess.run(['nvidia-smi'], capture_output=True).returncode == 0 else 1)" 2>nul
if %errorlevel% equ 0 (
    set "HAS_GPU=1"
    echo    [检测到 NVIDIA GPU，安装 CUDA 版 PyTorch]
    echo.
    echo    安装 PyTorch (CUDA 版本)...
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cu124
) else (
    set "HAS_GPU=0"
    echo    [未检测到 GPU，安装 CPU 版 PyTorch]
    echo.
    echo    安装 PyTorch (CPU 版本)...
    pip install torch torchvision torchaudio --index-url https://download.pytorch.org/whl/cpu
)

echo.
echo    安装核心依赖...
pip install sentence-transformers chromadb numpy tqdm huggingface_hub

echo.
echo ============================================================
echo   部署完成！
echo ============================================================
echo.
echo   下一步操作:
echo     1. 编辑 .env 配置文件（可选）
echo     2. 收集 JRXML 文件:   python collect_jrxml.py
echo     3. 语义分块:          python jrxml_banch_chunker.py
echo     4. 下载嵌入模型:      python down_embedding_model.py
echo     5. 向量化:            python embed_chunks.py
echo     6. 导入 Chroma:       python import_to_chroma.py
echo     7. 开始查询:          python query_chroma.py
echo.
echo   查看当前配置:           python config.py
echo.
pause