#!/bin/bash

# --- KONFIGURACJA GŁÓWNA ---
	# Ścieżka do katalogu, w którym trzymasz pobrane modele w formacie .gguf
MODEL_PATH="/home/models/gguf"

	# Ścieżka do pliku wykonywalnego serwera llama.cpp
LLAMA_SERVER_PATH="/home/doman/llama.cpp/build-master/bin/llama-server"

	# Rozmiar okna kontekstu (w tokenach) przypisywany do uruchamianego modelu
CONTEXT="-c 4000"

	# Typ kwantyzacji dla kluczy (K) w buforze kontekstu (KV cache)
	# Możliwe wartości: f32, f16 (domyślna), q8_0, q5_1, q5_0, q4_1, q4_0, iq4_nl
CACHE_TYPE_K="q4_0"

	# Typ kwantyzacji dla wartości (V) w buforze kontekstu (KV cache)
	# Możliwe wartości: f32, f16 (domyślna), q8_0, q5_1, q5_0, q4_1, q4_0, iq4_nl
CACHE_TYPE_V="q4_0"

	# Bezpieczny zapas VRAM (w MB) na bufor kontekstu (KV cache) i podstawowe procesy systemowe
OVERHEAD_MB=1536

# --- KOLORY BASH ---
# Zmienne używane do formatowania wyświetlanego tekstu w terminalu
COLOR_RESET='\e[0m'       # Resetuje kolor do domyślnego
COLOR_GREEN='\e[32m'      # Oznacza aktualnie uruchomiony model
COLOR_LIGHT_RED='\e[91m'  # Oznacza modele, które zmieszczą się po zwolnieniu VRAM
COLOR_RED='\e[31m'        # Oznacza modele, które są zbyt duże na tę kartę

# --- LOGIKA STEROWANIA ---

if [ "$1" == "stop" ]; then
    # Zatrzymuje działający proces llama-server i zwalnia pamięć karty
    echo "🛑 Zatrzymywanie serwera llama i zwalnianie VRAM (MI50)..."
    pkill -f "llama-server" > /dev/null 2>&1
    echo "✅ System czysty. VRAM został zwolniony."
    exit 0

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

else
    # --- POBIERANIE AKTYWNEGO MODELU ---
    ACTIVE_MODEL=""
    if pgrep -f "llama-server" > /dev/null; then
        # Pobieramy pełną komendę działającego procesu
        SERVER_CMD=$(pgrep -a -f "llama-server" | head -n 1)
        # Wyciągamy ścieżkę do modelu po fladze -m i bierzemy samą nazwę pliku
        if [[ "$SERVER_CMD" == *"-m "* ]]; then
            ACTIVE_MODEL_PATH=$(echo "$SERVER_CMD" | sed 's/.*-m \([^ ]*\).*/\1/')
            ACTIVE_MODEL=$(basename "$ACTIVE_MODEL_PATH")
        fi
    fi

    # --- SPRAWDZANIE VRAM (ROCm) ---
    # Odczyt całkowitej i zajętej pamięci VRAM bezpośrednio z rocm-smi w Bajtach
    TOTAL_VRAM_B=$(rocm-smi --showmeminfo vram | grep -i "Total Memory" | head -n 1 | awk '{print $NF}')
    USED_VRAM_B=$(rocm-smi --showmeminfo vram | grep -i "Total Used Memory" | head -n 1 | awk '{print $NF}')

    if [ -z "$TOTAL_VRAM_B" ] || [ -z "$USED_VRAM_B" ]; then
        echo "❌ BŁĄD: Nie udało się odczytać pamięci VRAM z rocm-smi! Sprawdź instalację ROCm."
        exit 1
    fi

    # Przeliczanie Bajtów na Megabajty dla czytelności i łatwiejszych obliczeń
    TOTAL_VRAM_MB=$((TOTAL_VRAM_B / 1024 / 1024))
    USED_VRAM_MB=$((USED_VRAM_B / 1024 / 1024))
    FREE_VRAM_MB=$((TOTAL_VRAM_MB - USED_VRAM_MB))

    echo "🧙‍♂️ Skanowanie katalogu $MODEL_PATH..."
    echo "📊 Całkowita pamięć VRAM: ${TOTAL_VRAM_MB} MB"
    echo "📊 Dostępna pamięć VRAM: ${FREE_VRAM_MB} MB (Zapas na kontekst: ${OVERHEAD_MB} MB)"
    echo ""

    # Zabezpieczenie przed brakiem plików w katalogu
    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

    # Puste tablice do grupowania indeksów modeli według ich wymagań pamięciowych
    declare -a fit_free_indices
    declare -a fit_total_indices
    declare -a nofit_indices

    # Sortowanie modeli do odpowiednich grup
    for i in "${!files[@]}"; do
        filepath="${files[$i]}"
        # Wyciąganie rozmiaru pliku .gguf w MB
        size_mb=$(du -m "$filepath" | awk '{print $1}')
        # Dodawanie obowiązkowego zapasu VRAM do rozmiaru modelu
        req_mb=$((size_mb + OVERHEAD_MB))
        
        if [ "$req_mb" -le "$FREE_VRAM_MB" ]; then
            fit_free_indices+=($i)
        elif [ "$req_mb" -le "$TOTAL_VRAM_MB" ]; then
            fit_total_indices+=($i)
        else
            nofit_indices+=($i)
        fi
    done

    # --- WYSWIETLANIE LISTY ---
    echo "✅ Modele mieszczące się w aktualnie wolnej pamięci (VRAM):"
    echo "------------------------------------------------------"
    if [ ${#fit_free_indices[@]} -eq 0 ] && [ ${#fit_total_indices[@]} -eq 0 ]; then
        echo "  Brak modeli w tej kategorii."
    else
        # Wyświetlanie modeli, które wchodzą w obecny wolny VRAM (zwykły tekst)
        for i in "${fit_free_indices[@]}"; do
            filepath="${files[$i]}"
            filename=$(basename "$filepath")
            size=$(ls -lh "$filepath" | awk '{print $5}')
            num=$((i+1))
            
            if [ "$filename" == "$ACTIVE_MODEL" ]; then
                printf "${COLOR_GREEN}  %2d) %-50s [%5s] (Aktualnie uruchomiony)${COLOR_RESET}\n" "$num" "$filename" "$size"
            else
                printf "  %2d) %-50s [%5s]\n" "$num" "$filename" "$size"
            fi
        done

        # Wyświetlanie modeli, które wejdą na kartę, ale trzeba użyć opcji stop (lekko czerwone)
        for i in "${fit_total_indices[@]}"; do
            filepath="${files[$i]}"
            filename=$(basename "$filepath")
            size=$(ls -lh "$filepath" | awk '{print $5}')
            num=$((i+1))

            if [ "$filename" == "$ACTIVE_MODEL" ]; then
                printf "${COLOR_GREEN}  %2d) %-50s [%5s] (Aktualnie uruchomiony)${COLOR_RESET}\n" "$num" "$filename" "$size"
            else
                printf "${COLOR_LIGHT_RED}  %2d) %-50s [%5s] (Wymaga zwolnienia VRAM)${COLOR_RESET}\n" "$num" "$filename" "$size"
            fi
        done
    fi

    # Wyświetlanie modeli zbyt dużych na fizyczną pojemność karty (mocno czerwone)
    if [ ${#nofit_indices[@]} -gt 0 ]; then
        echo "------------------------------------------------------"
        echo "⚠️ Pozostałe modele (przekraczają całkowity VRAM):"
        echo "------------------------------------------------------"
        for i in "${nofit_indices[@]}"; do
            filepath="${files[$i]}"
            filename=$(basename "$filepath")
            size=$(ls -lh "$filepath" | awk '{print $5}')
            num=$((i+1))
            
            if [ "$filename" == "$ACTIVE_MODEL" ]; then
                printf "${COLOR_GREEN}  %2d) %-50s [%5s] (Aktualnie uruchomiony)${COLOR_RESET}\n" "$num" "$filename" "$size"
            else
                printf "${COLOR_RED}  %2d) %-50s [%5s] (Brak pamięci na karcie)${COLOR_RESET}\n" "$num" "$filename" "$size"
            fi
        done
        echo "------------------------------------------------------"
    fi

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

    # Walidacja, czy wpisano poprawny numer z dostępnego zakresu
    if ! [[ "$choice" =~ ^[0-9]+$ ]] || [ "$choice" -lt 1 ] || [ "$choice" -gt "${#files[@]}" ]; then
        echo "❌ Nieprawidłowy wybór. Uruchom skrypt ponownie."
        exit 1
    fi

    # Przypisanie wybranego pliku .gguf
    selected_index=$((choice-1))
    export MODEL=$(basename "${files[$selected_index]}")

    # --- URUCHAMIANIE SERWERA ---
    echo ""
    echo "🧹 Zamykanie poprzednich instancji llama-server..."
    pkill -f "llama-server" > /dev/null 2>&1
    sleep 1

    echo "🚀 Uruchamiam serwer z modelem: $MODEL"

    # Start serwera w tle z dodanymi flagami kwantyzacji KV cache
    nohup "$LLAMA_SERVER_PATH" \
      -m "$MODEL_PATH/$MODEL" \
      -ngl 99 $CONTEXT \
      --cache-type-k "$CACHE_TYPE_K" \
      --cache-type-v "$CACHE_TYPE_V" \
      --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 &

    # Szybka weryfikacja, czy proces faktycznie wystartował
    sleep 2
    if pgrep -f "llama-server" > /dev/null; then
        echo "✅ Serwer z modelem $MODEL wystartował w tle na porcie 8080."
    else
        echo "❌ BŁĄD: Serwer nie wystartował. Sprawdź logi."
    fi
fi