#!/usr/bin/env bash
set -euo pipefail

ROOT=/home/daniel/aeon-spark-test/h3
STACK_DIR="$ROOT/spark-stack"
COMPOSE_FILE="$STACK_DIR/compose.spark-stack.yml"

usage() {
  cat <<'EOF'
Usage:
  spark_stack_spark.sh start h3|qwen|flux
  spark_stack_spark.sh stop
  spark_stack_spark.sh status

Ports:
  h3   -> http://192.168.1.162:8001
  qwen -> http://192.168.1.162:8002
  flux -> http://192.168.1.162:8192
EOF
}

managed_containers=(
  h3-hot-runtime
  h3-blackwell-runtime-h3-hot-runtime-1
  qwen38-vllm
  flux2-klein-comfy
  flux-nvfp4-test
)

stop_managed() {
  for name in "${managed_containers[@]}"; do
    if docker ps -a --format '{{.Names}}' | grep -qx "$name"; then
      docker stop "$name" >/dev/null 2>&1 || true
      docker rm "$name" >/dev/null 2>&1 || true
    fi
  done
}

case "${1:-}" in
  start)
    target="${2:-}"
    case "$target" in
      h3|qwen|flux) ;;
      *) usage; exit 2 ;;
    esac
    mkdir -p "$ROOT/vllm-cache"
    stop_managed
    docker compose -f "$COMPOSE_FILE" --profile "$target" up -d
    ;;
  stop)
    stop_managed
    ;;
  status)
    docker ps -a --format 'table {{.Names}}\t{{.Status}}\t{{.Ports}}' | grep -E 'NAMES|h3-hot-runtime|qwen38-vllm|flux2-klein-comfy|flux-nvfp4-test' || true
    nvidia-smi --query-compute-apps=pid,process_name,used_memory --format=csv,noheader,nounits || true
    ;;
  *)
    usage
    exit 2
    ;;
esac