@echo off
setlocal enabledelayedexpansion

REM 清理旧的资源目录
IF EXIST "resource" (
    echo 清理旧的资源目录...
    rmdir /s /q "resource"
)

REM 清理错误路径的资源
IF EXIST "NeuralHaircut" rmdir /s /q "NeuralHaircut"
IF EXIST "Matte-Anything" rmdir /s /q "Matte-Anything"
IF EXIST "openpose" rmdir /s /q "openpose"
IF EXIST "hyperIQA" rmdir /s /q "hyperIQA"

REM 定义模型信息
SET "NEURAL_HAIRCUT_FOLDER_URL=1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8"
SET "PIXIE_MODEL_URL=1mPcGu62YPc4MdkT8FFiOCP629xsENHZf"
SET "MATTE_MODEL_URL=1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW"
SET "SAM_MODEL_URL=https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth"
SET "DINO_MODEL_URL=https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth"
SET "OPENPOSE_MODEL_URL=1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV"
SET "HYPERIQA_MODEL_URL=1OOUmnbvpGea0LIGpIWEbOyxfWx6UCiiE"

REM 定义文件路径
SET "NEURAL_HAIRCUT_FOLDER_PATH=resource\NeuralHaircut"
SET "PIXIE_PATH=resource\NeuralHaircut\PIXIE\pixie_data.tar.gz"
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"

REM 下载设置
SET "RETRY_COUNT=3"
SET "TIMEOUT_SECONDS=600"
SET "CURL_OPTS=--connect-timeout 30 --max-time %TIMEOUT_SECONDS% --retry %RETRY_COUNT% -L -C -"
SET "MIN_FILE_SIZE=1000000"

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目录结构
mkdir resource 2>nul
cd resource
mkdir NeuralHaircut\pretrained_models\diffusion_prior 2>nul
mkdir NeuralHaircut\pretrained_models\strand_prior 2>nul
mkdir NeuralHaircut\PIXIE 2>nul
mkdir Matte-Anything 2>nul
mkdir openpose\models\pose\coco 2>nul
mkdir hyperIQA\pretrained 2>nul

REM 验证已有文件
echo 正在检查已有资源文件...
SET MISSING_FILES=0
SET /A TOTAL_FILES=8
SET /A EXISTING_FILES=0
SET MISSING_LIST=

REM Neural Haircut Diffusion Prior 检查
IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\model.pt" (
    echo [√] Neural Haircut Diffusion Prior已存在
    SET /A EXISTING_FILES+=1
) ELSE (
    echo [×] 缺少Neural Haircut Diffusion Prior
    SET /A MISSING_FILES+=1
    SET MISSING_LIST=!MISSING_LIST!Neural Haircut Diffusion Prior, 
)

REM Neural Haircut Strand Prior 检查
IF EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\model.pt" (
    echo [√] Neural Haircut Strand Prior已存在
    SET /A EXISTING_FILES+=1
) ELSE (
    echo [×] 缺少Neural Haircut Strand Prior
    SET /A MISSING_FILES+=1
    SET MISSING_LIST=!MISSING_LIST!Neural Haircut Strand Prior, 
)

REM PIXIE 模型检查
IF EXIST "%PIXIE_PATH%" (
    echo [√] PIXIE模型已存在
    SET /A EXISTING_FILES+=1
) ELSE (
    echo [×] 缺少PIXIE模型
    SET /A MISSING_FILES+=1
    SET MISSING_LIST=!MISSING_LIST!PIXIE模型, 
)

REM SAM 模型检查
IF EXIST "%SAM_PATH%" (
    echo [√] %MODEL_SAM_NAME%已存在
    SET /A EXISTING_FILES+=1
) ELSE (
    echo [×] 缺少%MODEL_SAM_NAME%
    SET /A MISSING_FILES+=1
    SET MISSING_LIST=!MISSING_LIST!%MODEL_SAM_NAME%, 
)

REM GroundingDINO 模型检查
IF EXIST "%DINO_PATH%" (
    echo [√] %MODEL_DINO_NAME%已存在
    SET /A EXISTING_FILES+=1
) ELSE (
    echo [×] 缺少%MODEL_DINO_NAME%
    SET /A MISSING_FILES+=1
    SET MISSING_LIST=!MISSING_LIST!%MODEL_DINO_NAME%, 
)

REM OpenPose 模型检查
IF EXIST "%OPENPOSE_PATH%" (
    echo [√] %MODEL_OPENPOSE_NAME%已存在
    SET /A EXISTING_FILES+=1
) ELSE (
    echo [×] 缺少%MODEL_OPENPOSE_NAME%
    SET /A MISSING_FILES+=1
    SET MISSING_LIST=!MISSING_LIST!%MODEL_OPENPOSE_NAME%, 
)

REM hyperIQA 模型检查
IF EXIST "%HYPERIQA_PATH%" (
    echo [√] %MODEL_HYPERIQA_NAME%已存在
    SET /A EXISTING_FILES+=1
) ELSE (
    echo [×] 缺少%MODEL_HYPERIQA_NAME%
    SET /A MISSING_FILES+=1
    SET MISSING_LIST=!MISSING_LIST!%MODEL_HYPERIQA_NAME%, 
)

echo.
echo 已存在 %EXISTING_FILES%/%TOTAL_FILES% 个模型文件
IF %MISSING_FILES% GTR 0 (
    echo 需要下载 %MISSING_FILES% 个模型文件
    echo 缺少的模型: !MISSING_LIST:~0,-2!
    echo.
    echo ************************************************************
    echo *                  开始下载所需资源文件                     *
    echo *            下载时间取决于您的网络连接速度                 *
    echo *                    请耐心等待...                         *
    echo ************************************************************
) ELSE (
    echo 所有模型文件已存在，无需下载
    exit /b 0
)

REM 设置代理选项（可选）
IF DEFINED HTTP_PROXY (
    SET CURL_OPTS=%CURL_OPTS% --proxy %HTTP_PROXY%
    SET REQUESTS_CA_BUNDLE=%PROJECT_DIR%\certs\cacert.pem
)

REM 下载 Neural Haircut 文件
cd resource\NeuralHaircut
echo 正在下载Neural Haircut预训练文件夹...
python -m gdown --folder "%NEURAL_HAIRCUT_FOLDER_URL%"

REM 下载 PIXIE 文件
cd NeuralHaircut\PIXIE
IF EXIST "pixie_data.tar" (
    echo PIXIE模型已存在，跳过下载...
) ELSE (
    echo 正在下载PIXIE模型...
    python -m gdown "%PIXIE_MODEL_URL%"
    tar -xvzf pixie_data.tar.gz ./
    del pixie_data.tar.gz
)

REM 下载 Matte-Anything 文件
cd ..\..\Matte-Anything
echo 正在下载Matte-Anything文件...
curl %CURL_OPTS% -o "%SAM_PATH%" "%SAM_MODEL_URL%"
curl -L --retry 5 -C - -o "%DINO_PATH%" "%DINO_MODEL_URL%"
python -m gdown "%MATTE_MODEL_URL%" -O model.pth

REM 下载 OpenPose 模型
cd ..\openpose
python -m gdown "%OPENPOSE_MODEL_URL%" -O models.tar.gz
tar -xvzf models.tar.gz
del models.tar.gz

REM 下载 hyperIQA 模型
cd ..\hyperIQA\pretrained
python -m gdown "%HYPERIQA_MODEL_URL%"

cd ..\..\..
echo.
IF %ERRORLEVEL% EQU 0 (
    echo ************************************************************
    echo *                    下载任务已完成                         *
    echo *              所有资源文件已下载到resource目录             *
    echo ************************************************************
) ELSE (
    echo ************************************************************
    echo *                    下载任务未完成                         *
    echo *                  请检查以上错误信息                       *
    echo ************************************************************
    exit /b 1
)

REM 验证下载的文件完整性
echo 正在验证下载文件...
cd %PROJECT_DIR%
SET VERIFY_FAILED=0

IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\diffusion_prior\model.pt" (
    echo [×] 缺少Neural Haircut Diffusion Prior
    SET /A VERIFY_FAILED+=1
) ELSE (
    echo [√] Neural Haircut Diffusion Prior验证成功
)

IF NOT EXIST "%NEURAL_HAIRCUT_FOLDER_PATH%\pretrained_models\strand_prior\model.pt" (
    echo [×] 缺少Neural Haircut Strand Prior
    SET /A VERIFY_FAILED+=1
) ELSE (
    echo [√] Neural Haircut Strand Prior验证成功
)

IF NOT EXIST "%PIXIE_PATH%" echo 警告：缺少PIXIE模型
IF NOT EXIST "%SAM_PATH%" echo 警告：缺少SAM模型
IF NOT EXIST "%DINO_PATH%" echo 警告：缺少GroundingDINO模型
IF NOT EXIST "%OPENPOSE_PATH%" echo 警告：缺少OpenPose模型
IF NOT EXIST "%HYPERIQA_PATH%" echo 警告：缺少hyperIQA模型

IF %VERIFY_FAILED% GTR 0 (
    echo.
    echo 文件完整性验证失败：%VERIFY_FAILED% 个文件缺失或不完整
    echo 请重新运行脚本下载缺失文件
    exit /b 1
) ELSE (
    echo.
    echo 文件完整性验证成功！
)