#!/bin/bash

# --- KONFIGURACJA GŁÓWNA ---
MODEL_PATH="/home/models/gguf"
LLAMA_SERVER_PATH="/home/doman/llama.cpp/build-master/bin/llama-server"
CONTEXT="-c 4000"
CACHE_TYPE_K="q4_0"
CACHE_TYPE_V="q4_0"
OVERHEAD_MB=1536 
MAX_WAIT=120

# --- NOWE OPCJE (MULTI-GPU / RPC) ---
# Zostaw puste "", aby skrypt pytał o backend przy każdym uruchomieniu. 
# Możesz wpisać na stałe: "cuda", "rocm" lub "vulkan".
BACKEND=""

# Cele dla serwerów RPC. Jeśli korzystasz, wpisz adresy (np. "127.0.0.1:50052,192.168.1.10:50052").
# Jeśli zostawisz puste, flaga --rpc nie zostanie dodana.
RPC_TARGETS=""

# --- KOLORY BASH ---
COLOR_RESET='\e[0m'
COLOR_GREEN='\e[32m'
COLOR_LIGHT_RED='\e[91m'
COLOR_RED='\e[31m'

# --- FUNKCJE ZLICZANIA VRAM ---
TOTAL_VRAM_MB=0
FREE_VRAM_MB=0

get_cuda_vram() {
    if command -v nvidia-smi &> /dev/null; then
        tot=$(nvidia-smi --query-gpu=memory.total --format=csv,noheader,nounits | awk '{s+=$1} END {print s}')
        free=$(nvidia-smi --query-gpu=memory.free --format=csv,noheader,nounits | awk '{s+=$1} END {print s}')
        TOTAL_VRAM_MB=$((TOTAL_VRAM_MB + tot))
        FREE_VRAM_MB=$((FREE_VRAM_MB + free))
    fi
}

get_rocm_vram() {
    if command -v rocm-smi &> /dev/null; then
        tot_b=$(rocm-smi --showmeminfo vram | grep -i "Total Memory" | awk '{s+=$NF} END {print s}')
        used_b=$(rocm-smi --showmeminfo vram | grep -i "Total Used Memory" | awk '{s+=$NF} END {print s}')
        if [ -n "$tot_b" ] && [ -n "$used_b" ]; then
            tot_mb=$((tot_b / 1024 / 1024))
            used_mb=$((used_b / 1024 / 1024))
            TOTAL_VRAM_MB=$((TOTAL_VRAM_MB + tot_mb))
            FREE_VRAM_MB=$((FREE_VRAM_MB + (tot_mb - used_mb)))
        fi
    fi
}

# --- 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
    # --- POBIERANIE AKTYWNEGO MODELU ---
    ACTIVE_MODEL=""
    if pgrep -f "llama-server" > /dev/null; then
        SERVER_CMD=$(pgrep -a -f "llama-server" | head -n 1)
        if [[ "$SERVER_CMD" == *"-m "* ]]; then
            ACTIVE_MODEL_PATH=$(echo "$SERVER_CMD" | sed 's/.*-m \([^ ]*\).*/\1/')
            ACTIVE_MODEL=$(basename "$ACTIVE_MODEL_PATH")
        fi
    fi

    # --- WYBÓR BACKENDU ---
    SELECTED_BACKEND="$BACKEND"
    if [ -z "$SELECTED_BACKEND" ]; then
        echo "🔧 Wybierz środowisko (Backend):"
        echo "  1) CUDA (NVIDIA)"
        echo "  2) ROCm (AMD)"
        echo "  3) Vulkan (Połączone GPU / Wszystkie dostępne w systemie)"
        read -p "👉 Twój wybór (1-3): " be_choice
        case $be_choice in
            1) SELECTED_BACKEND="cuda" ;;
            2) SELECTED_BACKEND="rocm" ;;
            3) SELECTED_BACKEND="vulkan" ;;
            *) echo "❌ Nieprawidłowy wybór."; exit 1 ;;
        esac
        echo ""
    fi

    # --- SPRAWDZANIE VRAM (MULTI-GPU) ---
    if [ "$SELECTED_BACKEND" == "cuda" ]; then
        get_cuda_vram
    elif [ "$SELECTED_BACKEND" == "rocm" ]; then
        get_rocm_vram
    elif [ "$SELECTED_BACKEND" == "vulkan" ]; then
        get_cuda_vram
        get_rocm_vram
    fi

    if [ "$TOTAL_VRAM_MB" -eq 0 ]; then
        echo "❌ BŁĄD: Nie udało się odczytać pamięci VRAM! Sprawdź sterowniki dla wybranego backendu ($SELECTED_BACKEND)."
        exit 1
    fi

    echo "🧙‍♂️ Skanowanie katalogu $MODEL_PATH..."
    echo "📊 Całkowita zsumowana pamięć VRAM: ${TOTAL_VRAM_MB} MB"
    echo "📊 Dostępna pamięć VRAM: ${FREE_VRAM_MB} MB (Zapas na kontekst: ${OVERHEAD_MB} MB)"
    if [ -n "$RPC_TARGETS" ]; then
        echo "🌐 Wykryto aktywne cele RPC: $RPC_TARGETS"
    fi
    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_free_indices
    declare -a fit_total_indices
    declare -a nofit_indices

    for i in "${!files[@]}"; do
        filepath="${files[$i]}"
        size_mb=$(du -m "$filepath" | awk '{print $1}')
        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
        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

        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

    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 w systemie)${COLOR_RESET}\n" "$num" "$filename" "$size"
            fi
        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]}")

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

    rm -f ~/server.log

    RPC_FLAG=""
    if [ -n "$RPC_TARGETS" ]; then
        RPC_FLAG="--rpc $RPC_TARGETS"
    fi

    echo "🚀 Uruchamiam serwer z modelem: $MODEL"

    nohup "$LLAMA_SERVER_PATH" \
      -m "$MODEL_PATH/$MODEL" \
      -ngl 99 $CONTEXT \
      --cache-type-k "$CACHE_TYPE_K" \
      --cache-type-v "$CACHE_TYPE_V" \
      $RPC_FLAG \
      --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 &

    echo "⏳ Oczekiwanie na załadowanie modelu i alokację bufora kontekstu w VRAM..."

    elapsed=0
    loaded=false

    while [ $elapsed -lt $MAX_WAIT ]; do
        if ! pgrep -f "llama-server" > /dev/null; then
            echo "❌ BŁĄD: Proces serwera został nagle zamknięty!"
            echo "   (Najprawdopodobniej model wraz z zadeklarowanym kontekstem nie zmieścił się w pamięci VRAM)"
            echo "💡 Uruchom './llama-run.py logs' aby sprawdzić dokładną przyczynę błędu."
            exit 1
        fi

        # NIEZAWODNY TEST: Uderzamy bezpośrednio w interfejs po HTTP. Zero opóźnień logowania.
        if command -v curl >/dev/null 2>&1; then
            if curl -s -f http://127.0.0.1:8080/ > /dev/null; then
                loaded=true
                break
            fi
        else
            # Zapasowo, jakby curl nie był zainstalowany
            if grep -E -i -q "HTTP server listening|llama server listening" ~/server.log 2>/dev/null; then
                loaded=true
                break
            fi
        fi

        sleep 2
        elapsed=$((elapsed + 2))
    done

    echo ""
    if [ "$loaded" = true ]; then
        echo "> Sprawdzanie wielkości kontekstu..."
        sleep 1
        echo "✅ Uruchomiono serwer z modelem $MODEL"
        echo "✅ Serwer wystartował w tle na porcie 8080. Poprawnie ustawiono wielkosc context tokenow."
    else
        echo "⚠️ Serwer działa, ale nie potwierdził pełnego załadowania w ciągu $MAX_WAIT sekund."
        echo "💡 Wpisz './llama-run.py logs' aby sprawdzić, czy proces ładowania nadal trwa."
    fi
fi