@echo off
SETLOCAL ENABLEDELAYEDEXPANSION
title Llama.cpp Librarian and Benchmark Tool v10

REM =========================================================
REM                  CONFIG & DEFAULTS
REM =========================================================
set "BASE=%~dp0"
set "MODELS_DIR=%BASE%models"
set "BENCH_DB=%BASE%model_library_benchmark.csv"
set "LLAMA_BENCH=%BASE%llama-bench.exe"

REM --- Keywords for Role Detection ---
set "KEYWORDS=NSFW Uncensored Amoral Dark Fallen Rivermind Snowpiercer Abyss Erebus Lewd Heretic"

REM Default Bench Settings
set THREADS=8
set GPU_LAYERS=999

REM =========================================================
REM           HARDWARE DETECTION
REM =========================================================
:GET_STATS
echo Detecting Hardware Context...

set "CPU_NAME=Unknown"
set "GPU_NAME=N/A"
set "VRAM_TOTAL_GB=0"
set "VRAM_FREE_GB=0"
set "RAM_TOTAL_GB=0"
set "DRIVER_VER=N/A"

REM --- CPU ---
for /f "skip=1 delims=" %%a in ('wmic cpu get name 2^>nul') do if not defined CPU_NAME set "CPU_NAME=%%a"
for /f "tokens=* delims= " %%a in ("%CPU_NAME%") do set "CPU_NAME=%%a"

REM --- RAM ---
for /f "skip=1" %%p in ('wmic os get TotalVisibleMemorySize') do if not defined RAM_RAW set "RAM_RAW=%%p"
set /a RAM_TOTAL_GB=%RAM_RAW%/1024/1024

REM --- GPU ---
nvidia-smi --query-gpu=name,memory.total,memory.free,driver_version --format=csv,noheader,nounits > "%BASE%gpu.tmp" 2>nul
if %errorlevel%==0 (
    for /f "tokens=1,2,3,4 delims=," %%a in ('type "%BASE%gpu.tmp"') do (
        set "GPU_NAME=%%a"
        set /a VRAM_TOTAL_GB=%%b/1024
        set /a VRAM_FREE_GB=%%c/1024
        set "DRIVER_VER=%%d"
    )
    del "%BASE%gpu.tmp"
)

REM =========================================================
REM                  MAIN MENU
REM =========================================================
:MAIN_MENU
cls
echo ========================================================
echo             LIBRARIAN AND BENCHMARK v10
echo ========================================================
echo  SYSTEM CONTEXT:
echo   GPU:  %GPU_NAME% (%VRAM_FREE_GB%GB / %VRAM_TOTAL_GB%GB)
echo   RAM:  %RAM_TOTAL_GB%GB System Memory
echo   Drv:  %DRIVER_VER%
echo ========================================================
echo.
echo  [1] Scan Library (Sorted by Size/Quant)
echo  [2] View Leaderboard (Top Speed)
echo  [3] View Full Database
echo  [4] Exit
echo.
set /p "CHOICE=Selection: "

if "%CHOICE%"=="1" goto :SCAN_MODELS
if "%CHOICE%"=="2" goto :LEADERBOARD
if "%CHOICE%"=="3" goto :VIEW_DB
if "%CHOICE%"=="4" exit
goto :MAIN_MENU

REM =========================================================
REM           INTELLIGENT SCANNING ENGINE
REM =========================================================
:SCAN_MODELS
cls
echo Indexing Library...
echo.

set count=0

REM --- Initialize CSV Header if missing ---
if not exist "%BENCH_DB%" (
    echo Model_Name,Size_GB,Quant,Role,Template_Guess,Fit_Status,Ingest_PP512,Gen_TG128,Rating,GPU,Date > "%BENCH_DB%"
)

for %%F in ("%MODELS_DIR%\*.gguf") do (
    set "fname=%%~nxF"
    set "fpath=%%F"
    
    echo !fname! | findstr /i "mmproj" >nul
    if !errorlevel! NEQ 0 (
        set /a count+=1
        
        REM --- 1. Size Calculation ---
        set "size_gb=0"
        for /f "usebackq" %%S in (`powershell -NoProfile -Command "[int][math]::Round(%%~zF / 1GB)"`) do set "size_gb=%%S"
        
        REM --- 2. Fit Analysis ---
        set "fit_status=Unknown"
        if !size_gb! LSS !VRAM_FREE_GB! (
            set "fit_status=GPU_FAST"
        ) else (
             set /a total_mem=!VRAM_FREE_GB! + !RAM_TOTAL_GB!
             if !size_gb! LSS !total_mem! (
                 set "fit_status=CPU_SPLIT"
             ) else (
                 set "fit_status=TOO_LARGE"
             )
        )

        REM --- 3. Role Detection ---
        set "role=General"
        echo !fname! | findstr /I "%KEYWORDS%" >nul
        if !errorlevel! EQU 0 set "role=Specialized/NSFW"

        REM --- 4. Template Heuristics ---
        set "tmpl=Unknown"
        echo !fname! | findstr /I "llama-3 llama3 dark-nexus rivermind snowpiercer" >nul
        if !errorlevel! EQU 0 set "tmpl=Llama3"
        echo !fname! | findstr /I "mistral v0.1 v0.2 v0.3" >nul
        if !errorlevel! EQU 0 set "tmpl=Mistral"
        echo !fname! | findstr /I "chatml heremes dolphin yi looking-glass" >nul
        if !errorlevel! EQU 0 set "tmpl=ChatML"
        echo !fname! | findstr /I "alpaca" >nul
        if !errorlevel! EQU 0 set "tmpl=Alpaca"
        echo !fname! | findstr /I "vicuna" >nul
        if !errorlevel! EQU 0 set "tmpl=Vicuna"
        echo !fname! | findstr /I "gemma" >nul
        if !errorlevel! EQU 0 set "tmpl=Gemma"
        echo !fname! | findstr /I "qwen" >nul
        if !errorlevel! EQU 0 set "tmpl=Qwen"
        echo !fname! | findstr /I "olmo" >nul
        if !errorlevel! EQU 0 set "tmpl=Olmo"
        echo !fname! | findstr /I "phi" >nul
        if !errorlevel! EQU 0 set "tmpl=Phi"
        echo !fname! | findstr /I "glm" >nul
        if !errorlevel! EQU 0 set "tmpl=GLM"

        REM --- 5. B-Parameter Tag (Reversed for Correctness) ---
        set "b_tag=Other"
        REM Loop from Smallest to Largest so exact matches overwrite partial matches
        REM e.g. "12B" will overwrite "1B" (from v1b)
        for %%B in (0.5B 1B 1.5B 2B 3B 4B 6B 7B 8B 9B 11B 12B 13B 14B 16B 20B 24B 27B 30B 32B 34B 70B 72B) do (
            echo !fname! | findstr /I "%%B" >nul
            if !errorlevel! EQU 0 set "b_tag=%%B"
        )

        REM --- 6. Quantization Detection ---
        set "quant=Unknown"
        echo !fname! | findstr /I "f32" >nul
        if !errorlevel! EQU 0 set "quant=F32"
        echo !fname! | findstr /I "f16" >nul
        if !errorlevel! EQU 0 set "quant=F16"
        echo !fname! | findstr /I "Q8_0" >nul
        if !errorlevel! EQU 0 set "quant=Q8_0"
        echo !fname! | findstr /I "Q6_K" >nul
        if !errorlevel! EQU 0 set "quant=Q6_K"
        echo !fname! | findstr /I "Q5_K_M" >nul
        if !errorlevel! EQU 0 set "quant=Q5_K_M"
        echo !fname! | findstr /I "Q5_K_S" >nul
        if !errorlevel! EQU 0 set "quant=Q5_K_S"
        echo !fname! | findstr /I "Q5_0" >nul
        if !errorlevel! EQU 0 set "quant=Q5_0"
        echo !fname! | findstr /I "Q4_K_M" >nul
        if !errorlevel! EQU 0 set "quant=Q4_K_M"
        echo !fname! | findstr /I "Q4_K_S" >nul
        if !errorlevel! EQU 0 set "quant=Q4_K_S"
        echo !fname! | findstr /I "Q4_0" >nul
        if !errorlevel! EQU 0 set "quant=Q4_0"
        echo !fname! | findstr /I "Q3_K_M" >nul
        if !errorlevel! EQU 0 set "quant=Q3_K_M"
        echo !fname! | findstr /I "IQ3" >nul
        if !errorlevel! EQU 0 set "quant=IQ3"
        echo !fname! | findstr /I "Q2_K" >nul
        if !errorlevel! EQU 0 set "quant=Q2_K"

        REM --- Store in Array ---
        set "m_name[!count!]=!fname!"
        set "m_path[!count!]=!fpath!"
        set "m_size[!count!]=!size_gb!"
        set "m_fit[!count!]=!fit_status!"
        set "m_role[!count!]=!role!"
        set "m_tmpl[!count!]=!tmpl!"
        set "m_quant[!count!]=!quant!"
        set "m_btag[!count!]=!b_tag!"
    )
)

if %count% EQU 0 (
    echo No models found.
    pause
    goto :MAIN_MENU
)

REM =========================================================
REM           NESTED SORT DISPLAY (Size -> Quant)
REM =========================================================
echo.
echo  ID   STATUS      QUANT     ROLE       MODEL
echo -------------------------------------------------------------

for %%S in (72B 70B 34B 32B 30B 27B 24B 20B 16B 14B 13B 12B 11B 9B 8B 7B 6B 4B 3B 2B 1.5B 1B 0.5B Other) do (
    call :PRINT_SIZE_GROUP "%%S"
)

echo.
echo [A] Benchmark ALL models
echo [B] Back to Menu
echo.
set /p "B_SEL=Select ID to benchmark (or A): "

if /I "%B_SEL%"=="B" goto :MAIN_MENU
if /I "%B_SEL%"=="A" goto :BENCH_ALL

if defined m_path[%B_SEL%] (
    call :RUN_BENCH "%B_SEL%"
    pause
    goto :SCAN_MODELS
)
goto :SCAN_MODELS

REM =========================================================
REM        HELPER: PRINT SIZE GROUP
REM =========================================================
:PRINT_SIZE_GROUP
set "size_found=0"
for /L %%i in (1,1,%count%) do (
    if "!m_btag[%%i]!"=="%~1" set "size_found=1"
)
if !size_found! EQU 1 (
    echo.
    echo --- %~1 CLASS MODELS ---
    for %%Q in (F32 F16 Q8_0 Q6_K Q5_K_M Q5_K_S Q5_0 Q4_K_M Q4_K_S Q4_0 Q3_K_M IQ3 Q2_K Unknown) do (
        for /L %%i in (1,1,%count%) do (
             if "!m_btag[%%i]!"=="%~1" (
                 if /I "!m_quant[%%i]!"=="%%Q" (
                     echo  [%%i] [!m_fit[%%i]!] [!m_quant[%%i]!] [!m_role[%%i]!] !m_name[%%i]!
                 )
             )
        )
    )
)
exit /b

REM =========================================================
REM                  BENCHMARK RUNNER (REGEX FIXED)
REM =========================================================
:BENCH_ALL
for /L %%i in (1,1,%count%) do (
    call :RUN_BENCH "%%i"
)
echo All tasks complete.
pause
goto :SCAN_MODELS

:RUN_BENCH
set "idx=%~1"
set "target_name=!m_name[%idx%]!"
set "target_path=!m_path[%idx%]!"
set "target_fit=!m_fit[%idx%]!"
set "target_quant=!m_quant[%idx%]!"
set "target_tmpl=!m_tmpl[%idx%]!"
set "target_role=!m_role[%idx%]!"
set "target_size=!m_size[%idx%]!"
set "target_btag=!m_btag[%idx%]!"

cls
echo ========================================================
echo  BENCHMARKING: %target_name%
echo ========================================================
echo  Class:    %target_btag%
echo  Quant:    %target_quant%
echo  Fit:      %target_fit%
echo --------------------------------------------------------

if exist "%BASE%bench_temp.txt" del "%BASE%bench_temp.txt"

REM Run llama-bench
"%LLAMA_BENCH%" -m "%target_path%" -t %THREADS% -ngl %GPU_LAYERS% > "%BASE%bench_temp.txt" 2>&1

type "%BASE%bench_temp.txt"

REM --- PARSING WITH POWERSHELL REGEX (ROBUST) ---
set "PP512=0"
set "TG128=0"

for /f "usebackq tokens=*" %%v in (`powershell -NoProfile -Command "$c = Get-Content '%BASE%bench_temp.txt' | Select-String '\|\s*pp512\s*\|\s*([0-9\.]+)'; if ($c) { $c.Matches.Groups[1].Value } else { '0' }"`) do set "PP512=%%v"

for /f "usebackq tokens=*" %%v in (`powershell -NoProfile -Command "$c = Get-Content '%BASE%bench_temp.txt' | Select-String '\|\s*tg128\s*\|\s*([0-9\.]+)'; if ($c) { $c.Matches.Groups[1].Value } else { '0' }"`) do set "TG128=%%v"

REM --- CALCULATE RATING ---
set "rating=SLOW"
for /f "tokens=1 delims=." %%n in ("!TG128!") do set "INT_TG=%%n"

if !INT_TG! GEQ 10 set "rating=USABLE"
if !INT_TG! GEQ 25 set "rating=SMOOTH"
if !INT_TG! GEQ 50 set "rating=GODLIKE"

echo.
echo  ================================
echo   SPEED RATING: %rating% (!TG128! t/s)
echo  ================================
echo.

REM Save to CSV
echo logging to %BENCH_DB%...
for /f "tokens=2 delims==" %%I in ('wmic os get localdatetime /value') do set datetime=%%I
set "TS=%datetime:~0,4%-%datetime:~4,2%-%datetime:~6,2%"

echo %target_name%,%target_size%,%target_quant%,%target_role%,%target_tmpl%,%target_fit%,!PP512!,!TG128!,%rating%,%GPU_NAME%,%TS% >> "%BENCH_DB%"

echo Done.
exit /b

REM =========================================================
REM                  LEADERBOARD
REM =========================================================
:LEADERBOARD
cls
echo ========================================================
echo                  SPEED LEADERBOARD
echo ========================================================
if not exist "%BENCH_DB%" (
    echo No database found. Please scan/benchmark first.
    pause
    goto :MAIN_MENU
)

echo.
echo TOP 15 FASTEST MODELS (Generation Speed):
echo -------------------------------------------------------------------------------
REM PowerShell magic with ErrorAction Ignore for bad data rows
powershell -NoProfile -Command ^
    "$data = Import-Csv '%BENCH_DB%';" ^
    "$clean_data = $data | Where-Object { $_.Gen_TG128 -match '^\d+(\.\d+)?$' };" ^
    "$sorted = $clean_data | Sort-Object -Property @{Expression={[double]$_.Gen_TG128}; Descending=$true};" ^
    "$sorted | Select-Object -First 15 Model_Name, Size_GB, Quant, Gen_TG128, Rating | Format-Table -AutoSize"

echo.
pause
goto :MAIN_MENU

REM =========================================================
REM                  VIEW DATABASE
REM =========================================================
:VIEW_DB
cls
if not exist "%BENCH_DB%" (
    echo Database not found.
    pause
    goto :MAIN_MENU
)
echo Opening Database in default CSV viewer...
start "" "%BENCH_DB%"
goto :MAIN_MENU