@echo off
setlocal enabledelayedexpansion

SET PROJECT_DIR=%~dp0
cd %PROJECT_DIR%

REM 定义文件路径
SET "NEURAL_HAIRCUT_FOLDER_PATH=resource\NeuralHaircut"
SET "PIXIE_PATH=resource\NeuralHaircut\PIXIE\pixie_data\pixie_data"
SET "MATTE_PATH=resource\Matte-Anything\model.pth"
SET "SAM_PATH=resource\Matte-Anything\sam_vit_h_4b8939.pth"
SET "DINO_PATH=resource\Matte-Anything\groundingdino_swint_ogc.pth"
SET "OPENPOSE_PATH=resource\openpose\models\pose\coco\pose_iter_584000.caffemodel"
SET "HYPERIQA_PATH=resource\hyperIQA\pretrained\hyperIQA.pth"

echo ======================================= 基础环境配置 ==============================================
echo.
echo **************************************************************************************************
echo *                                        基础环境配置                                             *
echo *                                   检查Python是否安装，更新pip，安装gdown                         *
echo *                                            请耐心等待...                                        *
echo **************************************************************************************************
echo.
REM 检查Python是否安装
python --version >nul 2>nul
IF %ERRORLEVEL% NEQ 0 (
    echo 错误：未找到Python
    echo 请安装Python并将其添加到系统环境变量PATH中
    exit /b 1
)

REM 更新pip
python -m pip install --upgrade pip

REM 检查gdown是否已安装
python -c "import gdown" >nul 2>nul
IF %ERRORLEVEL% EQU 0 (
    echo gdown已安装，跳过安装步骤...
) ELSE (
    REM 安装gdown
    echo 正在安装gdown...
    python -m pip install --user gdown --upgrade
    python -m pip install --user requests --upgrade
    IF %ERRORLEVEL% NEQ 0 (
        echo 错误：gdown安装失败
        echo 请检查网络连接后重试
        exit /b 1
    )
)

REM 添加Python Scripts目录到PATH
for /f "tokens=*" %%i in ('python -c "import sys; print(sys.executable)"') do set PYTHON_PATH=%%i
set PYTHON_SCRIPTS_PATH=%PYTHON_PATH:python.exe=Scripts%
set PYTHON_USER_SCRIPTS=%USERPROFILE%\AppData\Roaming\Python\Python3*\Scripts
set PATH=%PYTHON_SCRIPTS_PATH%;%PYTHON_USER_SCRIPTS%;%PATH%

REM 检查7-Zip是否安装
where 7z >nul 2>nul
IF %ERRORLEVEL% NEQ 0 (
    echo 错误：未找到7-Zip
    echo 请从 https://7-zip.org/ 下载并安装7-Zip
    echo 并将其添加到系统环境变量PATH中
    exit /b 1
)


REM 验证gdown是否可用
echo 正在验证gdown安装...
python -c "import gdown" >nul 2>nul
IF %ERRORLEVEL% NEQ 0 (
    echo 错误：gdown安装验证失败
    echo 请重新运行脚本
    exit /b 1
)
echo gdown安装成功！
echo.

echo ======================================= 开始下载所需资源文件 =======================================
echo.
echo **************************************************************************************************
echo *                                        开始下载所需资源文件                                      *
echo *                                    下载时间取决于您的网络连接速度                                 *
echo *                                            请耐心等待...                                        *
echo **************************************************************************************************
echo.

echo [1/8] 正在下载Neural Haircut文件...
mkdir %PROJECT_DIR%/resource/NeuralHaircut
cd %PROJECT_DIR%/resource/NeuralHaircut
IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckpt.pt" (
    IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\strand_ckpt.pt" (
        echo Neural Haircut文件已存在，跳过下载...
    ) ELSE (
        echo Neural Haircut Strand Prior模型不完整，重新下载...
        python -m gdown --folder "https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8"
    )
) ELSE (
    echo Neural Haircut Diffusion Prior模型不完整，重新下载...
    python -m gdown --folder "https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8"
)

REM 检查下载后的文件完整性
IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\dif_ckpt.pt" (
    echo 错误：Neural Haircut Diffusion Prior模型下载失败
    exit /b 1
)
IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\strand_ckpt.pt" (
    echo 错误：Neural Haircut Strand Prior模型下载失败
    exit /b 1
)

echo [2/8] 正在下载Neural HaircutPIXIE模型文件...
IF EXIST "%PIXIE_PATH%" (
    echo PIXIE模型已存在，跳过下载...
) ELSE (
    mkdir %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
    cd %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
    python -m gdown "1mPcGu62YPc4MdkT8FFiOCP629xsENHZf" && tar -xvzf pixie_data.tar.gz ./ && rm pixie_data.tar.gz
    cd %PROJECT_DIR%
)

echo [3/8] 正在下载hyperIQA模型文件...
mkdir %PROJECT_DIR%/resource/hyperIQA/pretrained
cd %PROJECT_DIR%/resource/hyperIQA/pretrained
IF EXIST "hyperIQA.pth" (
    echo hyperIQA模型已存在，跳过下载...
) ELSE (
    python -m gdown "1OOUmnbvpGea0LIGpIWEbOyxfWx6UCiiE"
    IF NOT EXIST "hyperIQA.pth" (
        echo 错误：hyperIQA模型下载失败
        exit /b 1
    )
)
cd %PROJECT_DIR%

echo [4/8] 正在下载Matte-Anything扩散先验模型文件...
mkdir %PROJECT_DIR%/resource/Matte-Anything/pretrained
cd %PROJECT_DIR%/resource/Matte-Anything/pretrained
IF EXIST "sam_vit_h_4b8939.pth" (
    echo SAM模型已存在，跳过下载...
) ELSE (
    wget "https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
    IF NOT EXIST "sam_vit_h_4b8939.pth" (
        echo 错误：SAM模型下载失败
        exit /b 1
    )
)

echo [5/8] 正在下载Matte-AnythingGroundingDINO模型文件...
IF EXIST "groundingdino_swint_ogc.pth" (
    echo GroundingDINO模型已存在，跳过下载...
) ELSE (
    wget "https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth"
    IF NOT EXIST "groundingdino_swint_ogc.pth" (
        echo 错误：GroundingDINO模型下载失败
        exit /b 1
    )
)

echo [6/8] 正在下载Matte-Anything模型文件...
IF EXIST "model.pth" (
    echo Matte-Anything模型已存在，跳过下载...
) ELSE (
    python -m gdown "1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW"
    IF NOT EXIST "model.pth" (
        echo 错误：Matte-Anything模型下载失败
        exit /b 1
    )
)
cd %PROJECT_DIR%

echo [7/8] 正在下载OpenPose模型文件...
mkdir %PROJECT_DIR%/resource/openpose/models/pose/coco
cd %PROJECT_DIR%/resource/openpose/models/pose/coco
IF EXIST "pose_iter_584000.caffemodel" (
    echo OpenPose模型已存在，跳过下载...
) ELSE (
    python -m gdown "1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV" && tar -xvzf models.tar.gz && rm models.tar.gz
    IF NOT EXIST "pose_iter_584000.caffemodel" (
        echo 错误：OpenPose模型下载失败
        exit /b 1
    )
)
cd %PROJECT_DIR%