#!/bin/bash

# --- MAIN CONFIGURATION ---
MODEL_PATH="/home/models/gguf"
LLAMA_SERVER_PATH="/home/doman/llama-serwer/llama-b8651/build/bin/llama-server"
CONTEXT="-c 131072"
CACHE_TYPE_K="q8_0"
CACHE_TYPE_V="q8_0"
OVERHEAD_MB=1536 
FIT="-fit off"
MAX_WAIT=60
PARALLEL=1
NGL="999"
MAIN_GPU="-mg 1"

# --- ENVIRONMENT VARIABLES ---
# Używamy $(dirname) na podstawie powyższej ścieżki serwera
LIB_PATH=$(dirname "$LLAMA_SERVER_PATH")
ENV_VARS="LD_LIBRARY_PATH=$LIB_PATH HIP_VISIBLE_DEVICES=1,0"

# --- PORT SEARCH CONFIGURATION ---
PORT_SEARCH_START=8080
PORT_SEARCH_END=8085

# --- NEW OPTIONS (MULTI-GPU / RPC) ---
BACKEND=""
RPC_TARGETS=""

# --- BASH COLORS ---
COLOR_RESET='\e[0m'
COLOR_GREEN='\e[32m'
COLOR_LIGHT_RED='\e[91m'
COLOR_RED='\e[31m'
COLOR_YELLOW='\e[33m'
COLOR_CYAN='\e[36m'

# --- VRAM COUNTING FUNCTIONS ---
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
}

# --- MMPROJ (VISION) FINDING FUNCTION ---
find_mmproj() {
    local filename="$1"
    local base_name=$(echo "$filename" | sed -E 's/[-.](Q[0-8]_[0-9K_A-Z]*|IQ[0-8]_[0-9A-Z_]*|bf16|f16)(\.gguf)?$//i; s/\.gguf$//i')
    local found=""
    
    found=$(ls "$MODEL_PATH/mmproj/"*"$base_name"* 2>/dev/null | grep -i "\.gguf$" | head -n 1)
    
    if [ -z "$found" ]; then
        local short_base=$(echo "$base_name" | cut -d'-' -f1-2)
        if [ -n "$short_base" ]; then
            found=$(ls "$MODEL_PATH/mmproj/"*"$short_base"* 2>/dev/null | grep -i "\.gguf$" | head -n 1)
        fi
    fi
    
    if [ -z "$found" ]; then
        local very_short=$(echo "$base_name" | cut -d'-' -f1)
        if [ -n "$very_short" ]; then
            found=$(ls "$MODEL_PATH/mmproj/"*"$very_short"* 2>/dev/null | grep -i "\.gguf$" | grep -i "mmproj" | head -n 1)
        fi
    fi
    
    echo "$found"
}

# --- COMMAND HANDLING ---

if [ "$1" == "help" ]; then
    echo -e "${COLOR_CYAN}--- LLAMA SERVER SCRIPT HELP ---${COLOR_RESET}"
    echo -e "${COLOR_YELLOW}Dostępne komendy:${COLOR_RESET}"
    echo "  ./llama-run.py               : Startuje interaktywny wybór modelu."
    echo "  ./llama-run.py stop          : Zatrzymuje serwer, jeśli działa tylko jedna instancja."
    echo "  ./llama-run.py stopall       : Zatrzymuje wszystkie działające instancje serwera."
    echo "  ./llama-run.py stop<port>    : Zatrzymuje konkretny serwer."
    echo "  ./llama-run.py logs          : Pokazuje logi domyślnego serwera."
    echo "  ./llama-run.py logs<port>    : Pokazuje logi konkretnego serwera."
    echo "  ./llama-run.py help          : Pokazuje tę pomoc."
    exit 0

elif [ "$1" == "stopall" ]; then
    echo -e "🛑 ${COLOR_YELLOW}Zatrzymywanie wszystkich instancji llama-server i zwalnianie pamięci VRAM...${COLOR_RESET}"
    pkill -f "llama-server" > /dev/null 2>&1
    echo "✅ System oczyszczony. Cała pamięć VRAM została zwolniona."
    exit 0

elif [[ "$1" =~ ^stop([0-9]+)$ ]]; then
    target_port="${BASH_REMATCH[1]}"
    echo -e "🛑 ${COLOR_YELLOW}Zatrzymywanie llama-server na porcie $target_port...${COLOR_RESET}"
    pkill -f "llama-server.*--port $target_port" > /dev/null 2>&1
    echo "✅ Serwer na porcie $target_port został zatrzymany."
    exit 0

elif [ "$1" == "stop" ]; then
    active_servers=$(pgrep -a -f "llama-server.*--port" | grep -v "$0" | wc -l)
    
    if [ "$active_servers" -eq 1 ]; then
        echo -e "🛑 ${COLOR_YELLOW}Zatrzymywanie jedynej aktywnej instancji llama-server...${COLOR_RESET}"
        pkill -f "llama-server" > /dev/null 2>&1
        echo "✅ System oczyszczony. Pamięć VRAM została zwolniona."
    elif [ "$active_servers" -eq 0 ]; then
        echo "ℹ️ Nie znaleziono aktywnych instancji llama-server."
    else
        echo -e "⚠️ ${COLOR_RED}Wykryto wiele instancji serwera ($active_servers).${COLOR_RESET}"
        echo "💡 Użyj './llama-run.py stopall' aby zamknąć wszystkie."
        pgrep -a -f "llama-server.*--port" | sed -n 's/.*--port \([0-9]*\).*-m \([^ ]*\).*/  Port: \1 | Model: \2/p'
    fi
    exit 0

elif [[ "$1" =~ ^logs([0-9]*)$ ]]; then
    target_port="${BASH_REMATCH[1]}"
    if [ -z "$target_port" ] || [ "$target_port" == "$PORT_SEARCH_START" ]; then
        log_target="$HOME/server.log"
    else
        log_target="$HOME/server${target_port}.log"
    fi
    
    if [ -f "$log_target" ]; then
        echo "📋 Podgląd $log_target..."
        tail -f "$log_target"
    else
        echo "❌ Nie znaleziono pliku logów: $log_target"
    fi
    exit 0

else
    # --- MAIN SCRIPT LOGIC ---
    
    # --- PORT CHECK & SELECTION ---
    echo "🔍 Sprawdzanie portów ($PORT_SEARCH_START-$PORT_SEARCH_END)..."
    FOUND_PORTS=()
    for p in $(seq $PORT_SEARCH_START $PORT_SEARCH_END); do
        if pgrep -a -f "llama-server.*--port $p" > /dev/null; then
            FOUND_PORTS+=($p)
        fi
    done

    DEFAULT_PORT=$PORT_SEARCH_START
    if [ ${#FOUND_PORTS[@]} -gt 0 ]; then
        ports_str=$(IFS=,; echo "${FOUND_PORTS[*]}")
        echo -e "🌐 ${COLOR_GREEN}Znaleziono serwer na porcie: $ports_str${COLOR_RESET}"
        
        for i in $(seq $PORT_SEARCH_START $PORT_SEARCH_END); do
            if ! [[ " ${FOUND_PORTS[*]} " =~ " ${i} " ]]; then
                DEFAULT_PORT=$i
                break
            fi
        done
    else
        echo "🌐 Brak aktywnych instancji w zdefiniowanym zakresie."
    fi

    echo ""
    read -p "👉 Wpisz port aby zaktualizować serwer lub uruchomić nowy [$DEFAULT_PORT]: " INPUT_PORT
    CHOSEN_PORT=${INPUT_PORT:-$DEFAULT_PORT}

    SERVER_LOG_FILE="$HOME/server.log"

    # --- GET ACTIVE MODEL FOR CHOSEN PORT ---
    ACTIVE_MODEL=""
    if pgrep -a -f "llama-server.*--port $CHOSEN_PORT" > /dev/null; then
        SERVER_CMD=$(pgrep -a -f "llama-server.*--port $CHOSEN_PORT" | head -n 1)
        ACTIVE_MODEL_PATH=$(echo "$SERVER_CMD" | sed -n 's/.*-m \([^ ]*\).*/\1/p')
        if [ -n "$ACTIVE_MODEL_PATH" ]; then
            ACTIVE_MODEL=$(basename "$ACTIVE_MODEL_PATH")
        fi
    fi

    # --- BACKEND SELECTION ---
    SELECTED_BACKEND="$BACKEND"
    if [ -z "$SELECTED_BACKEND" ]; then
        echo "🔧 Wybierz środowisko:"
        echo "  1) CUDA"
        echo "  2) ROCm"
        echo "  3) Vulkan"
        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

    # --- VRAM CHECK (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ć VRAM. Sprawdź sterowniki."
        exit 1
    fi

    echo "🧙‍♂️ Skanowanie katalogu $MODEL_PATH..."
    echo "📊 Całkowita pamięć VRAM: ${TOTAL_VRAM_MB} MB"
    echo "📊 Dostępna pamięć VRAM: ${FREE_VRAM_MB} MB (Bufor kontekstu: ${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 plików .gguf w $MODEL_PATH!"
        exit 1
    fi

    declare -a fit_free_indices
    declare -a fit_total_indices
    declare -a nofit_indices
    declare -a MODEL_VL_PATHS
    declare -a MODEL_VL_TEXTS

    for i in "${!files[@]}"; do
        filepath="${files[$i]}"
        filename=$(basename "$filepath")
        size_mb=$(du -m "$filepath" | awk '{print $1}')
        
        vl_path=$(find_mmproj "$filename")
        vl_size_mb=0
        vl_text=""
        
        if [ -n "$vl_path" ]; then
            vl_size_mb=$(du -m "$vl_path" | awk '{print $1}')
            vl_size_gb=$(awk "BEGIN {printf \"%.1f\", $vl_size_mb/1024}")
            vl_text="[wizja ✔ ${vl_size_gb}GB]"
            MODEL_VL_PATHS[$i]="$vl_path"
        else
            MODEL_VL_PATHS[$i]=""
        fi
        
        MODEL_VL_TEXTS[$i]="$vl_text"
        
        req_mb=$((size_mb + OVERHEAD_MB + vl_size_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

    # --- DISPLAY LIST ---
    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))
            vl_str="${COLOR_YELLOW}${MODEL_VL_TEXTS[$i]}${COLOR_RESET}"
            
            if [ "$filename" == "$ACTIVE_MODEL" ]; then
                printf "${COLOR_GREEN}  %2d) %-50s [%5s] %b (Uruchomiony na $CHOSEN_PORT)${COLOR_RESET}\n" "$num" "$filename" "$size" "$vl_str"
            else
                printf "  %2d) %-50s [%5s] %b\n" "$num" "$filename" "$size" "$vl_str"
            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))
            vl_str="${COLOR_YELLOW}${MODEL_VL_TEXTS[$i]}${COLOR_RESET}"

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

    if [ ${#nofit_indices[@]} -gt 0 ]; then
        echo "---------------------------------------------------------------------------------"
        echo "⚠️ Pozostałe modele (przekraczające całkowitą pamięć VRAM):"
        echo "---------------------------------------------------------------------------------"
        for i in "${nofit_indices[@]}"; do
            filepath="${files[$i]}"
            filename=$(basename "$filepath")
            size=$(ls -lh "$filepath" | awk '{print $5}')
            num=$((i+1))
            vl_str="${COLOR_YELLOW}${MODEL_VL_TEXTS[$i]}${COLOR_RESET}"
            
            if [ "$filename" == "$ACTIVE_MODEL" ]; then
                printf "${COLOR_GREEN}  %2d) %-50s [%5s] %b (Uruchomiony na $CHOSEN_PORT)${COLOR_RESET}\n" "$num" "$filename" "$size" "$vl_str"
            else
                printf "${COLOR_RED}  %2d) %-50s [%5s] %b (Niewystarczająca ilość pamięci)${COLOR_RESET}\n" "$num" "$filename" "$size" "$vl_str"
            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. Zrestartuj skrypt."
        exit 1
    fi

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

    # --- VISION MODULE SELECTION ---
    MMPROJ=""
    if [ -n "$selected_vl_path" ]; then
        vl_name=$(basename "$selected_vl_path")
        echo ""
        read -p "👉 Wykryto bibliotekę wizyjną ($vl_name). Załadować moduł wizyjny? [Y/n]: " load_vl
        if [[ -z "$load_vl" || "$load_vl" =~ ^[Yy]$ ]]; then
            MMPROJ="--mmproj $selected_vl_path"
            echo -e "✅ ${COLOR_GREEN}Moduł wizyjny zostanie załadowany.${COLOR_RESET}"
        else
            echo "❌ Moduł wizyjny zostanie pominięty."
        fi
    fi

    # --- STARTING SERVER ---
    echo ""
    echo "🧹 Zamykanie poprzedniej instancji serwera na porcie $CHOSEN_PORT..."
    pkill -f "llama-server.*--port $CHOSEN_PORT" > /dev/null 2>&1
    sleep 1

    rm -f "$SERVER_LOG_FILE"

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

    echo -e "🛠️ ${COLOR_CYAN}Wygenerowane polecenie startowe:${COLOR_RESET}"
    echo -e "${COLOR_YELLOW}env $ENV_VARS $LLAMA_SERVER_PATH -m $MODEL_PATH/$MODEL -ngl $NGL $CONTEXT $MMPROJ $FIT $MAIN_GPU --cache-type-k $CACHE_TYPE_K --cache-type-v $CACHE_TYPE_V $RPC_FLAG --jinja --parallel $PARALLEL --port $CHOSEN_PORT --host 0.0.0.0 --no-warmup --metrics --log-file ~/server.log --log-colors off --flash-attn on${COLOR_RESET}"
    echo ""

    echo "🚀 Uruchamianie serwera z modelem: $MODEL na porcie $CHOSEN_PORT"

    # Wykorzystujemy 'env' aby jednorazowo przekazać zmienne środowiskowe do komendy nohup
    nohup env $ENV_VARS "$LLAMA_SERVER_PATH" \
      -m "$MODEL_PATH/$MODEL" \
      -ngl "$NGL" \
      $CONTEXT \
      $MMPROJ \
      $FIT \
      $MAIN_GPU \
      --cache-type-k "$CACHE_TYPE_K" \
      --cache-type-v "$CACHE_TYPE_V" \
      $RPC_FLAG \
      --jinja \
      --parallel "$PARALLEL" \
      --port "$CHOSEN_PORT" \
      --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ę buforów VRAM..."

    elapsed=0
    loaded=false

    while [ $elapsed -lt $MAX_WAIT ]; do
        if ! pgrep -f "llama-server.*--port $CHOSEN_PORT" > /dev/null; then
            echo "❌ BŁĄD: Proces serwera został nieoczekiwanie zamknięty!"
            echo "   (Najprawdopodobniej model i kontekst nie zmieściły się w VRAM lub wystąpił błąd ładowania)"
            echo "💡 Uruchom './llama-run.py logs${CHOSEN_PORT}' aby sprawdzić dokładną przyczynę."
            exit 1
        fi

        if command -v curl >/dev/null 2>&1; then
            if curl -s -f http://127.0.0.1:$CHOSEN_PORT/ > /dev/null; then
                loaded=true
                break
            fi
        else
            if grep -E -i -q "HTTP server listening|llama server listening" "$SERVER_LOG_FILE" 2>/dev/null; then
                loaded=true
                break
            fi
        fi

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

    echo ""
    if [ "$loaded" = true ]; then
        echo "> Weryfikacja rozmiaru kontekstu..."
        sleep 1
        echo "✅ Serwer wystartował z modelem $MODEL"
        echo "✅ Serwer działa w tle na porcie $CHOSEN_PORT. Tokeny kontekstu ustawione prawidłowo."
        echo "📄 Logi zapisano do: ~/server.log"
    else
        echo "⚠️ Serwer działa, ale nie potwierdził pełnego załadowania w ciągu $MAX_WAIT sekund."
        echo "💡 Wpisz './llama-run.py logs${CHOSEN_PORT}' aby zobaczyć czy proces ładowania nadal trwa."
    fi
fi