#!/usr/bin/env bash
# ============================================================================
# BLACKROAD OS, INC. - PROPRIETARY AND CONFIDENTIAL
# Copyright (c) 2025-2026 BlackRoad OS, Inc. All Rights Reserved.
# ============================================================================
# br-fleet-ai - Fleet AI status and orchestration dashboard
# Shows all AI capabilities across the BlackRoad fleet at a glance
# Usage: br-fleet-ai [status|capacity|route|benchmark]
set -eo pipefail

source "$HOME/.blackroad/config/nodes.sh" 2>/dev/null || true

OLLAMA_NODES=(cecilia lucidia alice octavia)
HAILO_NODES=(cecilia octavia)

cmd_status() {
  printf '%b╔══════════════════════════════════════════════════════════════╗%b\n' "$PINK" "$RESET"
  printf '%b║            BlackRoad Fleet AI Status                         ║%b\n' "$PINK" "$RESET"
  printf '%b╚══════════════════════════════════════════════════════════════╝%b\n\n' "$PINK" "$RESET"

  local total_models=0
  local online_nodes=0
  local total_tops=0

  # Ollama nodes
  printf '  %b── Ollama Inference Nodes ──%b\n\n' "$BLUE" "$RESET"
  printf '  %-12s %-8s %-6s %-24s %s\n' "NODE" "STATUS" "MODELS" "TOP MODEL" "ENDPOINT"
  printf '  %-12s %-8s %-6s %-24s %s\n' "────" "──────" "──────" "─────────" "────────"

  for node in "${OLLAMA_NODES[@]}"; do
    local ip="${NODE_IP[$node]:-}"
    [[ -z "$ip" ]] && continue

    printf '  %-12s ' "$node"

    local tags
    tags=$(curl -sf --connect-timeout 3 "http://${ip}:11434/api/tags" 2>/dev/null)
    if [[ -n "$tags" ]]; then
      local count
      count=$(echo "$tags" | jq '.models | length' 2>/dev/null || echo "0")
      local top_model
      top_model=$(echo "$tags" | jq -r '.models[0].name // "—"' 2>/dev/null)
      total_models=$((total_models + count))
      ((online_nodes++))

      printf '%b%-8s%b %-6s %-24s %s\n' "$GREEN" "ONLINE" "$RESET" "$count" "$top_model" "${ip}:11434"
    else
      printf '%b%-8s%b %-6s %-24s %s\n' "$RED" "OFFLINE" "$RESET" "—" "—" "${ip}:11434"
    fi
  done

  echo ""

  # Hailo accelerators
  printf '  %b── AI Accelerators (Hailo-8) ──%b\n\n' "$BLUE" "$RESET"
  printf '  %-12s %-8s %-10s %s\n' "NODE" "STATUS" "TOPS" "DEVICE"
  printf '  %-12s %-8s %-10s %s\n' "────" "──────" "────" "──────"

  for node in "${HAILO_NODES[@]}"; do
    local ip="${NODE_IP[$node]:-}"
    local user="${NODE_USER[$node]:-}"
    [[ -z "$ip" ]] && continue

    printf '  %-12s ' "$node"

    if br_ping "$node" 2>/dev/null; then
      local has_hailo
      has_hailo=$(ssh -o ConnectTimeout=3 -o BatchMode=yes "${user}@${ip}" \
        "test -e /dev/hailo0 && echo yes || echo no" 2>/dev/null) || has_hailo="?"

      if [[ "$has_hailo" == "yes" ]]; then
        printf '%b%-8s%b %-10s %s\n' "$GREEN" "ACTIVE" "$RESET" "26" "/dev/hailo0"
        total_tops=$((total_tops + 26))
      else
        printf '%b%-8s%b %-10s %s\n' "$AMBER" "NO DEV" "$RESET" "—" "—"
      fi
    else
      printf '%b%-8s%b %-10s %s\n' "$RED" "OFFLINE" "$RESET" "—" "—"
    fi
  done

  echo ""

  # Cloud AI providers
  printf '  %b── Cloud AI Providers ──%b\n\n' "$BLUE" "$RESET"

  local providers=("ANTHROPIC_API_KEY:Anthropic:claude-sonnet-4-20250514" "OPENAI_API_KEY:OpenAI:gpt-4o" "GOOGLE_AI_API_KEY:Gemini:gemini-1.5-pro" "XAI_API_KEY:xAI/Grok:grok-beta")
  for entry in "${providers[@]}"; do
    IFS=':' read -r key name model <<< "$entry"
    printf '  %-12s ' "$name"
    if [[ -n "${!key:-}" ]]; then
      printf '%b%-8s%b %s\n' "$GREEN" "KEY SET" "$RESET" "$model"
    else
      printf '%b%-8s%b\n' "$AMBER" "NO KEY" "$RESET"
    fi
  done

  echo ""

  # Summary
  printf '  %b── Summary ──%b\n\n' "$PINK" "$RESET"
  printf '  Ollama nodes online:  %b%d%b\n' "$GREEN" "$online_nodes" "$RESET"
  printf '  Total models loaded:  %b%d%b\n' "$GREEN" "$total_models" "$RESET"
  printf '  Hailo-8 TOPS:         %b%d%b\n' "$GREEN" "$total_tops" "$RESET"
  printf '  Cloud providers:      %b%d%b configured\n' "$GREEN" \
    "$(env | grep -cE 'ANTHROPIC_API_KEY|OPENAI_API_KEY|GOOGLE_AI_API_KEY|XAI_API_KEY' 2>/dev/null || echo 0)" "$RESET"
  echo
}

cmd_capacity() {
  printf '%bFleet AI Capacity:%b\n\n' "$PINK" "$RESET"

  for node in "${OLLAMA_NODES[@]}"; do
    local ip="${NODE_IP[$node]:-}"
    local user="${NODE_USER[$node]:-}"
    [[ -z "$ip" ]] && continue

    printf '  %b=== %s (%s) ===%b\n' "$BLUE" "$node" "$ip" "$RESET"

    # Try to get system info via SSH
    if br_ping "$node" 2>/dev/null; then
      local info
      info=$(ssh -o ConnectTimeout=5 -o BatchMode=yes "${user}@${ip}" \
        "echo \"CPU: \$(nproc) cores\"; \
         echo \"RAM: \$(free -m 2>/dev/null | awk '/Mem:/{printf \"%dMB/%dMB (%.0f%%)\", \$3, \$2, \$3*100/\$2}')\"; \
         echo \"Temp: \$(vcgencmd measure_temp 2>/dev/null | grep -oP '[0-9.]+' || echo '?')°C\"; \
         echo \"Load: \$(cat /proc/loadavg 2>/dev/null | awk '{print \$1, \$2, \$3}')\"" 2>/dev/null) || info="SSH failed"
      echo "$info" | sed 's/^/    /'
    else
      echo "    OFFLINE"
    fi

    # Models
    local tags
    tags=$(curl -sf --connect-timeout 3 "http://${ip}:11434/api/tags" 2>/dev/null)
    if [[ -n "$tags" ]]; then
      printf '    Models: '
      echo "$tags" | jq -r '[.models[].name] | join(", ")' 2>/dev/null
    fi
    echo
  done
}

cmd_route() {
  local prompt="${1:-test}"
  printf '%bRoute decision for: "%s"%b\n\n' "$AMBER" "${prompt:0:50}" "$RESET"

  printf '  Checking nodes...\n'

  local best_node="" best_lat=999999 best_ip=""

  for node in "${OLLAMA_NODES[@]}"; do
    local ip="${NODE_IP[$node]:-}"
    [[ -z "$ip" ]] && continue

    printf '    %-12s ' "$node"

    local start_ms=$(date +%s%3N 2>/dev/null || python3 -c 'import time; print(int(time.time()*1000))')
    if curl -sf --connect-timeout 2 "http://${ip}:11434/api/tags" &>/dev/null; then
      local end_ms=$(date +%s%3N 2>/dev/null || python3 -c 'import time; print(int(time.time()*1000))')
      local lat=$((end_ms - start_ms))
      printf '%b%dms%b\n' "$GREEN" "$lat" "$RESET"

      if [[ $lat -lt $best_lat ]]; then
        best_lat=$lat
        best_node="$node"
        best_ip="$ip"
      fi
    else
      printf '%bDOWN%b\n' "$RED" "$RESET"
    fi
  done

  echo ""
  if [[ -n "$best_node" ]]; then
    printf '  %bSelected: %s (%s) — %dms latency%b\n' "$GREEN" "$best_node" "$best_ip" "$best_lat" "$RESET"
  else
    printf '  %bNo nodes available%b\n' "$RED" "$RESET"
  fi
}

cmd_benchmark() {
  local model="${1:-$DEFAULT_MODEL}"
  local iterations="${2:-5}"

  printf '%bBenchmarking %s across fleet (%d iterations)...%b\n\n' "$AMBER" "$model" "$iterations" "$RESET"

  local test_prompt="Reply with exactly one word: hello"

  for node in "${OLLAMA_NODES[@]}"; do
    local ip="${NODE_IP[$node]:-}"
    [[ -z "$ip" ]] && continue

    printf '  %-12s ' "$node"

    if ! curl -sf --connect-timeout 2 "http://${ip}:11434/api/tags" &>/dev/null; then
      printf '%bOFFLINE%b\n' "$RED" "$RESET"
      continue
    fi

    local total_lat=0 success=0
    for i in $(seq 1 "$iterations"); do
      local start_ms=$(date +%s%3N 2>/dev/null || python3 -c 'import time; print(int(time.time()*1000))')

      local resp
      resp=$(curl -sf --max-time 30 "http://${ip}:11434/api/generate" \
        -d "$(jq -n --arg m "$model" --arg p "$test_prompt" '{model: $m, prompt: $p, stream: false}')" 2>/dev/null | \
        jq -r '.response // empty' 2>/dev/null)

      local end_ms=$(date +%s%3N 2>/dev/null || python3 -c 'import time; print(int(time.time()*1000))')

      if [[ -n "$resp" ]]; then
        local lat=$((end_ms - start_ms))
        total_lat=$((total_lat + lat))
        ((success++))
      fi
    done

    if [[ $success -gt 0 ]]; then
      local avg=$((total_lat / success))
      printf '%b%d/%d OK  avg: %dms%b\n' "$GREEN" "$success" "$iterations" "$avg" "$RESET"
    else
      printf '%bAll failed%b\n' "$RED" "$RESET"
    fi
  done
  echo
}

usage() {
  cat <<EOF
${PINK}br-fleet-ai${RESET} - Fleet AI Status & Orchestration

${BLUE}COMMANDS:${RESET}
  status                Show full fleet AI status (default)
  capacity              Show compute capacity per node
  route [prompt]        Show routing decision for a prompt
  benchmark [model] [n] Benchmark model across fleet

${GREEN}EXAMPLES:${RESET}
  br-fleet-ai                     Full status dashboard
  br-fleet-ai capacity            RAM/CPU/models per node
  br-fleet-ai route "hello"       See which node would handle it
  br-fleet-ai benchmark llama3.2  Speed test across fleet
EOF
}

case "${1:-status}" in
  status|s)     cmd_status ;;
  capacity|c)   cmd_capacity ;;
  route|r)      shift; cmd_route "$@" ;;
  benchmark|b)  shift; cmd_benchmark "$@" ;;
  -h|--help|help) usage ;;
  *)            echo "Unknown: $1"; usage; exit 1 ;;
esac