@echo off
setlocal EnableExtensions
cd /d "%~dp0"

rem NLTK 3.10.1+ inisec blocks imports whose files live under CWD. With the venv at
rem .venv that wrongly blocks site-packages when we run from the project root.
set "NLTK_DISABLE_IMPORT_SECURITY=1"

title NovelAI preset optimize
echo === NovelAI preset optimize ===
echo Project: %CD%
echo.

for %%F in (examv2.txt examv2.toml novelai_preset_optimize.py text_metrics.py .env) do (
  if not exist "%%F" (
    echo MISSING: %%F ^(must sit next to this .bat^)
    goto :fail
  )
)

where python >nul 2>&1
if errorlevel 1 (
  echo Python not found on PATH.
  echo Install Python 3 from https://www.python.org/downloads/
  echo and check "Add python.exe to PATH", then re-run this .bat.
  goto :fail
)

rem optunahub/GitPython needs a real git.exe to fetch AutoSampler
where git >nul 2>&1
if errorlevel 1 (
  echo git not found on PATH.
  echo Install Git for Windows from https://git-scm.com/download/win
  echo and ensure "git" is on PATH, then re-run this .bat.
  goto :fail
)

if not exist ".venv\Scripts\python.exe" (
  echo Creating virtualenv at .venv
  python -m venv .venv
  if errorlevel 1 goto :fail
)

set "PY=.venv\Scripts\python.exe"
set "MARKER=.venv\.deps_ok"
rem Bump when install recipe changes (forces reinstall path).
set "DEPS_STAMP=5"

rem Happy path: marker matches stamp + light import check (no spaCy model load)
if exist "%MARKER%" (
  findstr /x /c:"%DEPS_STAMP%" "%MARKER%" >nul 2>&1
  if not errorlevel 1 (
    "%PY%" -c "import importlib.util as u; mods=('optuna','optunahub','cmaes','scipy','torch','spacy','textdescriptives','openai','tokenizers','numpy','dotenv','textstat','lexicalrichness','nrclex','requests','lexical_diversity'); missing=[m for m in mods if u.find_spec(m) is None]; raise SystemExit(1 if missing else 0)" >nul 2>&1
    if not errorlevel 1 (
      echo Dependencies ready — starting.
      goto :run
    )
  )
)

rem Diagnose what is missing (shown so reinstall is not a mystery)
echo Checking dependencies...
"%PY%" -c "import importlib.util as u; mods=['optuna','optunahub','cmaes','scipy','torch','spacy','textdescriptives','openai','tokenizers','numpy','dotenv','textstat','lexicalrichness','nrclex','requests','lexical_diversity']; miss=[m for m in mods if u.find_spec(m) is None]; print('missing packages:', ', '.join(miss) if miss else '(none)'); import sys; sys.exit(1 if miss else 0)"
set "NEED_PIP=%ERRORLEVEL%"

"%PY%" -c "import spacy; spacy.load('en_core_web_sm')" >nul 2>&1
set "NEED_SPACY=%ERRORLEVEL%"
if not "%NEED_SPACY%"=="0" echo missing spaCy model: en_core_web_sm

if "%NEED_PIP%"=="0" if "%NEED_SPACY%"=="0" (
  echo All deps present — writing marker.
  >"%MARKER%" echo %DEPS_STAMP%
  goto :run
)

if not "%NEED_PIP%"=="0" (
  echo Upgrading pip...
  "%PY%" -m pip install --upgrade pip "setuptools>=69,<81" wheel
  if errorlevel 1 goto :fail

  rem Never build numpy from source: old sdists on py3.13 can ICE the compiler.
  echo Installing numpy ^(binary wheel only^)...
  "%PY%" -m pip install --only-binary=:all: "numpy>=2.1,<3"
  if errorlevel 1 goto :fail

  echo Installing Python deps ^(first run / repair^)...
  rem textstat → nltk → regex; pin nltk around the 3.10.1 inisec false positive
  rem when the venv lives under the project root.
  "%PY%" -m pip install --prefer-binary ^
    "optuna>=4.0" ^
    optunahub ^
    cmaes ^
    scipy ^
    python-dotenv ^
    textstat ^
    "nltk>=3.9,<3.10.1" ^
    regex ^
    lexicalrichness ^
    "lexical-diversity" ^
    nrclex ^
    spacy ^
    textdescriptives ^
    openai ^
    "tokenizers>=0.20" ^
    requests ^
    "numpy>=2.1,<3"
  if errorlevel 1 goto :fail

  rem AutoSampler ^(optunahub^) needs torch; CPU wheel is enough.
  echo Installing torch ^(CPU^) for AutoSampler...
  "%PY%" -m pip install --prefer-binary torch --index-url https://download.pytorch.org/whl/cpu
  if errorlevel 1 goto :fail

  rem setuptools 81+ drops pkg_resources (needed by lexical-diversity)
  "%PY%" -m pip install "setuptools>=69,<81" >nul
)

if not "%NEED_SPACY%"=="0" (
  echo Ensuring spaCy en_core_web_sm model...
  "%PY%" -c "import spacy; spacy.load('en_core_web_sm')" >nul 2>&1
  if errorlevel 1 (
    "%PY%" -m spacy download en_core_web_sm
    if errorlevel 1 goto :fail
  ) else (
    echo spaCy model already installed.
  )
) else if not "%NEED_PIP%"=="0" (
  rem Packages reinstalled — model may still be fine; only download if gone
  "%PY%" -c "import spacy; spacy.load('en_core_web_sm')" >nul 2>&1
  if errorlevel 1 (
    echo Ensuring spaCy en_core_web_sm model...
    "%PY%" -m spacy download en_core_web_sm
    if errorlevel 1 goto :fail
  )
)

rem Full validate once after install, then stamp marker for next starts
echo Validating install...
"%PY%" -c "import optuna, optunahub, cmaes, scipy, torch, spacy, textdescriptives, openai, tokenizers, numpy, dotenv, textstat, lexicalrichness, nrclex, requests; from lexical_diversity import lex_div; spacy.load('en_core_web_sm'); optunahub.load_module('samplers/auto_sampler').AutoSampler"
if errorlevel 1 goto :fail
>"%MARKER%" echo %DEPS_STAMP%
echo Install OK.

:run
echo.
echo Starting novelai_preset_optimize.py ^(Ctrl+C to stop^)...
echo.
"%PY%" "novelai_preset_optimize.py"
set "ERR=%ERRORLEVEL%"
echo.
if not "%ERR%"=="0" (
  echo FAILED ^(exit %ERR%^).
) else (
  echo Done.
)
pause
exit /b %ERR%

:fail
echo.
echo FAILED.
pause
exit /b 1
