#!/bin/bash

# --- KONFIGURACJA ŚCIEŻEK ---
MODEL_PATH="/home/models/gguf"
MASTER_BIN="$HOME/llama-master/build-cuda-master/bin/llama-server"
WORKER_BIN="$HOME/llama-master/build-rpc-mi50/bin/rpc-server"
LOG_WORKER="$HOME/start_llama/worker_host.log"
# Wymuszona ścieżka logowania Mastera
LOG_MASTER="$HOME/server.log"

# --- ZMIENNE DLA MI50 ---
export HSA_OVERRIDE_GFX_VERSION=9.0.6
export HSA_ENABLE_SDMA=0

# Przejście do katalogu roboczego
cd ~/start_llama

# --- LOGIKA STEROWANIA ---

if [ "$1" == "stop" ]; then
    echo "🛑 Zatrzymywanie procesów hosta..."
    pkill -9 -f llama-server
    pkill -9 -f rpc-server
    echo "✅ Procesy ubite. VRAM na RTX 4070 i MI50 został zwolniony."
    exit 0

elif [ "$1" == "logs" ]; then
    echo "📋 Podgląd logów Mastera (RTX 4070)..."
    tail -f "$LOG_MASTER"
    exit 0

else
    # --- MENU WYBORU MODELU ---
    echo "🧙‍♂️ Skanowanie hosta: $MODEL_PATH..."
    shopt -s nullglob
    files=("${MODEL_PATH}"/*.gguf)

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

    echo ""
    echo "📚 Modele na dysku (HOST MODE):"
    echo "------------------------------------------------------"
    for i in "${!files[@]}"; do
        size=$(ls -lh "${files[$i]}" | awk '{print $5}')
        printf "  %2d) %-50s [%5s]\n" "$((i+1))" "$(basename "${files[$i]}")" "$size"
    done
    echo "------------------------------------------------------"

    read -p "👉 Wybierz numer (1-${#files[@]}): " choice

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

    SELECTED_MODEL="${files[$((choice-1))]}"
    MODEL_NAME=$(basename "$SELECTED_MODEL")

    # --- CZYSZCZENIE PORTÓW ---
    echo "🧹 Czyszczenie portów 8080 i 50052..."
    pkill -9 -f llama-server 2>/dev/null
    pkill -9 -f rpc-server 2>/dev/null
    sleep 1

    # --- START WORKERA (MI50) ---
    echo "🚀 Start WORKERA (MI50) na porcie 50052..."
    nohup "$WORKER_BIN" -p 50052 -H 0.0.0.0 > "$LOG_WORKER" 2>&1 &

    # Czekamy chwilę na inicjalizację RPC
    sleep 2

    # --- START MASTERA (RTX 4070) ---
    echo "🚀 Start MASTERA z modelem: $MODEL_NAME"
    # Ustawienia zoptymalizowane pod Twój sprzęt + wymuszone flagi
    nohup "$MASTER_BIN" \
        -m "$SELECTED_MODEL" \
        -c 8096 \
        --port 8080 \
        --host 0.0.0.0 \
        -ngl 99 \
        --rpc 127.0.0.1:50052 \
        --flash-attn on \
        --jinja \
        --metrics \
        --log-file "$LOG_MASTER" \
        --log-colors off \
        > /dev/null 2>&1 &

    # Weryfikacja startu
    sleep 3
    if pgrep -f "llama-server" > /dev/null; then
        echo "✅ Klaster HOST-RPC działa! Port: 8080."
        echo "💡 Logi: ./hostrpc.sh logs"
    else
        echo "❌ BŁĄD: Master nie wystartował. Sprawdź $LOG_MASTER"
    fi
fi