@echo off
setlocal enabledelayedexpansion

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

REM 检查7-Zip安装路径
SET "SEVENZIP_PATH=C:\Program Files\7-Zip\7z.exe"
IF NOT EXIST "%SEVENZIP_PATH%" (
    SET "SEVENZIP_PATH=C:\Program Files (x86)\7-Zip\7z.exe"
)
IF NOT EXIST "%SEVENZIP_PATH%" (
    echo ERROR: 7-Zip not found
    echo Please install 7-Zip from https://7-zip.org/
    echo Default install locations checked:
    echo - C:\Program Files\7-Zip\7z.exe
    echo - C:\Program Files (x86)\7-Zip\7z.exe
    exit /b 1
)

REM 检查Python是否安装
python --version >nul 2>nul
IF %ERRORLEVEL% NEQ 0 (
    echo ERROR: Python not found
    echo Please install Python and add it to PATH
    exit /b 1
)

REM 安装gdown
echo Installing gdown...
python -m pip install --user gdown --quiet
IF %ERRORLEVEL% NEQ 0 (
    echo ERROR: Failed to install gdown
    echo Please check your internet connection and try again
    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是否可用
python -m gdown --version >nul 2>nul
IF %ERRORLEVEL% NEQ 0 (
    echo ERROR: gdown installation found but not accessible
    echo Please try reinstalling gdown: python -m pip install --user --upgrade gdown
    exit /b 1
)

REM 创建resource目录结构
mkdir resource 2>nul
mkdir resource\NeuralHaircut\diffusion_prior 2>nul
mkdir resource\NeuralHaircut\PIXIE 2>nul
mkdir resource\Matte-Anything 2>nul
mkdir resource\openpose\models 2>nul
mkdir resource\hyperIQA\pretrained 2>nul

echo.
echo Starting downloads...
echo This may take a while depending on your internet connection...
echo.

REM 下载 Neural Haircut 文件
cd resource\NeuralHaircut
echo Downloading Neural Haircut files...
python -m gdown --folder https://drive.google.com/drive/folders/1TCdJ0CKR3Q6LviovndOkJaKm8S1T9F_8
cd diffusion_prior
python -m gdown 1_9EOUXHayKiGH5nkrayncln3d6m1uV7f
cd ..\PIXIE
python -m gdown 1mPcGu62YPc4MdkT8FFiOCP629xsENHZf
REM 使用7-Zip解压tar.gz文件
"%SEVENZIP_PATH%" x pixie_data.tar.gz -y
"%SEVENZIP_PATH%" x pixie_data.tar -y
del pixie_data.tar.gz
del pixie_data.tar

REM 下载 Matte-Anything 文件
cd ..\..\Matte-Anything
echo Downloading Matte-Anything files...
curl -LO https://dl.fbaipublicfiles.com/segment_anything/sam_vit_h_4b8939.pth
curl -LO https://github.com/IDEA-Research/GroundingDINO/releases/download/v0.1.0-alpha/groundingdino_swint_ogc.pth
gdown 1d97oKuITCeWgai2Tf3iNilt6rMSSYzkW -O model.pth

REM 下载 OpenPose 模型
cd ..\openpose
echo Downloading OpenPose models...
python -m gdown 1Yn03cKKfVOq4qXmgBMQD20UMRRRkd_tV -O models.tar.gz
"%SEVENZIP_PATH%" x models.tar.gz -y
"%SEVENZIP_PATH%" x models.tar -y -omodels
del models.tar.gz
del models.tar

REM 下载 hyperIQA 模型
cd ..\hyperIQA\pretrained
echo Downloading hyperIQA models...
python -m gdown 1OOUmnbvpGea0LIGpIWEbOyxfWx6UCiiE

cd ..\..
echo.
echo Download completed successfully!
echo Resource files have been downloaded to the resource directory.
echo.

IF %ERRORLEVEL% NEQ 0 (
    echo.
    echo ERROR: Download failed
    echo Please check your internet connection and try again
    exit /b 1
)