@echo off
cd /d %~dp0
set "version_title=AI-Toolkit Installer- Py 3.10.11 embedded"
title %version_title%

REM =========================================================
REM  AI-Toolkit - Windows 1-Click Installer Aitrepreneur
REM  SHOULD:
REM  - Installs Git + Node automatically
REM  - Uses Python 3.10.11 EMBEDDED!!!
REM  - CUDA-aware (12.8 vs 12.6)
REM  - Downloads the 2 launcher bat files for easy launch
REM =========================================================

REM ----------------------------
REM Arguments / retries
REM ----------------------------
set "PIPargs=--no-cache-dir --no-warn-script-location --timeout=1000 --retries 200"
set "CURLargs=--retry 200 --retry-all-errors"
set "UVargs=--no-cache --link-mode=copy"

REM ----------------------------
REM CRITICAL: Clear Python/Conda env vars
REM ----------------------------
set PYTHONPATH=
set PYTHONHOME=
set PYTHON=
set PYTHONSTARTUP=
set PYTHONUSERBASE=
set PIP_CONFIG_FILE=
set PIP_REQUIRE_VIRTUALENV=
set VIRTUAL_ENV=
set CONDA_PREFIX=
set CONDA_DEFAULT_ENV=
set PYENV_ROOT=
set PYENV_VERSION=

REM ----------------------------
REM IMPORTANT NOTE (display only)
REM ----------------------------
echo.
echo *****************************************************
echo  IMPORTANT:
echo  Place this installer in a folder whose FULL path
echo  contains NO SPACES.
echo.
echo  Good:  C:\AI-Toolkit
echo  Bad :  C:\Users\JohnDoe\AI Toolkit
echo *****************************************************
echo.

REM ----------------------------
REM Capture local git/node paths (if already installed)
REM ----------------------------
for /f "delims=" %%G in ('cmd /c "where git.exe 2>nul"') do (set "GIT_PATH=%%~dpG")
for /f "delims=" %%G in ('cmd /c "where node.exe 2>nul"') do (set "NODE_PATH=%%~dpG")

REM ----------------------------
REM Build minimal PATH WITHOUT system Python
REM ----------------------------
set "path="
if defined GIT_PATH set "path=%GIT_PATH%"
if defined NODE_PATH set "path=%path%;%NODE_PATH%"
if exist %windir%\system32 set "path=%PATH%;%windir%\System32"
if exist %windir%\system32\WindowsPowerShell\v1.0 set "path=%PATH%;%windir%\system32\WindowsPowerShell\v1.0"
if exist %localappdata%\Microsoft\WindowsApps set "path=%PATH%;%localappdata%\Microsoft\WindowsApps"

REM ----------------------------
REM Prevent overwriting existing install
REM ----------------------------
if exist AI-Toolkit (
    echo.
    echo WARNING: 'AI-Toolkit' folder already exists!
    echo Move this installer to a new empty folder and run again.
    echo Press any key to Exit...
    pause >nul
    goto :eof
)
if exist ai-toolkit (
    echo.
    echo WARNING: 'ai-toolkit' folder already exists!
    echo Move this installer to a new empty folder and run again.
    echo Press any key to Exit...
    pause >nul
    goto :eof
)
if exist python_embeded (
    echo.
    echo WARNING: 'python_embeded' folder already exists!
    echo Move this installer to a new empty folder and run again.
    echo Press any key to Exit...
    pause >nul
    goto :eof
)

REM ----------------------------
REM Capture start time
REM ----------------------------
for /f %%i in ('powershell -command "Get-Date -Format HH:mm:ss"') do set start=%%i

REM ----------------------------
REM Skip downloading Git LFS files
REM ----------------------------
set GIT_LFS_SKIP_SMUDGE=1

REM ----------------------------
REM Choose GPU series / CUDA index
REM ----------------------------
echo.
echo --------------------------------------------------
echo Select your NVIDIA GPU generation:
echo --------------------------------------------------
echo  1 ^> RTX 50-series  (PyTorch cu128, CUDA Toolkit 12.8)
echo  2 ^> RTX 40/30/20   (PyTorch cu126, CUDA Toolkit 12.6)
echo --------------------------------------------------
set /p gpu_choice="Enter choice [1/2]: "

if not "%gpu_choice%"=="1" if not "%gpu_choice%"=="2" (
    echo.
    echo ERROR: Invalid choice. Enter 1 or 2.
    echo Press any key to exit...
    pause >nul
    exit /b
)

if "%gpu_choice%"=="1" (
    set "TORCH_INDEX=https://download.pytorch.org/whl/cu128"
    set "CUDA_REQ=12.8"
    set "CUDA_URL=https://developer.nvidia.com/cuda-12-8-0-download-archive"
) else (
    set "TORCH_INDEX=https://download.pytorch.org/whl/cu126"
    set "CUDA_REQ=12.6"
    set "CUDA_URL=https://developer.nvidia.com/cuda-12-6-0-download-archive"
)

REM Torch versions (change whenever you want)
set "TORCH_VER=2.7.0"
set "TV_VER=0.22.0"
set "TA_VER=2.7.0"

echo.
echo ==================================================
echo CUDA Toolkit requirement: %CUDA_REQ%
echo Download: %CUDA_URL%
echo ==================================================
echo.
pause

REM ----------------------------
REM Clear pip/uv cache (helps reliability)
REM ----------------------------
call :clear_pip_uv_cache

REM ----------------------------
REM Install/Update Git + Node.js via winget
REM ----------------------------
call :install_git
call :nodejs_install

REM ----------------------------
REM Check git is installed
REM ----------------------------
git.exe --version >nul 2>&1
if errorlevel 1 (
    echo.
    echo ERROR: git is NOT installed.
    echo Please install git manually from https://git-scm.com/ and rerun.
    echo Press any key to Exit...
    pause >nul
    exit /b
)

REM ----------------------------
REM Clone AI-Toolkit repo directly into "AI-Toolkit"
REM ----------------------------
call :ai_toolkit_clone

REM ----------------------------
REM Install Python 3.10.11 embedded inside AI-Toolkit\python_embeded
REM ----------------------------
call :python_embedded_install_31011

REM ----------------------------
REM Setup venv + deps + UI inside AI-Toolkit
REM ----------------------------
call :ai_toolkit_install

REM ----------------------------
REM Clear caches again
REM ----------------------------
call :clear_pip_uv_cache

REM ----------------------------
REM Capture end time / runtime
REM ----------------------------
for /f %%i in ('powershell -command "Get-Date -Format HH:mm:ss"') do set end=%%i
for /f %%i in ('powershell -command "(New-TimeSpan -Start (Get-Date ''%start%'') -End (Get-Date ''%end%'')).TotalSeconds"') do set diff=%%i

echo.
echo ::::::::::::::: Installation Complete :::::::::::::::
echo Total Running Time: %diff% seconds
echo.
echo Launchers downloaded into: .\AI-Toolkit\
echo - LAUNCHER-TOOLKIT.bat
echo - SECURE_LAUNCHER-TOOLKIT.bat
echo.

echo Starting WebUI...
echo.

REM =========================
REM CHANGE: Actually launch the WebUI now
REM =========================
if exist "%~dp0AI-Toolkit\LAUNCHER-TOOLKIT.bat" (
    cd /d "%~dp0AI-Toolkit"
    call "LAUNCHER-TOOLKIT.bat"
) else (
    echo WARNING: LAUNCHER-TOOLKIT.bat not found in AI-Toolkit folder.
    echo Press any key to exit...
    pause >nul
)

goto :eof

REM =========================
REM FUNCTIONS
REM =========================

:clear_pip_uv_cache
if exist "%localappdata%\pip\cache" rd /s /q "%localappdata%\pip\cache" && md "%localappdata%\pip\cache"
if exist "%localappdata%\uv\cache"  rd /s /q "%localappdata%\uv\cache"  && md "%localappdata%\uv\cache"
echo.
echo Clearing pip and uv cache: Done
echo.
goto :eof

:install_git
echo Installing/Updating Git...
where winget.exe >nul 2>&1
if errorlevel 1 (
    echo.
    echo ERROR: winget.exe not found.
    echo Install "App Installer" from Microsoft Store, then rerun.
    echo Press any key to exit...
    pause >nul
    exit /b
)
winget.exe install --id Git.Git -e --source winget
set "path=%PATH%;%ProgramFiles%\Git\cmd"
echo.
goto :eof

:nodejs_install
echo Installing/Updating Node.js...
where winget.exe >nul 2>&1
if errorlevel 1 (
    echo.
    echo ERROR: winget.exe not found.
    echo Install "App Installer" from Microsoft Store, then rerun.
    echo Press any key to exit...
    pause >nul
    exit /b
)
winget.exe install --id=OpenJS.NodeJS -e
set "path=%PATH%;%ProgramFiles%\nodejs"
title %version_title%
echo.
goto :eof

:ai_toolkit_clone
echo Installing AI-Toolkit...
echo.

cd /d "%~dp0"

echo Cloning AI-Toolkit repo into "AI-Toolkit"...
git.exe clone https://github.com/ostris/ai-toolkit.git "AI-Toolkit"
if errorlevel 1 (
    echo.
    echo ERROR: git clone failed.
    echo Press any key to exit...
    pause >nul
    exit /b
)

goto :eof

:python_embedded_install_31011
echo Installing Python embedded 3.10.11...
echo.

where curl.exe >nul 2>&1
if errorlevel 1 (
    echo ERROR: curl.exe not found. You need Windows 10 v1803+ or add curl to PATH.
    echo Press any key to exit...
    pause >nul
    exit /b
)

set "PY_VER=3.10.11"
set "PY_ZIP=python-%PY_VER%-embed-amd64.zip"
set "PY_URL=https://www.python.org/ftp/python/%PY_VER%/%PY_ZIP%"

REM work inside AI-Toolkit folder (repo root)
pushd "AI-Toolkit"

curl.exe -OL "%PY_URL%" --ssl-no-revoke %CURLargs%
if not exist "%PY_ZIP%" (
    echo.
    echo ERROR: Failed downloading Python embedded zip:
    echo %PY_URL%
    echo Press any key to exit...
    pause >nul
    popd
    exit /b
)

md python_embeded && cd python_embeded
tar.exe -xf "..\%PY_ZIP%"
erase "..\%PY_ZIP%"

echo.
echo Installing pip...
echo.
curl.exe -sSL https://bootstrap.pypa.io/get-pip.py -o get-pip.py --ssl-no-revoke %CURLargs%

REM Make embedded python allow site-packages (minimal + reliable)
(
  echo ..
  echo DLLs
  echo Lib\site-packages
  echo Lib
  echo Scripts
  echo python310.zip
  echo .
  echo import site
) > python310._pth

"%CD%\python.exe" -I get-pip.py %PIPargs%
"%CD%\python.exe" -I -m pip install uv==0.9.7 %PIPargs%
"%CD%\python.exe" -I -m pip install --upgrade pip %PIPargs%
"%CD%\python.exe" -I -m pip install virtualenv %PIPargs%

echo.
echo Embedded Python version:
"%CD%\python.exe" -I -c "import sys; print(sys.version)"
echo.

popd
goto :eof

:ai_toolkit_install
echo Setting up AI-Toolkit (venv + deps + UI)...
echo.

cd /d "%~dp0AI-Toolkit"

REM Download your launcher BATs (repo root)
echo Downloading launcher BAT files...
curl.exe --ssl-no-revoke -L -o LAUNCHER-TOOLKIT.bat ^
  https://huggingface.co/Aitrepreneur/FLX/resolve/main/LAUNCHER-TOOLKIT.bat %CURLargs%
if errorlevel 1 (
    echo ERROR: Failed to download LAUNCHER-TOOLKIT.bat
    echo Press any key to exit...
    pause >nul
    exit /b
)

curl.exe --ssl-no-revoke -L -o SECURE_LAUNCHER-TOOLKIT.bat ^
  https://huggingface.co/Aitrepreneur/FLX/resolve/main/SECURE_LAUNCHER-TOOLKIT.bat %CURLargs%
if errorlevel 1 (
    echo ERROR: Failed to download SECURE_LAUNCHER-TOOLKIT.bat
    echo Press any key to exit...
    pause >nul
    exit /b
)

REM Create venv from embedded Python (AI-Toolkit\python_embeded)
"python_embeded\python.exe" -I -m virtualenv --clear --no-download venv
if errorlevel 1 (
    echo.
    echo ERROR: Failed to create venv.
    echo Press any key to exit...
    pause >nul
    exit /b
)

call venv\Scripts\activate.bat

REM Clear env again inside venv
set PYTHONPATH=
set PYTHONHOME=
set PIP_CONFIG_FILE=

REM Ensure uv exists
if exist "python_embeded\Scripts\uv.exe" pip install uv==0.9.7 %PIPargs%

if exist "venv\Scripts\uv.exe" (
    echo Using uv for package installation...

    uv pip install torch==%TORCH_VER% torchvision==%TV_VER% torchaudio==%TA_VER% --index-url %TORCH_INDEX% %UVargs%
    if errorlevel 1 goto :toolkit_fail

    uv pip install -r requirements.txt %UVargs%
    if errorlevel 1 goto :toolkit_fail

    uv pip install poetry-core %UVargs%
    if errorlevel 1 goto :toolkit_fail
) else (
    echo uv not available. Using pip...

    pip install torch==%TORCH_VER% torchvision==%TV_VER% torchaudio==%TA_VER% --index-url %TORCH_INDEX% %PIPargs%
    if errorlevel 1 goto :toolkit_fail

    pip install -r requirements.txt %PIPargs%
    if errorlevel 1 goto :toolkit_fail

    pip install poetry-core %PIPargs%
    if errorlevel 1 goto :toolkit_fail
)

python -c "import torch, torchvision, torchaudio; print('OK', torch.__version__)"
if errorlevel 1 goto :toolkit_fail

REM Build and start UI
cd ui
if not exist node_modules (
    echo First run: npm install...
    npm install
    if errorlevel 1 goto :toolkit_fail
)

echo Starting UI...
npm run build_and_start
goto :eof

:toolkit_fail
echo.
echo ERROR: Installation failed during Python/Node dependency setup.
echo Common causes:
echo - wrong GPU choice (cu126 vs cu128)
echo - missing CUDA Toolkit %CUDA_REQ%
echo - antivirus blocking downloads
echo - unstable internet
echo.
echo Press any key to exit...
pause >nul
exit /b