#!/usr/bin/env bash
# check-gpu-availability.sh — Check Scaleway GPU availability and select fastest path
# Usage:
#   ./scripts/check-gpu-availability.sh           # show all + recommend fastest
#   ./scripts/check-gpu-availability.sh --best     # output only the best pool name (for CI)
#   ./scripts/check-gpu-availability.sh H100       # filter by type
#
# Exit codes: 0 = found available GPU, 1 = nothing available, 2 = error
set -euo pipefail

CLUSTER_ZONE="fr-par-2"
BEST_ONLY=false
FILTER=""

for arg in "$@"; do
    case "$arg" in
        --best) BEST_ONLY=true ;;
        *)      FILTER="$arg" ;;
    esac
done

# GPU types ranked by speed (fastest first)
# Format: scw_type|pool_name|label|gpu_count|approx_tflops
GPU_RANK=(
    "H100-SXM-8-80G|ci-training-h100-sxm8|H100 SXM 8x80GB|8|26400"
    "H100-SXM-4-80G|ci-training-h100-sxm4|H100 SXM 4x80GB|4|13200"
    "H100-SXM-2-80G|ci-training-h100-sxm|H100 SXM 2x80GB|2|6600"
    "H100-2-80G|ci-training-h100x2|H100 PCIe 2x80GB|2|5280"
    "H100-1-80G|ci-training-h100|H100 PCIe 1x80GB|1|2640"
    "L40S-1-48G|ci-training|L40S 1x48GB|1|733"
)

get_availability() {
    local type=$1
    local avail
    avail=$(scw instance server-type list zone="$CLUSTER_ZONE" 2>/dev/null | grep "^${type} " || true)
    if [[ -z "$avail" ]]; then
        echo "unavailable"
        return
    fi
    if echo "$avail" | grep -q "out of stock"; then
        echo "out_of_stock"
    elif echo "$avail" | grep -q "low stock"; then
        echo "low_stock"
    elif echo "$avail" | grep -q "available"; then
        echo "available"
    else
        echo "unknown"
    fi
}

get_pool_status() {
    local pool_name=$1
    local cluster_id
    cluster_id=$(scw k8s cluster list -o json 2>/dev/null | jq -r '.[0].id // empty')
    [[ -z "$cluster_id" ]] && echo "no_cluster" && return
    local status
    status=$(scw k8s pool list cluster-id="$cluster_id" -o json 2>/dev/null | \
        jq -r ".[] | select(.name==\"$pool_name\") | .status" 2>/dev/null)
    echo "${status:-no_pool}"
}

# Find the fastest available GPU
best_pool=""
best_label=""
best_type=""

if ! $BEST_ONLY; then
    echo "╔══════════════════════════════════════════════════════════════╗"
    echo "║          Scaleway GPU — Fastest Path Selection              ║"
    echo "╠══════════════════════════════════════════════════════════════╣"
    printf "║  Zone: %-52s║\n" "$CLUSTER_ZONE"
    printf "║  Time: %-52s║\n" "$(date -u +%Y-%m-%dT%H:%M:%SZ)"
    echo "╚══════════════════════════════════════════════════════════════╝"
    echo ""
    printf "  %-22s %-6s %-13s %-12s %s\n" "TYPE" "GPUs" "TFLOPS(FP16)" "STOCK" "POOL"
    echo "  ─────────────────────────────────────────────────────────────"
fi

for entry in "${GPU_RANK[@]}"; do
    IFS='|' read -r scw_type pool_name label gpu_count tflops <<< "$entry"

    # Apply filter
    if [[ -n "$FILTER" && "$scw_type" != *"$FILTER"* && "$label" != *"$FILTER"* ]]; then
        continue
    fi

    avail=$(get_availability "$scw_type")

    # Color and symbol
    symbol="✗"
    color="\033[31m"
    case "$avail" in
        available)    symbol="✓"; color="\033[32m" ;;
        low_stock)    symbol="~"; color="\033[33m" ;;
        out_of_stock) symbol="✗"; color="\033[31m" ;;
        *)            symbol="?"; color="\033[90m" ;;
    esac

    # Track best available
    if [[ -z "$best_pool" && ("$avail" == "available" || "$avail" == "low_stock") ]]; then
        best_pool="$pool_name"
        best_label="$label"
        best_type="$scw_type"
    fi

    if ! $BEST_ONLY; then
        printf "  %b%s %-21s\033[0m %-6s %-13s %b%-12s\033[0m %s\n" \
            "$color" "$symbol" "$scw_type" "$gpu_count" "$tflops" "$color" "$avail" "$pool_name"
    fi
done

if ! $BEST_ONLY; then
    echo ""

    # Kapsule pool status
    echo "  ─── Kapsule GPU Pools ───"
    CLUSTER_ID=$(scw k8s cluster list -o json 2>/dev/null | jq -r '.[0].id // empty')
    if [[ -n "$CLUSTER_ID" ]]; then
        scw k8s pool list cluster-id="$CLUSTER_ID" -o json 2>/dev/null | \
            jq -r '.[] | select(.node_type | test("h100|l40|l4_|gpu")) | "  \(.name): \(.node_type) size=\(.size)/\(.max_size) status=\(.status)"'
    fi
    echo ""

    # Active GPU nodes
    echo "  ─── Active GPU Nodes ───"
    gpu_nodes=$(kubectl get nodes -l k8s.scaleway.com/pool-name -o json 2>/dev/null | \
        jq -r '.items[] | select(.metadata.labels["k8s.scaleway.com/pool-name"] | test("training|gpu")) | "  \(.metadata.name) (\(.metadata.labels["k8s.scaleway.com/pool-name"]))"' 2>/dev/null || true)
    echo "${gpu_nodes:-  (none running)}"
    echo ""

    # Node errors
    if [[ -n "$CLUSTER_ID" ]]; then
        errors=$(scw k8s node list cluster-id="$CLUSTER_ID" -o json 2>/dev/null | \
            jq -r '.[] | select(.status=="creation_error") | "  ⚠ \(.name): \(.error_message)"' 2>/dev/null || true)
        if [[ -n "$errors" ]]; then
            echo "  ─── Provisioning Errors ───"
            echo "$errors"
            echo ""
        fi
    fi

    # Recommendation
    if [[ -n "$best_pool" ]]; then
        echo "  ══════════════════════════════════════════════════════════"
        echo "  FASTEST AVAILABLE: $best_label"
        echo "  Pool: $best_pool  (type: $best_type)"
        echo "  ══════════════════════════════════════════════════════════"
    else
        echo "  ⚠ No GPU instances available in $CLUSTER_ZONE"
    fi
fi

# --best mode: just output pool name for CI scripting
if $BEST_ONLY; then
    if [[ -n "$best_pool" ]]; then
        echo "$best_pool"
        exit 0
    else
        exit 1
    fi
fi

[[ -n "$best_pool" ]] && exit 0 || exit 1