@echo off
setlocal EnableExtensions EnableDelayedExpansion
cd /d "%~dp0"
chcp 65001 >nul 2>nul
title OmniVoice 1-Click Installer

set "APP_NAME=OmniVoice"
set "INSTALL_DIR=%~dp0OmniVoice"
set "VENV_DIR=%INSTALL_DIR%\venv"
set "LAUNCHER=%INSTALL_DIR%\run_omnivoice_demo.bat"

echo =====================================================
echo                 OmniVoice 1-Click Installer
echo =====================================================
echo.
echo This script will:
echo   1. Find a compatible Python installation ^(3.10 / 3.11 / 3.12^)
echo   2. Create a fresh virtual environment
echo   3. Install PyTorch CUDA 12.8 and torchaudio
echo   4. Install OmniVoice
echo   5. Create a launcher BAT for the web UI
echo.
echo IMPORTANT:
echo   This installer is made for NVIDIA GPUs.
echo   You do NOT need to install the CUDA Toolkit manually.
echo   You only need a recent NVIDIA driver.
echo.
echo Install folder:
echo   "%INSTALL_DIR%"
echo.

call :find_python
if errorlevel 1 call :try_install_python
if errorlevel 1 goto :python_missing

echo =====================================================
echo                    SYSTEM CHECK
echo =====================================================
echo.

echo Python:
echo   FOUND
echo   Path:
echo     "%PYTHON_EXE%"

for /f "delims=" %%i in ('"%PYTHON_EXE%" -c "import sys; print(str(sys.version_info[0]) + '.' + str(sys.version_info[1]) + '.' + str(sys.version_info[2]))" 2^>nul') do set "PYTHON_VER=%%i"

if defined PYTHON_VER (
    echo   Version:
    echo     !PYTHON_VER!
) else (
    echo   Version:
    echo     Could not read version, but Python was found.
)

echo.

echo FFmpeg:
where ffmpeg >nul 2>nul
if errorlevel 1 (
    echo   NOT FOUND
    echo   Trying to install FFmpeg with winget...
    where winget >nul 2>nul
    if errorlevel 1 (
        echo.
        echo ERROR: winget was not found, so FFmpeg cannot be installed automatically.
        echo Please install FFmpeg manually, then run this installer again.
        echo.
        goto :fail
    )
    winget install -e --id Gyan.FFmpeg --accept-package-agreements --accept-source-agreements --disable-interactivity
    if errorlevel 1 (
        echo.
        echo ERROR: Failed to install FFmpeg.
        echo Please install FFmpeg manually, then run this installer again.
        echo.
        goto :fail
    )
) else (
    echo   FOUND
)

echo.

echo NVIDIA:
where nvidia-smi >nul 2>nul
if errorlevel 1 (
    echo   NOT FOUND
    echo.
    echo ERROR: NVIDIA GPU / driver was not detected.
    echo This installer is NVIDIA-only.
    echo Please install/update your NVIDIA driver, reboot, and run this again.
    echo.
    goto :fail
)

echo   FOUND
echo.
nvidia-smi
echo.

echo =====================================================
echo Starting OmniVoice install now...
echo =====================================================
echo.

if exist "%VENV_DIR%" (
    echo Removing old virtual environment...
    rmdir /s /q "%VENV_DIR%"
    if exist "%VENV_DIR%" (
        echo.
        echo ERROR: Failed to remove old venv.
        echo Close any open terminals using OmniVoice and try again.
        echo.
        goto :fail
    )
)

if not exist "%INSTALL_DIR%" mkdir "%INSTALL_DIR%"
if not exist "%INSTALL_DIR%\hf_cache" mkdir "%INSTALL_DIR%\hf_cache"
if not exist "%INSTALL_DIR%\torch_cache" mkdir "%INSTALL_DIR%\torch_cache"

set "HF_HOME=%INSTALL_DIR%\hf_cache"
set "HF_HUB_DISABLE_SYMLINKS_WARNING=1"
set "TORCH_HOME=%INSTALL_DIR%\torch_cache"
set "PIP_DISABLE_PIP_VERSION_CHECK=1"
set "PYTHONUTF8=1"
set "PYTHONIOENCODING=utf-8"

echo Creating virtual environment...
"%PYTHON_EXE%" -m venv "%VENV_DIR%"
if errorlevel 1 (
    echo.
    echo ERROR: Failed to create the virtual environment.
    echo.
    goto :fail
)

set "VENV_PY=%VENV_DIR%\Scripts\python.exe"

if not exist "%VENV_PY%" (
    echo.
    echo ERROR: Could not find the venv Python executable.
    echo.
    goto :fail
)

echo.
echo Upgrading pip, setuptools, and wheel...
"%VENV_PY%" -m ensurepip --upgrade
if errorlevel 1 goto :fail

"%VENV_PY%" -m pip install --upgrade pip setuptools wheel
if errorlevel 1 goto :fail

echo.
echo Installing PyTorch CUDA 12.8 build...
"%VENV_PY%" -m pip install torch==2.8.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
if errorlevel 1 goto :fail

echo.
echo Installing OmniVoice from PyPI...
"%VENV_PY%" -m pip install --prefer-binary --upgrade omnivoice
if errorlevel 1 goto :fail

echo.
echo Re-applying PyTorch CUDA 12.8 build...
"%VENV_PY%" -m pip install torch==2.8.0 torchaudio==2.8.0 --index-url https://download.pytorch.org/whl/cu128
if errorlevel 1 goto :fail

echo.
echo Running CUDA test...
"%VENV_PY%" -c "import torch, omnivoice; print('Torch:', torch.__version__); print('CUDA available:', torch.cuda.is_available()); print('CUDA:', torch.version.cuda); assert torch.cuda.is_available(), 'CUDA is not available'; print('GPU:', torch.cuda.get_device_name(0)); print('OmniVoice import OK')"
if errorlevel 1 goto :fail

echo.
echo Creating launcher...
(
    echo @echo off
    echo setlocal EnableExtensions
    echo cd /d "%%~dp0"
    echo chcp 65001 ^>nul 2^>nul
    echo title OmniVoice
    echo.
    echo if not exist "%%~dp0hf_cache" mkdir "%%~dp0hf_cache"
    echo if not exist "%%~dp0torch_cache" mkdir "%%~dp0torch_cache"
    echo.
    echo set "HF_HOME=%%~dp0hf_cache"
    echo set "HF_HUB_DISABLE_SYMLINKS_WARNING=1"
    echo set "TORCH_HOME=%%~dp0torch_cache"
    echo set "PIP_DISABLE_PIP_VERSION_CHECK=1"
    echo set "PYTHONUTF8=1"
    echo set "PYTHONIOENCODING=utf-8"
    echo.
    echo echo =====================================================
    echo echo Starting OmniVoice web UI...
    echo echo =====================================================
    echo echo.
    echo echo First launch can take a while because models download from Hugging Face.
    echo echo.
    echo echo Open this in your browser:
    echo echo   http://127.0.0.1:8001
    echo echo.
    echo start "" "http://127.0.0.1:8001"
    echo.
    echo if exist "%%~dp0venv\Scripts\omnivoice-demo.exe" ^(
    echo     "%%~dp0venv\Scripts\omnivoice-demo.exe" --ip 127.0.0.1 --port 8001
    echo ^) else ^(
    echo     "%%~dp0venv\Scripts\python.exe" -m omnivoice.cli.demo --ip 127.0.0.1 --port 8001
    echo ^)
    echo.
    echo echo.
    echo echo OmniVoice has stopped or closed.
    echo pause
) > "%LAUNCHER%"
if errorlevel 1 goto :fail

echo.
echo =====================================================
echo Installation completed successfully.
echo.
echo To launch OmniVoice, run:
echo   "%LAUNCHER%"
echo =====================================================
echo.
pause
exit /b 0


:find_python
set "PYTHON_EXE="

for %%V in (3.11 3.10 3.12) do (
    py -%%V -c "import sys, struct; assert sys.version_info >= (3,10); assert sys.version_info < (3,13); assert struct.calcsize('P') * 8 == 64; print(sys.executable)" > "%TEMP%\omnivoice_py_path.txt" 2>nul
    if not errorlevel 1 (
        set /p PYTHON_EXE=<"%TEMP%\omnivoice_py_path.txt"
        del "%TEMP%\omnivoice_py_path.txt" >nul 2>nul
        exit /b 0
    )
)

for /f "delims=" %%P in ('where python 2^>nul') do (
    "%%P" -c "import sys, struct; assert sys.version_info >= (3,10); assert sys.version_info < (3,13); assert struct.calcsize('P') * 8 == 64; print(sys.executable)" > "%TEMP%\omnivoice_py_path.txt" 2>nul
    if not errorlevel 1 (
        set /p PYTHON_EXE=<"%TEMP%\omnivoice_py_path.txt"
        del "%TEMP%\omnivoice_py_path.txt" >nul 2>nul
        exit /b 0
    )
)

del "%TEMP%\omnivoice_py_path.txt" >nul 2>nul
exit /b 1


:try_install_python
where winget >nul 2>nul
if errorlevel 1 exit /b 1

echo Python 3.10 / 3.11 / 3.12 was not found.
echo Attempting to install Python 3.11 automatically with winget...
winget install -e --id Python.Python.3.11 --accept-package-agreements --accept-source-agreements --disable-interactivity
if errorlevel 1 exit /b 1

call :find_python
if errorlevel 1 exit /b 1

exit /b 0


:python_missing
echo.
echo ERROR: Python 3.10 / 3.11 / 3.12 could not be found or installed automatically.
echo Please install Python 3.11 manually, then run this BAT again.
echo.
pause
exit /b 1


:fail
echo.
echo =====================================================
echo Installation failed.
echo Check the error message above, then run this installer again.
echo =====================================================
echo.
pause
exit /b 1