#!/usr/bin/env bash
# ============================================================================
# BLACKROAD OS, INC. - PROPRIETARY AND CONFIDENTIAL
# Copyright (c) 2025-2026 BlackRoad OS, Inc. All Rights Reserved.
# ============================================================================
# br-models - BlackRoad AI Model Registry
# Discover, catalog, and manage models across the fleet
# Usage: br-models <list|scan|pull|remove|info>
set -eo pipefail

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

MODELS_DIR="$HOME/.blackroad/ai"
MODELS_DB="$MODELS_DIR/models.db"
mkdir -p "$MODELS_DIR"

_sql() { sqlite3 "$MODELS_DB" "$@" 2>/dev/null; }
_sql_escape() { echo "$1" | sed "s/'/''/g"; }

OLLAMA_NODES=(cecilia lucidia alice octavia)

init_db() {
  _sql <<'SQL'
CREATE TABLE IF NOT EXISTS models (
    id TEXT PRIMARY KEY,
    name TEXT,
    node TEXT,
    size TEXT,
    parameter_size TEXT,
    quantization TEXT,
    family TEXT,
    format TEXT,
    modified DATETIME,
    discovered DATETIME DEFAULT CURRENT_TIMESTAMP
);
CREATE TABLE IF NOT EXISTS model_nodes (
    model TEXT,
    node TEXT,
    status TEXT DEFAULT 'available',
    last_checked DATETIME DEFAULT CURRENT_TIMESTAMP,
    PRIMARY KEY (model, node)
);
CREATE INDEX IF NOT EXISTS idx_model_name ON models(name);
PRAGMA journal_mode=WAL;
PRAGMA busy_timeout=5000;
SQL
}

cmd_scan() {
  printf '%bScanning fleet for AI models...%b\n\n' "$AMBER" "$RESET"

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

    printf '  %-10s (%s): ' "$node" "$ip"

    local tags
    tags=$(curl -sf --connect-timeout 3 "http://${ip}:11434/api/tags" 2>/dev/null)
    if [[ -z "$tags" ]]; then
      printf '%boffline%b\n' "$RED" "$RESET"
      continue
    fi

    local count
    count=$(echo "$tags" | jq '.models | length' 2>/dev/null || echo "0")
    printf '%b%s models%b\n' "$GREEN" "$count" "$RESET"
    total=$((total + count))

    # Store each model
    echo "$tags" | jq -r '.models[] | [.name, .size // 0, .details.parameter_size // "", .details.quantization_level // "", .details.family // "", .details.format // "", .modified_at // ""] | @tsv' 2>/dev/null | \
    while IFS=$'\t' read -r name size param_size quant family format modified; do
      local id="${node}:${name}"
      local size_human
      if [[ "$size" -gt 1073741824 ]]; then
        size_human="$(echo "scale=1; $size/1073741824" | bc 2>/dev/null || echo "?")GB"
      elif [[ "$size" -gt 1048576 ]]; then
        size_human="$(echo "scale=0; $size/1048576" | bc 2>/dev/null || echo "?")MB"
      else
        size_human="${size}B"
      fi

      _sql "INSERT OR REPLACE INTO models (id, name, node, size, parameter_size, quantization, family, format, modified)
            VALUES ('$(_sql_escape "$id")', '$(_sql_escape "$name")', '$node', '$size_human', '$(_sql_escape "$param_size")', '$(_sql_escape "$quant")', '$(_sql_escape "$family")', '$(_sql_escape "$format")', '$(_sql_escape "$modified")')"
      _sql "INSERT OR REPLACE INTO model_nodes (model, node, status, last_checked)
            VALUES ('$(_sql_escape "$name")', '$node', 'available', datetime('now'))"
    done
  done

  printf '\n  %bTotal: %d models across fleet%b\n' "$PINK" "$total" "$RESET"
}

cmd_list() {
  local filter="${1:-}"
  printf '%b%-30s %-10s %-8s %-10s %-12s%b\n' "$BLUE" "MODEL" "NODE" "SIZE" "PARAMS" "QUANT" "$RESET"
  printf '%-30s %-10s %-8s %-10s %-12s\n' "─────" "────" "────" "──────" "─────"

  local where=""
  [[ -n "$filter" ]] && where="WHERE name LIKE '%$(_sql_escape "$filter")%'"

  _sql "SELECT name, node, size, parameter_size, quantization FROM models $where ORDER BY name, node" | \
  while IFS='|' read -r name node size params quant; do
    printf '%-30s %-10s %-8s %-10s %-12s\n' "$name" "$node" "$size" "${params:-—}" "${quant:-—}"
  done
}

cmd_where() {
  local model="${1:?model name required}"
  printf '%bNodes with %s:%b\n' "$BLUE" "$model" "$RESET"
  _sql "SELECT node, status, last_checked FROM model_nodes WHERE model LIKE '%$(_sql_escape "$model")%' ORDER BY node" | \
  while IFS='|' read -r node status checked; do
    local status_color="$GREEN"
    [[ "$status" != "available" ]] && status_color="$RED"
    printf '  %-10s %b%-10s%b (checked: %s)\n' "$node" "$status_color" "$status" "$RESET" "$checked"
  done
}

cmd_pull() {
  local model="${1:?model name required}"
  local node="${2:?node required}"
  local ip="${NODE_IP[$node]:-}"
  [[ -z "$ip" ]] && { echo "Unknown node: $node" >&2; return 1; }

  printf '%bPulling %s on %s (%s)...%b\n' "$AMBER" "$model" "$node" "$ip" "$RESET"
  curl -sf "http://${ip}:11434/api/pull" \
    -d "{\"name\":\"$model\"}" 2>/dev/null | \
    jq -r '.status // empty' 2>/dev/null | tail -1

  printf '%bDone. Run: br-models scan%b\n' "$GREEN" "$RESET"
}

cmd_remove() {
  local model="${1:?model name required}"
  local node="${2:?node required}"
  local ip="${NODE_IP[$node]:-}"
  [[ -z "$ip" ]] && { echo "Unknown node: $node" >&2; return 1; }

  printf '%bRemoving %s from %s...%b\n' "$AMBER" "$model" "$node" "$RESET"
  curl -sf -X DELETE "http://${ip}:11434/api/delete" \
    -d "{\"name\":\"$model\"}" 2>/dev/null

  _sql "DELETE FROM models WHERE name='$(_sql_escape "$model")' AND node='$node'"
  _sql "DELETE FROM model_nodes WHERE model='$(_sql_escape "$model")' AND node='$node'"
  printf '%bRemoved%b\n' "$GREEN" "$RESET"
}

cmd_summary() {
  printf '%bModel Registry Summary%b\n\n' "$PINK" "$RESET"

  printf '  Unique models:  %s\n' "$(_sql "SELECT COUNT(DISTINCT name) FROM models")"
  printf '  Total instances: %s\n' "$(_sql "SELECT COUNT(*) FROM models")"
  printf '  Fleet nodes:    %s\n\n' "$(_sql "SELECT COUNT(DISTINCT node) FROM models")"

  printf '  %bBy family:%b\n' "$BLUE" "$RESET"
  _sql "SELECT COALESCE(family,'unknown'), COUNT(*) FROM models GROUP BY family ORDER BY COUNT(*) DESC LIMIT 10" | \
  while IFS='|' read -r family cnt; do
    printf '    %-20s %d\n' "$family" "$cnt"
  done

  echo ""
  printf '  %bBy node:%b\n' "$BLUE" "$RESET"
  _sql "SELECT node, COUNT(*) FROM models GROUP BY node ORDER BY COUNT(*) DESC" | \
  while IFS='|' read -r node cnt; do
    printf '    %-10s %d models\n' "$node" "$cnt"
  done
}

usage() {
  cat <<EOF
${PINK}br-models${RESET} - BlackRoad AI Model Registry

${BLUE}COMMANDS:${RESET}
  scan                         Discover models across fleet
  list [filter]                List all models (optional filter)
  where <model>                Show which nodes have a model
  pull <model> <node>          Pull model to a node
  remove <model> <node>        Remove model from a node
  summary                     Registry statistics

${GREEN}EXAMPLES:${RESET}
  br-models scan               Survey all Ollama nodes
  br-models list llama         Find all llama models
  br-models where mistral      Which nodes have mistral?
  br-models pull phi3 lucidia  Pull phi3 to lucidia
  br-models summary            Fleet model stats
EOF
}

[[ -f "$MODELS_DB" ]] || init_db

case "${1:-}" in
  scan)    cmd_scan ;;
  list|ls) cmd_list "${2:-}" ;;
  where|w) cmd_where "${2:?model name required}" ;;
  pull)    cmd_pull "${2:?model}" "${3:?node}" ;;
  remove|rm) cmd_remove "${2:?model}" "${3:?node}" ;;
  summary|s) cmd_summary ;;
  -h|--help|help|"") usage ;;
  *)       cmd_list "$1" ;;  # default to search
esac