#!/bin/bash
# Gestion du homelab Funk (storage-01 + gpu-01)
# Usage: funk-cluster [start|stop|status]

GPU01_IP="192.168.10.20"
GPU01_MAC="a8:a1:59:72:a6:11"  # WOL — adresse MAC de gpu-01 (NIC principal)
GPU01_SSH="ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no gpu-01"

GREEN='\033[0;32m'
RED='\033[0;31m'
YELLOW='\033[1;33m'
BOLD='\033[1m'
NC='\033[0m'

ok()   { echo -e "  ${GREEN}✓${NC} $*"; }
fail() { echo -e "  ${RED}✗${NC} $*"; }
warn() { echo -e "  ${YELLOW}!${NC} $*"; }
info() { echo -e "  → $*"; }

CHECKS_OK=0
CHECKS_FAIL=0

check_ok()   { ok "$*";   CHECKS_OK=$((CHECKS_OK + 1)); }
check_fail() { fail "$*"; CHECKS_FAIL=$((CHECKS_FAIL + 1)); }

# --- Wake-on-LAN (python3, pas de dépendance externe) ---
wol() {
    local mac="${1//:}"
    python3 -c "
import socket
mac = '$mac'
payload = bytes.fromhex('ff' * 6 + mac * 16)
s = socket.socket(socket.AF_INET, socket.SOCK_DGRAM)
s.setsockopt(socket.SOL_SOCKET, socket.SO_BROADCAST, 1)
s.sendto(payload, ('192.168.10.255', 9))
s.close()
"
    echo -e "  → Paquet WOL envoyé vers $1"
}

# --- Vérification service systemd ---
check_service() {
    local svc=$1
    if systemctl is-active --quiet "$svc" 2>/dev/null; then
        check_ok "$svc"
        return 0
    else
        check_fail "$svc (inactif)"
        return 1
    fi
}

# --- Vérification PostgreSQL (connexion réelle via postgres) ---
check_postgresql() {
    if sudo -u postgres psql -c "SELECT 1" hermes &>/dev/null 2>&1; then
        check_ok "PostgreSQL — base hermes accessible"
    else
        check_fail "PostgreSQL — base hermes inaccessible"
    fi
}

# --- Vérification Qdrant ---
check_qdrant() {
    local result
    result=$(curl -sf --max-time 5 "http://127.0.0.1:6333/collections" \
        | python3 -c "
import sys, json
d = json.load(sys.stdin)
cols = d.get('result', {}).get('collections', [])
print(len(cols))
" 2>/dev/null)
    if [ -n "$result" ]; then
        check_ok "Qdrant — $result collection(s)"
    else
        check_fail "Qdrant non joignable (port 6333)"
    fi
}

# --- Vérification LiteLLM ---
check_litellm() {
    local models
    models=$(curl -sf --max-time 5 \
        -H "Authorization: Bearer lm-studio" \
        "http://127.0.0.1:4000/v1/models" \
        | python3 -c "
import sys, json
d = json.load(sys.stdin).get('data', [])
print(', '.join(m['id'] for m in d))
" 2>/dev/null)
    if [ -n "$models" ]; then
        check_ok "LiteLLM — modèles : $models"
    else
        check_fail "LiteLLM non joignable (port 4000)"
    fi
}

# --- Vérification Hermes HTTP ---
check_hermes() {
    if curl -sf --max-time 5 "http://127.0.0.1:8080/health" &>/dev/null; then
        check_ok "Hermes HTTP (port 8080)"
    elif systemctl is-active --quiet hermes-agent 2>/dev/null; then
        check_ok "Hermes service actif (pas d'endpoint HTTP)"
    else
        check_fail "Hermes inactif"
    fi
}

# --- Vérification NFS export ---
check_nfs_export() {
    if showmount -e localhost 2>/dev/null | grep -q "/srv/data/models"; then
        check_ok "NFS export /srv/data/models OK"
    else
        check_fail "NFS export /srv/data/models manquant"
    fi
}

# --- Vérification llama-server sur gpu-01 ---
check_llama_server() {
    local model
    model=$(curl -sf --max-time 5 "http://${GPU01_IP}:1234/v1/models" \
        | python3 -c "import sys,json; d=json.load(sys.stdin); print(d['data'][0]['id'])" 2>/dev/null)
    if [ -n "$model" ]; then
        check_ok "llama-server — modèle : $model"
        return 0
    else
        check_fail "llama-server (non joignable)"
        return 1
    fi
}

# --- Vérification embeddings ---
check_embeddings() {
    local dim
    dim=$(curl -sf --max-time 15 -X POST "http://${GPU01_IP}:1234/v1/embeddings" \
        -H "Content-Type: application/json" \
        -d '{"model":"qwen3-8b","input":["test"]}' \
        | python3 -c "import sys,json; print(len(json.load(sys.stdin)['data'][0]['embedding']))" 2>/dev/null)
    if [ -n "$dim" ]; then
        check_ok "Embeddings OK — dimension : $dim"
    else
        check_fail "Embeddings non disponibles"
    fi
}

# --- Attendre qu'un hôte réponde au ping ---
wait_ping() {
    local host=$1 timeout=${2:-120} elapsed=0
    while ! ping -c1 -W1 "$host" &>/dev/null; do
        sleep 3; elapsed=$((elapsed + 3))
        [ $elapsed -ge $timeout ] && return 1
        echo -ne "  → Attente ping $host... ${elapsed}s\r"
    done
    echo ""
    return 0
}

# --- Attendre llama-server HTTP ---
wait_llama() {
    local timeout=${1:-240} elapsed=0
    while ! curl -sf --max-time 3 "http://${GPU01_IP}:1234/v1/models" &>/dev/null; do
        sleep 5; elapsed=$((elapsed + 5))
        [ $elapsed -ge $timeout ] && return 1
        echo -ne "  → Attente llama-server... ${elapsed}s\r"
    done
    echo ""
    return 0
}

# --- Afficher le score final ---
print_score() {
    local total=$((CHECKS_OK + CHECKS_FAIL))
    echo ""
    if [ $CHECKS_FAIL -eq 0 ]; then
        echo -e "  ${GREEN}${BOLD}[$CHECKS_OK/$total] Tous les checks OK${NC}"
    else
        echo -e "  ${YELLOW}${BOLD}[$CHECKS_OK/$total] — $CHECKS_FAIL check(s) en échec${NC}"
    fi
}

# ============================================================
# STATUS
# ============================================================
cmd_status() {
    CHECKS_OK=0; CHECKS_FAIL=0
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  FUNK — État du homelab"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    echo ""
    echo "[ storage-01 — services ]"
    check_service postgresql-16
    check_service qdrant
    check_service litellm
    check_service hermes-agent

    echo ""
    echo "[ storage-01 — santé ]"
    check_postgresql
    check_qdrant
    check_litellm
    check_hermes
    check_nfs_export

    echo ""
    echo "[ gpu-01 — 192.168.10.20 ]"
    if ping -c1 -W2 "$GPU01_IP" &>/dev/null; then
        check_ok "ping OK"
        check_llama_server
        check_embeddings
        if $GPU01_SSH "mountpoint -q /mnt/models" 2>/dev/null; then
            check_ok "NFS /mnt/models monté"
        else
            check_fail "NFS /mnt/models non monté"
        fi
    else
        check_fail "gpu-01 non joignable (éteint ?)"
    fi

    print_score
    echo ""
}

# ============================================================
# START
# ============================================================
cmd_start() {
    CHECKS_OK=0; CHECKS_FAIL=0
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  FUNK — Démarrage"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    # 1. Services storage-01
    echo ""
    echo "[ storage-01 — démarrage des services ]"
    for svc in postgresql-16 qdrant litellm hermes-agent; do
        systemctl start "$svc" 2>/dev/null
        sleep 1
        if systemctl is-active --quiet "$svc"; then
            ok "$svc démarré"
        else
            warn "$svc — problème au démarrage (voir: journalctl -u $svc -n 20)"
        fi
    done

    # 2. Vérifications santé storage-01
    echo ""
    echo "[ storage-01 — vérifications ]"
    sleep 2  # laisser les services s'initialiser
    check_postgresql
    check_qdrant
    check_litellm
    check_hermes
    check_nfs_export

    # 3. gpu-01 — Wake-on-LAN si éteint
    echo ""
    echo "[ gpu-01 — démarrage ]"
    if ping -c1 -W2 "$GPU01_IP" &>/dev/null; then
        check_ok "gpu-01 déjà allumé"
    else
        wol "$GPU01_MAC"
        info "Attente du démarrage de gpu-01..."
        if ! wait_ping "$GPU01_IP" 120; then
            check_fail "gpu-01 n'a pas répondu au ping en 120s"
            print_score
            return 1
        fi
        check_ok "gpu-01 joignable"
        # Attendre que SSH soit disponible
        local ssh_tries=0
        while ! $GPU01_SSH "true" 2>/dev/null; do
            sleep 3; ssh_tries=$((ssh_tries + 1))
            [ $ssh_tries -ge 10 ] && break
            echo -ne "  → Attente SSH gpu-01... $((ssh_tries * 3))s\r"
        done
        echo ""
    fi

    # 4. Vérifier NFS sur gpu-01
    info "Vérification NFS sur gpu-01..."
    local nfs_tries=0
    while ! $GPU01_SSH "mountpoint -q /mnt/models" 2>/dev/null; do
        sleep 3; nfs_tries=$((nfs_tries + 1))
        [ $nfs_tries -ge 20 ] && break
        echo -ne "  → Attente montage NFS... $((nfs_tries * 3))s\r"
    done
    echo ""
    if $GPU01_SSH "mountpoint -q /mnt/models" 2>/dev/null; then
        check_ok "NFS /mnt/models monté sur gpu-01"
    else
        check_fail "NFS /mnt/models non monté (modèle non accessible)"
    fi

    # 5. Attendre llama-server
    echo ""
    echo "[ gpu-01 — llama-server ]"
    info "Attente chargement du modèle LLM (jusqu'à 4 min)..."
    if wait_llama 240; then
        check_llama_server
        check_embeddings
    else
        check_fail "llama-server n'a pas répondu en 4 minutes"
    fi

    # 6. Score final
    print_score
    if [ $CHECKS_FAIL -eq 0 ]; then
        echo -e "  ${GREEN}Homelab Funk opérationnel — bonne session !${NC}"
    else
        echo -e "  ${YELLOW}Homelab démarré avec des avertissements${NC}"
    fi
    echo ""
}

# ============================================================
# STOP
# ============================================================
cmd_stop() {
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  FUNK — Arrêt"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    # 1. Arrêter Hermes et LiteLLM proprement
    echo ""
    echo "[ storage-01 — arrêt des services applicatifs ]"
    for svc in hermes-agent litellm; do
        systemctl stop "$svc" 2>/dev/null && ok "$svc arrêté" || warn "$svc déjà inactif"
    done

    # 2. Éteindre gpu-01
    echo ""
    echo "[ gpu-01 — arrêt ]"
    if ping -c1 -W2 "$GPU01_IP" &>/dev/null; then
        info "Arrêt de gpu-01..."
        $GPU01_SSH "sudo shutdown -h now" 2>/dev/null
        local tries=0
        while ping -c1 -W1 "$GPU01_IP" &>/dev/null; do
            sleep 2; tries=$((tries + 1))
            [ $tries -ge 30 ] && break
            echo -ne "  → Extinction en cours... $((tries * 2))s\r"
        done
        echo ""
        ok "gpu-01 éteint"
    else
        warn "gpu-01 déjà éteint"
    fi

    # 3. Arrêter Qdrant et PostgreSQL
    echo ""
    echo "[ storage-01 — arrêt des services données ]"
    for svc in qdrant postgresql-16; do
        systemctl stop "$svc" 2>/dev/null && ok "$svc arrêté" || warn "$svc déjà inactif"
    done

    echo ""
    ok "Homelab Funk arrêté proprement"
    echo ""
}

# ============================================================
case "${1:-status}" in
    start)  cmd_start ;;
    stop)   cmd_stop  ;;
    status) cmd_status ;;
    *)
        echo "Usage: funk-cluster [start|stop|status]"
        echo ""
        echo "  start   Démarre gpu-01 (WOL) + vérifie tous les services"
        echo "  stop    Arrêt propre gpu-01 + services storage-01"
        echo "  status  État de tous les services"
        exit 1
        ;;
esac