#!/bin/bash

MODEL_PATH="/home/models/gguf"
CONTEXT="-c 4000"
# Bezpieczny zapas VRAM (w MB) na bufor kontekstu (KV cache) i obciążenie systemowe
OVERHEAD_MB=1536 

# --- LOGIKA STEROWANIA ---

if [ "$1" == "stop" ]; then
    echo "🛑 Zatrzymywanie serwera llama i zwalnianie VRAM..."
    pkill -f "llama-server" > /dev/null 2>&1
    echo "✅ System czysty. VRAM został zwolniony."
    exit 0

elif [ "$1" == "logs" ]; then
    echo "📋 Podgląd logów serwera..."
    tail -f ~/server.log
    exit 0

else
    # --- SPRAWDZANIE VRAM ---
    # Pobieranie wolnego VRAM w MB bezpośrednio z narzędzia NVIDIA
    FREE_VRAM_MB=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | head -n 1)

    if [ -z "$FREE_VRAM_MB" ]; then
        echo "❌ BŁĄD: Nie udało się odczytać pamięci VRAM! Sprawdź sterowniki NVIDIA."
        exit 1
    fi

    # --- MAGICZNE MENU WYBORU MODELU ---
    echo "🧙‍♂️ Skanowanie katalogu $MODEL_PATH..."
    echo "📊 Dostępna pamięć VRAM: ${FREE_VRAM_MB} MB (Zapas na kontekst: ${OVERHEAD_MB} MB)"
    echo ""

    shopt -s nullglob
    files=("$MODEL_PATH"/*.gguf)

    if [ ${#files[@]} -eq 0 ]; then
        echo "❌ BŁĄD: Nie znaleziono żadnych plików .gguf w $MODEL_PATH!"
        exit 1
    fi

    declare -a fit_indices
    declare -a nofit_indices

    # Pętla sortująca modele
    for i in "${!files[@]}"; do
        filepath="${files[$i]}"
        
        # Obliczanie wagi pliku w megabajtach
        size_mb=$(du -m "$filepath" | awk '{print $1}')
        req_mb=$((size_mb + OVERHEAD_MB))
        
        # Przydział do odpowiedniej grupy
        if [ "$req_mb" -le "$FREE_VRAM_MB" ]; then
            fit_indices+=($i)
        else
            nofit_indices+=($i)
        fi
    done

    echo "✅ Modele mieszczące się w pamięci (VRAM):"
    echo "------------------------------------------------------"
    if [ ${#fit_indices[@]} -eq 0 ]; then
        echo "  Brak modeli mieszczących się w aktualnie wolnej pamięci VRAM."
    else
        for i in "${fit_indices[@]}"; do
            filepath="${files[$i]}"
            filename=$(basename "$filepath")
            size=$(ls -lh "$filepath" | awk '{print $5}')
            num=$((i+1))
            printf "  %2d) %-50s [%5s]\n" "$num" "$filename" "$size"
        done
    fi
    echo "------------------------------------------------------"

    if [ ${#nofit_indices[@]} -gt 0 ]; then
        echo ""
        echo "⚠️ Modele wymagające więcej VRAM (ryzyko braku pamięci):"
        echo "------------------------------------------------------"
        for i in "${nofit_indices[@]}"; do
            filepath="${files[$i]}"
            filename=$(basename "$filepath")
            size=$(ls -lh "$filepath" | awk '{print $5}')
            num=$((i+1))
            printf "  %2d) %-50s [%5s]\n" "$num" "$filename" "$size"
        done
        echo "------------------------------------------------------"
    fi

    echo ""
    read -p "👉 Wybierz numer modelu (1-${#files[@]}) lub wciśnij CTRL+C aby anulować: " choice

    if ! [[ "$choice" =~ ^[0-9]+$ ]] || [ "$choice" -lt 1 ] || [ "$choice" -gt "${#files[@]}" ]; then
        echo "❌ Nieprawidłowy wybór. Uruchom skrypt ponownie."
        exit 1
    fi

    selected_index=$((choice-1))
    export MODEL=$(basename "${files[$selected_index]}")

    # --- WŁAŚCIWY START ---
    echo ""
    echo "🧹 Zamykanie poprzednich instancji llama-server..."
    pkill -f "llama-server" > /dev/null 2>&1
    sleep 1

    echo "🚀 Uruchamiam serwer z modelem: $MODEL"

    nohup /home/doman/llama.cpp/build-cuda/bin/llama-server \
      -m "$MODEL_PATH/$MODEL" \
      -ngl 99 $CONTEXT \
      --jinja \
      --parallel 1 \
      --port 8080 \
      --host 0.0.0.0 \
      --no-warmup \
      --metrics \
      --log-file ~/server.log \
      --log-colors off \
      --flash-attn on > /dev/null 2>&1 &

    sleep 2
    if pgrep -f "llama-server" > /dev/null; then
        echo "✅ Serwer z modelem $MODEL wystartował w tle na porcie 8080."
        echo "💡 Aby sprawdzić postęp, wpisz: ./rtx.sh logs"
    else
        echo "❌ BŁĄD: Serwer nie wystartował. Sprawdź logi."
    fi
fi