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

GPU01_IP="192.168.10.20"
GPU01_MAC="a8:a1:59:72:a6:11"
GPU01_SSH="ssh -o ConnectTimeout=5 -o StrictHostKeyChecking=no gpu-01"

K8S_IPS=(192.168.10.11 192.168.10.12 192.168.10.13)
K8S_MACS=("6c:4b:90:82:8e:47" "6c:4b:90:cf:7f:c5" "6c:4b:90:b6:49:20")
K8S_NAMES=(compute-01 compute-02 compute-03)
TALOSCTL="sudo -u ansible talosctl --talosconfig /home/ansible/.config/talos/config"
KUBECTL="sudo -u ansible kubectl"

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)); }

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()
"
}

check_service() {
    local svc=$1
    if systemctl is-active --quiet "$svc" 2>/dev/null; then
        check_ok "$svc"
    else
        check_fail "$svc (inactif)"
    fi
}

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
}

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); print(len(d.get('result',{}).get('collections',[])))" 2>/dev/null)
    if [ -n "$result" ]; then
        check_ok "Qdrant — $result collection(s)"
    else
        check_fail "Qdrant non joignable (port 6333)"
    fi
}

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; print(len(json.load(sys.stdin).get('data',[])))" 2>/dev/null)
    if [ -n "$models" ]; then
        check_ok "LiteLLM — $models modèle(s) enregistré(s)"
    else
        check_fail "LiteLLM non joignable (port 4000)"
    fi
}

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
}

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
}

check_llama_server() {
    local port=${1:-1234} label=${2:-"GPU"}
    local model
    model=$(curl -sf --max-time 5 "http://${GPU01_IP}:${port}/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 $label (:$port) — $model"
    else
        check_fail "llama-server $label (:$port) — non joignable"
    fi
}

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
}

check_k8s() {
    echo ""
    echo "[ cluster k8s — compute-01/02/03 ]"
    local up=0
    for i in 0 1 2; do
        local ip="${K8S_IPS[$i]}" name="${K8S_NAMES[$i]}"
        if ping -c1 -W2 "$ip" &>/dev/null; then
            check_ok "$name ($ip) — ping OK"
            up=$((up + 1))
        else
            warn "$name ($ip) — éteint"
        fi
    done
    if [ $up -gt 0 ]; then
        local ready
        ready=$($KUBECTL get nodes --no-headers 2>/dev/null | grep -c " Ready" || echo 0)
        if [ "$ready" -gt 0 ]; then
            check_ok "kubectl — $ready nœud(s) Ready"
        else
            warn "kubectl — nœuds non Ready (encore en démarrage ?)"
        fi
    fi
}

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
}

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

stop_service() {
    local svc=$1
    systemctl stop --timeout=10 "$svc" 2>/dev/null
    systemctl kill --signal=SIGKILL "$svc" 2>/dev/null
    ok "$svc arrêté"
}

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"
        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
        check_llama_server 1234 "GPU  (Qwen3-8B)"
        check_embeddings
        check_llama_server 1236 "CPU1 (Qwen3-1.7B system)"
        check_llama_server 1237 "CPU2 (Qwen3-1.7B monitor)"
    else
        check_fail "gpu-01 non joignable (éteint ?)"
    fi

    check_k8s
    print_score
    echo ""
}

# ============================================================
# START
# ============================================================
cmd_start() {
    local with_k8s=$1
    CHECKS_OK=0; CHECKS_FAIL=0
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  FUNK — Démarrage${with_k8s:+ (+ cluster k8s)}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    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
        systemctl is-active --quiet "$svc" \
            && ok "$svc démarré" \
            || warn "$svc — problème (journalctl -u $svc -n 20)"
    done

    echo ""
    echo "[ storage-01 — vérifications ]"
    sleep 2
    check_postgresql; check_qdrant; check_litellm; check_hermes; check_nfs_export

    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 "Paquet WOL envoyé → gpu-01"
        if ! wait_ping "$GPU01_IP" 120; then
            check_fail "gpu-01 n'a pas répondu en 120s"
            print_score; return 1
        fi
        check_ok "gpu-01 joignable"
        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... $((ssh_tries * 3))s\r"
        done; echo ""
    fi

    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 NFS... $((nfs_tries * 3))s\r"
    done; echo ""
    $GPU01_SSH "mountpoint -q /mnt/models" 2>/dev/null \
        && check_ok "NFS /mnt/models monté" \
        || check_fail "NFS /mnt/models non monté"

    echo ""
    echo "[ gpu-01 — llama-server GPU (port 1234) ]"
    info "Attente chargement modèle GPU (jusqu'à 4 min)..."
    if wait_llama 1234 240; then
        check_llama_server 1234 "GPU  (Qwen3-8B)"; check_embeddings
    else
        check_fail "llama-server GPU n'a pas répondu en 4 min"
    fi

    echo ""
    echo "[ gpu-01 — llama-server CPU (ports 1236/1237) ]"
    info "Attente chargement modèles CPU (jusqu'à 2 min)..."
    wait_llama 1236 120; check_llama_server 1236 "CPU1 (Qwen3-1.7B system)"
    wait_llama 1237 60;  check_llama_server 1237 "CPU2 (Qwen3-1.7B monitor)"

    if [ -n "$with_k8s" ]; then
        echo ""
        echo "[ cluster k8s — démarrage compute-01/02/03 ]"
        for i in 0 1 2; do
            local ip="${K8S_IPS[$i]}" name="${K8S_NAMES[$i]}" mac="${K8S_MACS[$i]}"
            if ping -c1 -W2 "$ip" &>/dev/null; then
                check_ok "$name déjà allumé"
            else
                wol "$mac"
                ok "$name — WOL envoyé"
            fi
        done
        info "Attente ping des nœuds..."
        for ip in "${K8S_IPS[@]}"; do
            wait_ping "$ip" 120 && ok "$ip joignable" || warn "$ip n'a pas répondu"
        done
        info "Attente nœuds Ready (jusqu'à 3 min)..."
        local k8s_elapsed=0
        until [ $($KUBECTL get nodes --no-headers 2>/dev/null | grep -c " Ready" || echo 0) -ge 3 ] \
              || [ $k8s_elapsed -ge 180 ]; do
            sleep 10; k8s_elapsed=$((k8s_elapsed + 10))
            echo -ne "  → Attente nœuds Ready... ${k8s_elapsed}s\r"
        done; echo ""
        local ready=$($KUBECTL get nodes --no-headers 2>/dev/null | grep -c " Ready" || echo 0)
        [ "$ready" -ge 1 ] \
            && check_ok "kubectl — $ready nœud(s) Ready" \
            || warn "Nœuds pas encore Ready (vérifier avec kubectl get nodes)"
    fi

    print_score
    [ $CHECKS_FAIL -eq 0 ] \
        && echo -e "  ${GREEN}Homelab Funk opérationnel — bonne session !${NC}" \
        || echo -e "  ${YELLOW}Homelab démarré avec des avertissements${NC}"
    echo ""
}

# ============================================================
# STOP
# ============================================================
cmd_stop() {
    local with_k8s=$1
    echo ""
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"
    echo "  FUNK — Arrêt${with_k8s:+ (+ cluster k8s)}"
    echo "━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━━"

    echo ""
    echo "[ storage-01 — arrêt des services applicatifs ]"
    for svc in hermes-agent litellm; do stop_service "$svc"; done

    if [ -n "$with_k8s" ]; then
        echo ""
        echo "[ cluster k8s — arrêt propre compute-01/02/03 ]"

        # Quels nœuds sont en vie ?
        local alive_ips=() alive_names=()
        for i in 0 1 2; do
            if ping -c1 -W2 "${K8S_IPS[$i]}" &>/dev/null; then
                alive_ips+=("${K8S_IPS[$i]}")
                alive_names+=("${K8S_NAMES[$i]}")
            fi
        done

        if [ ${#alive_ips[@]} -eq 0 ]; then
            warn "Nœuds k8s déjà éteints"
        else
            # 1. Cordon de tous les nœuds — empêche tout rescheduling pendant le drain
            info "Cordon (plus aucun pod ne peut être replanifié)..."
            for name in "${alive_names[@]}"; do
                $KUBECTL cordon "$name" 2>/dev/null \
                    && ok "$name — cordonné" \
                    || warn "$name — cordon échoué (nœud peut-être inaccessible)"
            done

            # 2. Drain workers en premier (compute-02, compute-03)
            #    Les pods reçoivent SIGTERM → flush WAL/SQLite → PVCs NFS démontés
            echo ""
            info "Drain workers (SIGTERM pods, 30s grace, 90s max)..."
            for i in 1 2; do
                if ping -c1 -W2 "${K8S_IPS[$i]}" &>/dev/null; then
                    info "  Drain ${K8S_NAMES[$i]}..."
                    $KUBECTL drain "${K8S_NAMES[$i]}" \
                        --ignore-daemonsets \
                        --delete-emptydir-data \
                        --grace-period=30 \
                        --timeout=90s \
                        --force 2>&1 | grep -vE "^(Warning|I[0-9])" | tail -3 || true
                    ok "${K8S_NAMES[$i]} — pods évacués, PVCs NFS libérés"
                fi
            done

            # 3. Drain control-plane (compute-01) — ArgoCD, kube-state-metrics, etc.
            if ping -c1 -W2 "${K8S_IPS[0]}" &>/dev/null; then
                echo ""
                info "Drain control-plane (compute-01)..."
                $KUBECTL drain compute-01 \
                    --ignore-daemonsets \
                    --delete-emptydir-data \
                    --grace-period=30 \
                    --timeout=90s \
                    --force 2>&1 | grep -vE "^(Warning|I[0-9])" | tail -3 || true
                ok "compute-01 — drainé"
            fi

            # 4. Vérification NFS : plus aucune session active depuis les nœuds k8s
            echo ""
            info "Vérification montages NFS..."
            local nfs_wait=0
            while ss -tn 2>/dev/null | grep ':2049' | grep -qE '10\.1[123]\.'; do
                sleep 3; nfs_wait=$((nfs_wait + 3))
                [ $nfs_wait -ge 30 ] && break
                echo -ne "  → NFS encore actif depuis nœuds k8s... ${nfs_wait}s\r"
            done; echo ""
            if ss -tn 2>/dev/null | grep ':2049' | grep -qE '10\.1[123]\.'; then
                warn "Sessions NFS toujours actives après 30s — extinction quand même"
            else
                ok "NFS libéré — aucune écriture en cours depuis le cluster"
            fi

            # 5. Extinction via talosctl (kernel halt propre — plus rien ne tourne)
            echo ""
            local nodes=$(IFS=,; echo "${alive_ips[*]}")
            info "talosctl shutdown → ${alive_ips[*]}"
            if $TALOSCTL shutdown --nodes "$nodes" 2>/dev/null; then
                ok "Shutdown envoyé aux nœuds"
            else
                warn "Erreur talosctl shutdown — extinction manuelle requise"
                return
            fi

            # 6. Attente extinction complète (ping)
            local tries=0
            while [ $tries -lt 20 ]; do
                local still_up=0
                for ip in "${alive_ips[@]}"; do
                    ping -c1 -W1 "$ip" &>/dev/null && still_up=$((still_up + 1))
                done
                [ $still_up -eq 0 ] && break
                sleep 3; tries=$((tries + 1))
                echo -ne "  → Extinction k8s... $((tries * 3))s\r"
            done; echo ""
            ok "Cluster k8s éteint proprement"
        fi
    fi

    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 gpu-01... $((tries * 2))s\r"
        done; echo ""
        ok "gpu-01 éteint"
    else
        warn "gpu-01 déjà éteint"
    fi

    echo ""
    echo "[ storage-01 — arrêt des services données ]"
    for svc in qdrant postgresql-16; do stop_service "$svc"; done

    echo ""
    ok "Homelab Funk arrêté proprement"
    echo ""
    echo "[ storage-01 — extinction dans 3s ]"
    info "La connexion SSH va être coupée..."
    sleep 3
    sudo shutdown -h now
}

# ============================================================
CMD="${1:-status}"
WITH_K8S=""
for arg in "$@"; do [ "$arg" = "--k8s" ] && WITH_K8S=1; done

case "$CMD" in
    start)  cmd_start "$WITH_K8S" ;;
    stop)   cmd_stop  "$WITH_K8S" ;;
    status) cmd_status ;;
    *)
        echo "Usage: funk-cluster [start|stop|status] [--k8s]"
        echo ""
        echo "  start          Démarre gpu-01 (WOL) + vérifie tous les services"
        echo "  start --k8s    Idem + démarre le cluster k8s (compute-01/02/03)"
        echo "  stop           Arrêt propre gpu-01 + services storage-01"
        echo "  stop  --k8s    Idem + arrêt propre k8s (drain → NFS check → talosctl shutdown)"
        echo "  status         État complet (storage-01 + gpu-01 + k8s)"
        exit 1
        ;;
esac