@echo off
setlocal enabledelayedexpansion

REM 清理旧的资源目录
IF EXIST "resource" (
    echo 清理旧的资源目录...
    rmdir /s /q "resource"
)
else (
    echo 资源目录不存在，创建资源目录...
    mkdir resource
)



REM 设置当前目录
SET PROJECT_DIR=%~dp0
cd %PROJECT_DIR%

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

REM 检查resource目录是否存在
if not exist %PROJECT_DIR%/resource (
    echo 错误：resource目录不存在
    echo 请检查目录是否存在
    exit /b 1
)

REM 批量创建resource目录结构
echo 正在创建resource目录结构...
mkdir %PROJECT_DIR%/resource/NeuralHaircut/
mkdir %PROJECT_DIR%/resource/NeuralHaircut/PIXIE/
mkdir %PROJECT_DIR%/resource/hyperIQA/pretrained/
mkdir %PROJECT_DIR%/resource/Matte-Anything/pretrained/
mkdir %PROJECT_DIR%/resource/openpose

REM Download Neural Haircut files
echo 正在下载Neural Haircut文件...
cd %PROJECT_DIR%/resource/NeuralHaircut
python -m gdown --folder https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8

REM 更新 diffusion_prior 模型
cd %PROJECT_DIR%/resource/NeuralHaircut/pretrained_models/diffusion_prior
echo 正在下载Neural Haircut扩散先验模型文件...
python -m gdown 1_9EOUXHayKiGH5nkrayncln3d6m1uV7f

REM 下载 PIXIE 文件
cd %PROJECT_DIR%/resource/NeuralHaircut/PIXIE
echo 正在下载Neural HaircutPIXIE模型文件...
python -m gdown 1mPcGu62YPc4MdkT8FFiOCP629xsENHZf && tar -xvzf pixie_data.tar.gz ./ && rm pixie_data.tar.gz

REM 下载 hyperIQA 模型
cd %PROJECT_DIR%/resource/hyperIQA/pretrained
echo 正在下载hyperIQA模型文件...
python -m gdown 1OOUmnbvpGea0LIGpIWEbOyxfWx6UCiiE
cd %PROJECT_DIR%

REM Download Matte-Anything files
cd %PROJECT_DIR%/resource/Matte-Anything/pretrained
echo 正在下载Matte-Anything扩散先验模型文件...
wget https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
echo 正在下载Matte-AnythingGroundingDINO模型文件...
wget https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
echo 正在下载Matte-Anything模型文件...
python -m gdown 1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW

REM Download OpenPose files
cd %PROJECT_DIR%/resource/openpose
echo 正在下载OpenPose模型文件...
python -m gdown 1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV && tar -xvzf models.tar.gz && rm models.tar.gz

cd %PROJECT_DIR%
echo 所有模型文件下载完成！