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

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"

IMAGES_DIR="${IMAGES_DIR:-$SCRIPT_DIR/images}"
COMPOSE_TGZ="${COMPOSE_TGZ:-$SCRIPT_DIR/docker-compose.tgz}"
COMPOSE_DIR="${COMPOSE_DIR:-$SCRIPT_DIR/docker-compose}"
PROJECT_NAME="${PROJECT_NAME:-sscc_models}"

SIMAI_GPU="${SIMAI_GPU:-0}"
RECOPILOT_GPU="${RECOPILOT_GPU:-1}"
CRYPTO_GPU="${CRYPTO_GPU:-2}"
MILVUS_GPU="${MILVUS_GPU:-3}"
QWEN_GPUS="${QWEN_GPUS:-4,5,6,7}"
GPU_BUSY_THRESHOLD_MB="${GPU_BUSY_THRESHOLD_MB:-500}"
SERVICE_CHECK_TIMEOUT="${SERVICE_CHECK_TIMEOUT:-1800}"
SERVICE_CHECK_INTERVAL="${SERVICE_CHECK_INTERVAL:-15}"
CURL_TIMEOUT="${CURL_TIMEOUT:-300}"

STOP_EXISTING=0
FORCE=0
REEXTRACT_DATA=0
SKIP_IMAGE_LOAD=0
SKIP_GPU_CHECK=0
SKIP_PORT_CHECK=0
SKIP_SERVICE_CHECK=0
INSTALL_COMPOSE=1
CHECK_ONLY=0
PREPARE_ONLY=0

COMPOSE_BIN=""
STACK_FILE=""
QWEN_TP_SIZE=0

SIMAI_IMAGE="${SIMAI_IMAGE:-}"
VLLM_073_IMAGE="${VLLM_073_IMAGE:-}"
VLLM_017_IMAGE="${VLLM_017_IMAGE:-}"
MILVUS_IMAGE="${MILVUS_IMAGE:-}"
ETCD_IMAGE="${ETCD_IMAGE:-quay.io/coreos/etcd:v3.5.5}"
MINIO_IMAGE="${MINIO_IMAGE:-minio/minio:RELEASE.2023-03-20T20-16-18Z}"

DOCKER=()
COMPOSE=()
QWEN_GPU_LIST=()
ALL_GPUS=()

log() {
  printf '[INFO] %s\n' "$*"
}

warn() {
  printf '[WARN] %s\n' "$*" >&2
}

die() {
  printf '[ERROR] %s\n' "$*" >&2
  exit 1
}

usage() {
  cat <<'EOF'
Usage:
  ./deploy_models_oneclick.sh [options]

Default layout:
  ./deploy_models_oneclick.sh
  ./docker-compose.tgz
  ./images/
    simai-model.base.tar.gz
    vllm-openai_v0.7.3.tar.gz
    vllm-openai_v0.17.1-amd64.tar.gz
    milvus_v2.4.24-gpu.tar.gz or milvus-gpu-simai.tar.gz
    etcd.tar.gz
    minio.tar.gz
    volumes.tgz
    Qwen3-32B.tgz

Default GPU allocation:
  simai-model: 0
  recopilot:   1
  crypto:      2
  milvus-gpu:  3
  qwen3-32b:   4,5,6,7

Options:
  --images-dir DIR          Image/model data directory. Default: ./images
  --compose-tgz FILE        docker-compose.tgz path. Default: ./docker-compose.tgz
  --compose-dir DIR         Extracted compose directory. Default: ./docker-compose
  --project-name NAME       Docker compose project name. Default: sscc_models
  --simai-gpu ID            GPU for simai-model. Default: 0
  --recopilot-gpu ID        GPU for recopilot. Default: 1
  --crypto-gpu ID           GPU for crypto-model. Default: 2
  --milvus-gpu ID           GPU for milvus standalone. Default: 3
  --qwen-gpus CSV           GPUs for Qwen3-32B. Default: 4,5,6,7
  --stop-existing           Stop known containers before checking and starting
  --force                   Ignore busy GPU and busy port checks
  --reextract-data          Re-extract volumes.tgz and Qwen3-32B.tgz even if target dirs exist
  --skip-image-load         Do not run docker load; images must already exist
  --skip-gpu-check          Skip GPU count and GPU memory checks
  --skip-port-check         Skip host port checks
  --skip-service-check      Skip post-deployment container/API checks
  --service-check-timeout S Wait timeout for post-deployment checks. Default: 1800
  --service-check-interval S Retry interval for post-deployment checks. Default: 15
  --curl-timeout S          Max seconds for each API check request. Default: 300
  --no-install-compose      Do not copy docker-compose to /usr/local/bin
  --check-only              Validate inputs/host checks only; do not load/start
  --prepare-only            Load/extract/generate compose, then stop before starting containers
  -h, --help                Show this help

Environment overrides:
  SIMAI_IMAGE, VLLM_073_IMAGE, VLLM_017_IMAGE, MILVUS_IMAGE,
  ETCD_IMAGE, MINIO_IMAGE, GPU_BUSY_THRESHOLD_MB,
  SERVICE_CHECK_TIMEOUT, SERVICE_CHECK_INTERVAL, CURL_TIMEOUT
EOF
}

parse_args() {
  while [[ $# -gt 0 ]]; do
    case "$1" in
      --images-dir)
        IMAGES_DIR="$2"
        shift 2
        ;;
      --images-dir=*)
        IMAGES_DIR="${1#*=}"
        shift
        ;;
      --compose-tgz)
        COMPOSE_TGZ="$2"
        shift 2
        ;;
      --compose-tgz=*)
        COMPOSE_TGZ="${1#*=}"
        shift
        ;;
      --compose-dir)
        COMPOSE_DIR="$2"
        shift 2
        ;;
      --compose-dir=*)
        COMPOSE_DIR="${1#*=}"
        shift
        ;;
      --project-name)
        PROJECT_NAME="$2"
        shift 2
        ;;
      --project-name=*)
        PROJECT_NAME="${1#*=}"
        shift
        ;;
      --simai-gpu)
        SIMAI_GPU="$2"
        shift 2
        ;;
      --simai-gpu=*)
        SIMAI_GPU="${1#*=}"
        shift
        ;;
      --recopilot-gpu)
        RECOPILOT_GPU="$2"
        shift 2
        ;;
      --recopilot-gpu=*)
        RECOPILOT_GPU="${1#*=}"
        shift
        ;;
      --crypto-gpu)
        CRYPTO_GPU="$2"
        shift 2
        ;;
      --crypto-gpu=*)
        CRYPTO_GPU="${1#*=}"
        shift
        ;;
      --milvus-gpu)
        MILVUS_GPU="$2"
        shift 2
        ;;
      --milvus-gpu=*)
        MILVUS_GPU="${1#*=}"
        shift
        ;;
      --qwen-gpus)
        QWEN_GPUS="$2"
        shift 2
        ;;
      --qwen-gpus=*)
        QWEN_GPUS="${1#*=}"
        shift
        ;;
      --stop-existing)
        STOP_EXISTING=1
        shift
        ;;
      --force)
        FORCE=1
        shift
        ;;
      --reextract-data)
        REEXTRACT_DATA=1
        shift
        ;;
      --skip-image-load)
        SKIP_IMAGE_LOAD=1
        shift
        ;;
      --skip-gpu-check)
        SKIP_GPU_CHECK=1
        shift
        ;;
      --skip-port-check)
        SKIP_PORT_CHECK=1
        shift
        ;;
      --skip-service-check)
        SKIP_SERVICE_CHECK=1
        shift
        ;;
      --service-check-timeout)
        SERVICE_CHECK_TIMEOUT="$2"
        shift 2
        ;;
      --service-check-timeout=*)
        SERVICE_CHECK_TIMEOUT="${1#*=}"
        shift
        ;;
      --service-check-interval)
        SERVICE_CHECK_INTERVAL="$2"
        shift 2
        ;;
      --service-check-interval=*)
        SERVICE_CHECK_INTERVAL="${1#*=}"
        shift
        ;;
      --curl-timeout)
        CURL_TIMEOUT="$2"
        shift 2
        ;;
      --curl-timeout=*)
        CURL_TIMEOUT="${1#*=}"
        shift
        ;;
      --no-install-compose)
        INSTALL_COMPOSE=0
        shift
        ;;
      --check-only)
        CHECK_ONLY=1
        shift
        ;;
      --prepare-only)
        PREPARE_ONLY=1
        shift
        ;;
      -h|--help)
        usage
        exit 0
        ;;
      *)
        die "Unknown option: $1"
        ;;
    esac
  done

  COMPOSE_BIN="$COMPOSE_DIR/docker-compose"
  STACK_FILE="$COMPOSE_DIR/compose-oneclick.yaml"
}

is_uint() {
  [[ "$1" =~ ^[0-9]+$ ]]
}

join_by_comma() {
  local IFS=,
  printf '%s' "$*"
}

parse_gpu_layout() {
  local qwen_csv="${QWEN_GPUS//[[:space:]]/}"
  [[ -n "$qwen_csv" ]] || die "QWEN_GPUS cannot be empty"

  IFS=',' read -r -a QWEN_GPU_LIST <<< "$qwen_csv"
  QWEN_TP_SIZE="${#QWEN_GPU_LIST[@]}"
  ALL_GPUS=("$SIMAI_GPU" "$RECOPILOT_GPU" "$CRYPTO_GPU" "$MILVUS_GPU" "${QWEN_GPU_LIST[@]}")
}

validate_gpu_layout() {
  parse_gpu_layout

  local id
  for id in "${ALL_GPUS[@]}"; do
    [[ -n "$id" ]] || die "GPU id cannot be empty"
    is_uint "$id" || die "GPU id must be a non-negative integer: $id"
  done

  [[ "$QWEN_TP_SIZE" -eq 4 ]] || die "Qwen3-32B requires exactly 4 GPUs, got: $QWEN_GPUS"

  local -A seen=()
  for id in "${ALL_GPUS[@]}"; do
    if [[ -n "${seen[$id]:-}" ]]; then
      die "GPU id $id is assigned more than once. Current layout: $(join_by_comma "${ALL_GPUS[@]}")"
    fi
    seen["$id"]=1
  done

  [[ "${#ALL_GPUS[@]}" -eq 8 ]] || die "This stack should reserve exactly 8 GPUs, got ${#ALL_GPUS[@]}"
}

validate_runtime_options() {
  is_uint "$SERVICE_CHECK_TIMEOUT" || die "SERVICE_CHECK_TIMEOUT must be a positive integer: $SERVICE_CHECK_TIMEOUT"
  is_uint "$SERVICE_CHECK_INTERVAL" || die "SERVICE_CHECK_INTERVAL must be a positive integer: $SERVICE_CHECK_INTERVAL"
  is_uint "$CURL_TIMEOUT" || die "CURL_TIMEOUT must be a positive integer: $CURL_TIMEOUT"

  (( SERVICE_CHECK_TIMEOUT > 0 )) || die "SERVICE_CHECK_TIMEOUT must be greater than 0"
  (( SERVICE_CHECK_INTERVAL > 0 )) || die "SERVICE_CHECK_INTERVAL must be greater than 0"
  (( CURL_TIMEOUT > 0 )) || die "CURL_TIMEOUT must be greater than 0"
}

max_required_gpu_id() {
  local max=0
  local id
  for id in "${ALL_GPUS[@]}"; do
    (( id > max )) && max="$id"
  done
  printf '%s\n' "$max"
}

require_file() {
  [[ -f "$1" ]] || die "Missing file: $1"
  [[ -r "$1" ]] || die "File is not readable: $1"
}

require_dir() {
  [[ -d "$1" ]] || die "Missing directory: $1"
  [[ -r "$1" ]] || die "Directory is not readable: $1"
}

check_payload_files() {
  require_dir "$IMAGES_DIR"

  if [[ ! -f "$COMPOSE_BIN" ]]; then
    require_file "$COMPOSE_TGZ"
  fi

  require_file "$IMAGES_DIR/volumes.tgz"
  require_file "$IMAGES_DIR/Qwen3-32B.tgz"

  if [[ "$SKIP_IMAGE_LOAD" -eq 0 ]]; then
    local required_images=(
      "simai-model.base.tar.gz"
      "vllm-openai_v0.7.3.tar.gz"
      "vllm-openai_v0.17.1-amd64.tar.gz"
      "etcd.tar.gz"
      "minio.tar.gz"
    )

    local image_file
    for image_file in "${required_images[@]}"; do
      require_file "$IMAGES_DIR/$image_file"
    done

    if [[ -f "$IMAGES_DIR/milvus_v2.4.24-gpu.tar.gz" ]]; then
      require_file "$IMAGES_DIR/milvus_v2.4.24-gpu.tar.gz"
    elif [[ -f "$IMAGES_DIR/milvus-gpu-simai.tar.gz" ]]; then
      require_file "$IMAGES_DIR/milvus-gpu-simai.tar.gz"
    else
      die "Missing Milvus image file: $IMAGES_DIR/milvus_v2.4.24-gpu.tar.gz or $IMAGES_DIR/milvus-gpu-simai.tar.gz"
    fi
  fi
}

setup_docker() {
  command -v docker >/dev/null 2>&1 || die "docker command not found"

  if docker info >/dev/null 2>&1; then
    DOCKER=(docker)
    return
  fi

  if command -v sudo >/dev/null 2>&1 && sudo docker info >/dev/null 2>&1; then
    DOCKER=(sudo docker)
    return
  fi

  die "Cannot access Docker daemon. Run as a docker user or with sudo privileges."
}

ensure_compose_payload() {
  if [[ -f "$COMPOSE_BIN" ]]; then
    return
  fi

  require_file "$COMPOSE_TGZ"
  local compose_parent
  compose_parent="$(dirname "$COMPOSE_DIR")"
  mkdir -p "$compose_parent"
  log "Extracting docker-compose payload: $COMPOSE_TGZ to $compose_parent"
  tar -xf "$COMPOSE_TGZ" -C "$compose_parent"
  [[ -f "$COMPOSE_BIN" ]] || die "docker-compose binary was not found after extracting $COMPOSE_TGZ"
}

setup_compose() {
  ensure_compose_payload
  if [[ ! -x "$COMPOSE_BIN" ]]; then
    chmod +x "$COMPOSE_BIN" 2>/dev/null || {
      if command -v sudo >/dev/null 2>&1; then
        sudo chmod +x "$COMPOSE_BIN"
      else
        die "Cannot chmod +x $COMPOSE_BIN"
      fi
    }
  fi

  if [[ "${DOCKER[0]}" == "sudo" ]]; then
    COMPOSE=(sudo "$COMPOSE_BIN")
  else
    COMPOSE=("$COMPOSE_BIN")
  fi
}

ensure_compose_dir_writable() {
  [[ -d "$COMPOSE_DIR" ]] || die "Compose directory does not exist: $COMPOSE_DIR"
  [[ -w "$COMPOSE_DIR" ]] || die "Compose directory is not writable: $COMPOSE_DIR. Run this script as the directory owner/root, or set --compose-dir to a writable extracted docker-compose directory."
}

install_compose_binary() {
  [[ "$INSTALL_COMPOSE" -eq 1 ]] || return 0

  local target="/usr/local/bin/docker-compose"
  if [[ -x "$target" ]]; then
    log "docker-compose already exists: $target"
    return
  fi

  log "Installing docker-compose to $target"
  if [[ "$(id -u)" -eq 0 ]]; then
    if cp "$COMPOSE_BIN" "$target" && chmod +x "$target"; then
      return
    fi
  elif command -v sudo >/dev/null 2>&1; then
    if sudo cp "$COMPOSE_BIN" "$target" && sudo chmod +x "$target"; then
      return
    fi
  fi

  warn "Failed to install docker-compose to $target; continuing with local binary: $COMPOSE_BIN"
}

load_image_if_present() {
  local file="$1"
  [[ -f "$file" ]] || return 0

  log "Loading docker image: $file"
  "${DOCKER[@]}" load -i "$file"
}

load_images() {
  [[ "$SKIP_IMAGE_LOAD" -eq 0 ]] || {
    log "Skipping docker image load"
    return
  }

  load_image_if_present "$IMAGES_DIR/simai-model.base.tar.gz"
  load_image_if_present "$IMAGES_DIR/vllm-openai_v0.7.3.tar.gz"
  load_image_if_present "$IMAGES_DIR/vllm-openai_v0.17.1-amd64.tar.gz"
  load_image_if_present "$IMAGES_DIR/milvus_v2.4.24-gpu.tar.gz"
  load_image_if_present "$IMAGES_DIR/milvus-gpu-simai.tar.gz"
  load_image_if_present "$IMAGES_DIR/etcd.tar.gz"
  load_image_if_present "$IMAGES_DIR/minio.tar.gz"
}

image_exists() {
  "${DOCKER[@]}" image inspect "$1" >/dev/null 2>&1
}

pick_image() {
  local var_name="$1"
  shift
  local current="${!var_name:-}"
  local candidate

  if [[ -n "$current" ]]; then
    printf -v "$var_name" '%s' "$current"
    return
  fi

  for candidate in "$@"; do
    if image_exists "$candidate"; then
      printf -v "$var_name" '%s' "$candidate"
      return
    fi
  done

  printf -v "$var_name" '%s' "$1"
}

require_image_present() {
  local image="$1"
  if ! image_exists "$image"; then
    die "Docker image is not loaded: $image. Check image tar names in $IMAGES_DIR or set the matching *_IMAGE environment variable."
  fi
}

select_and_verify_images() {
  pick_image SIMAI_IMAGE "simai-mode:base" "simai-model:base"
  pick_image VLLM_073_IMAGE "vllm/vllm-openai:v0.7.3"
  pick_image VLLM_017_IMAGE "vllm/vllm-openai:v0.17.1" "vllm/vllm-openai:v0.17.1-amd64"
  pick_image MILVUS_IMAGE "milvusdb/milvus:v2.4.24-gpu" "milvusdb/milvus:v2.4.9-gpu-simai"

  require_image_present "$SIMAI_IMAGE"
  require_image_present "$VLLM_073_IMAGE"
  require_image_present "$VLLM_017_IMAGE"
  require_image_present "$MILVUS_IMAGE"
  require_image_present "$ETCD_IMAGE"
  require_image_present "$MINIO_IMAGE"
}

dir_has_content() {
  [[ -d "$1" ]] && [[ -n "$(find "$1" -mindepth 1 -print -quit 2>/dev/null)" ]]
}

extract_runtime_data() {
  mkdir -p "$COMPOSE_DIR/model"

  if dir_has_content "$COMPOSE_DIR/volumes" && [[ "$REEXTRACT_DATA" -eq 0 ]]; then
    log "Milvus volumes already exist, skip extracting: $COMPOSE_DIR/volumes"
  else
    log "Extracting Milvus volumes to $COMPOSE_DIR"
    tar -xf "$IMAGES_DIR/volumes.tgz" -C "$COMPOSE_DIR"
  fi

  if dir_has_content "$COMPOSE_DIR/model/Qwen3-32B" && [[ "$REEXTRACT_DATA" -eq 0 ]]; then
    log "Qwen3-32B model already exists, skip extracting: $COMPOSE_DIR/model/Qwen3-32B"
  else
    log "Extracting Qwen3-32B model to $COMPOSE_DIR/model"
    tar -xf "$IMAGES_DIR/Qwen3-32B.tgz" -C "$COMPOSE_DIR/model"
  fi

  dir_has_content "$COMPOSE_DIR/volumes" || die "Milvus volumes were not found after extraction: $COMPOSE_DIR/volumes"
  dir_has_content "$COMPOSE_DIR/model/Qwen3-32B" || die "Qwen3-32B model was not found after extraction: $COMPOSE_DIR/model/Qwen3-32B"
}

gpu_yaml_list() {
  local out="["
  local id
  for id in "${QWEN_GPU_LIST[@]}"; do
    out+="\"$id\", "
  done
  out="${out%, }]"
  printf '%s\n' "$out"
}

generate_compose_file() {
  local qwen_gpu_yaml
  qwen_gpu_yaml="$(gpu_yaml_list)"

  log "Generating merged compose file: $STACK_FILE"
  cat > "$STACK_FILE" <<EOF
services:
  simai-model:
    image: "$SIMAI_IMAGE"
    container_name: simai-model
    restart: unless-stopped
    ports:
      - "27015:27015"
    environment:
      - RPC_PORT=27015
    volumes:
      - ./model/simai:/root/model
      - ./model/model_http_server.py:/root/model_http_server.py
    command: python3 /root/model_http_server.py
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: ["gpu"]
              device_ids: ["$SIMAI_GPU"]
        limits:
          cpus: "20"

  recopilot:
    image: "$VLLM_073_IMAGE"
    container_name: recopilot
    restart: unless-stopped
    ports:
      - "30856:8000"
    volumes:
      - ./model/recopilot-model:/model
    command:
      - "--model"
      - "/model"
      - "--max-model-len"
      - "16384"
      - "--served-model-name"
      - "recopilot-v0.1-alpha"
      - "--api-key"
      - "none"
      - "--dtype"
      - "half"
      - "--trust-remote-code"
      - "--enforce-eager"
      - "--disable-log-requests"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: ["gpu"]
              device_ids: ["$RECOPILOT_GPU"]
        limits:
          cpus: "20"

  crypto-model:
    image: "$VLLM_073_IMAGE"
    container_name: crypto-model
    restart: unless-stopped
    ports:
      - "8001:8000"
    volumes:
      - ./model/crypto-model:/model
    command:
      - "--model"
      - "/model"
      - "--max-model-len"
      - "16384"
      - "--served-model-name"
      - "crypto-model"
      - "--api-key"
      - "none"
      - "--dtype"
      - "half"
      - "--trust-remote-code"
      - "--enforce-eager"
      - "--disable-log-requests"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: ["gpu"]
              device_ids: ["$CRYPTO_GPU"]
        limits:
          cpus: "20"

  qwen3-32b:
    image: "$VLLM_017_IMAGE"
    container_name: qwen3-32b
    restart: unless-stopped
    ports:
      - "8000:8000"
    shm_size: "16gb"
    environment:
      - VLLM_ATTENTION_BACKEND=XFORMERS
    volumes:
      - ./model/Qwen3-32B:/model
    command:
      - "--model"
      - "/model"
      - "--tensor-parallel-size"
      - "$QWEN_TP_SIZE"
      - "--max-model-len"
      - "32768"
      - "--served-model-name"
      - "qwen3-32b"
      - "--api-key"
      - "none"
      - "--dtype"
      - "half"
      - "--trust-remote-code"
      - "--disable-log-requests"
      - "--gpu-memory-utilization"
      - "0.90"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              device_ids: $qwen_gpu_yaml
              capabilities: ["gpu"]
        limits:
          cpus: "20"

  etcd:
    image: "$ETCD_IMAGE"
    container_name: milvus-etcd
    restart: unless-stopped
    environment:
      - ETCD_AUTO_COMPACTION_MODE=revision
      - ETCD_AUTO_COMPACTION_RETENTION=1000
      - ETCD_QUOTA_BACKEND_BYTES=4294967296
      - ETCD_SNAPSHOT_COUNT=50000
    volumes:
      - ./volumes/etcd:/etcd
    command: etcd -advertise-client-urls=http://127.0.0.1:2379 -listen-client-urls http://0.0.0.0:2379 --data-dir /etcd
    healthcheck:
      test: ["CMD", "etcdctl", "endpoint", "health"]
      interval: 30s
      timeout: 20s
      retries: 3

  minio:
    image: "$MINIO_IMAGE"
    container_name: milvus-minio
    restart: unless-stopped
    environment:
      MINIO_ACCESS_KEY: "vfp"
      MINIO_SECRET_KEY: "vfpAllUser@)@@"
    ports:
      - "9001:9001"
      - "9002:9000"
    volumes:
      - ./volumes/minio:/minio_data
    command: minio server /minio_data --console-address ":9001"
    healthcheck:
      test: ["CMD", "curl", "-f", "http://localhost:9000/minio/health/live"]
      interval: 30s
      timeout: 20s
      retries: 3

  milvus-standalone:
    image: "$MILVUS_IMAGE"
    container_name: milvus-standalone
    restart: unless-stopped
    command: ["milvus", "run", "standalone"]
    security_opt:
      - seccomp:unconfined
    environment:
      ETCD_ENDPOINTS: etcd:2379
      MINIO_ADDRESS: minio:9000
    volumes:
      - ./volumes/milvus:/var/lib/milvus
    ports:
      - "19530:19530"
      - "9091:9091"
    deploy:
      resources:
        reservations:
          devices:
            - driver: nvidia
              capabilities: ["gpu"]
              device_ids: ["$MILVUS_GPU"]
        limits:
          cpus: "20"
    depends_on:
      - etcd
      - minio

networks:
  default:
    name: sscc-models
EOF
}

check_gpu_host() {
  [[ "$SKIP_GPU_CHECK" -eq 0 ]] || {
    log "Skipping GPU check"
    return
  }

  command -v nvidia-smi >/dev/null 2>&1 || die "nvidia-smi not found; NVIDIA driver is not ready"

  local gpu_count
  gpu_count="$(nvidia-smi -L | awk '/^GPU [0-9]+:/{count++} END{print count+0}')"
  local max_id
  max_id="$(max_required_gpu_id)"
  if (( gpu_count <= max_id )); then
    die "Not enough GPUs. Required ids: $(join_by_comma "${ALL_GPUS[@]}"), detected GPU count: $gpu_count"
  fi

  local busy_lines=()
  local idx mem
  while IFS=',' read -r idx mem; do
    idx="${idx//[[:space:]]/}"
    mem="${mem//[[:space:]]/}"
    [[ -n "$idx" && -n "$mem" ]] || continue

    local required=0
    local id
    for id in "${ALL_GPUS[@]}"; do
      if [[ "$id" == "$idx" ]]; then
        required=1
        break
      fi
    done

    if [[ "$required" -eq 1 && "$mem" =~ ^[0-9]+$ && "$mem" -gt "$GPU_BUSY_THRESHOLD_MB" ]]; then
      busy_lines+=("GPU $idx uses ${mem}MiB")
    fi
  done < <(nvidia-smi --query-gpu=index,memory.used --format=csv,noheader,nounits)

  if [[ "${#busy_lines[@]}" -gt 0 ]]; then
    printf '[WARN] Required GPUs look busy:\n' >&2
    printf '  %s\n' "${busy_lines[@]}" >&2
    if [[ "$FORCE" -eq 0 ]]; then
      die "Stop existing GPU workloads, or rerun with --stop-existing for known containers, or --force to ignore this check."
    fi
  fi
}

port_in_use() {
  local port="$1"
  if command -v ss >/dev/null 2>&1; then
    ss -ltn | awk '{print $4}' | grep -Eq "[:.]${port}$"
  elif command -v netstat >/dev/null 2>&1; then
    netstat -ltn | awk '{print $4}' | grep -Eq "[:.]${port}$"
  else
    return 1
  fi
}

check_ports() {
  [[ "$SKIP_PORT_CHECK" -eq 0 ]] || {
    log "Skipping port check"
    return
  }

  local ports=(27015 30856 8001 8000 19530 9091 9001 9002)
  local busy_ports=()
  local port
  for port in "${ports[@]}"; do
    if port_in_use "$port"; then
      busy_ports+=("$port")
    fi
  done

  if [[ "${#busy_ports[@]}" -gt 0 ]]; then
    warn "Host ports are already in use: $(join_by_comma "${busy_ports[@]}")"
    if [[ "$FORCE" -eq 0 ]]; then
      die "Free these ports or rerun with --force if they are occupied by the existing target stack."
    fi
  fi
}

container_exists() {
  local name="$1"
  "${DOCKER[@]}" ps -a --format '{{.Names}}' | grep -Fxq "$name"
}

stop_existing_containers() {
  [[ "$STOP_EXISTING" -eq 1 ]] || return 0

  log "Stopping known model/vector containers"
  if [[ -f "$STACK_FILE" && -x "$COMPOSE_BIN" ]]; then
    (
      cd "$COMPOSE_DIR"
      "${COMPOSE[@]}" -p "$PROJECT_NAME" -f "$STACK_FILE" down --remove-orphans
    ) || warn "compose down failed; continuing with direct container cleanup"
  fi

  local names=(
    simai-model
    recopilot
    crypto-model
    qwen3-32b
    milvus-standalone
    milvus-etcd
    milvus-minio
  )
  local name
  for name in "${names[@]}"; do
    if container_exists "$name"; then
      "${DOCKER[@]}" rm -f "$name"
    fi
  done
}

start_stack() {
  log "Starting all services with compose project: $PROJECT_NAME"
  (
    cd "$COMPOSE_DIR"
    "${COMPOSE[@]}" -p "$PROJECT_NAME" -f "$STACK_FILE" up -d
  )
}

container_state() {
  local name="$1"
  "${DOCKER[@]}" inspect -f '{{.State.Status}}' "$name" 2>/dev/null || true
}

container_health() {
  local name="$1"
  "${DOCKER[@]}" inspect -f '{{if .State.Health}}{{.State.Health.Status}}{{else}}none{{end}}' "$name" 2>/dev/null || true
}

print_container_debug() {
  local name="$1"
  warn "Container status for $name:"
  "${DOCKER[@]}" inspect -f '  name={{.Name}} status={{.State.Status}} health={{if .State.Health}}{{.State.Health.Status}}{{else}}none{{end}} exit={{.State.ExitCode}} error={{.State.Error}}' "$name" 2>/dev/null || true
  warn "Last logs for $name:"
  "${DOCKER[@]}" logs --tail 80 "$name" 2>&1 || true
}

show_stack_status() {
  log "Current compose status:"
  (
    cd "$COMPOSE_DIR"
    "${COMPOSE[@]}" -p "$PROJECT_NAME" -f "$STACK_FILE" ps
  ) || true
}

wait_for_containers_ready() {
  local containers=(
    simai-model
    recopilot
    crypto-model
    qwen3-32b
    milvus-etcd
    milvus-minio
    milvus-standalone
  )
  local deadline=$((SECONDS + SERVICE_CHECK_TIMEOUT))
  local name state health ready failed_name

  log "Waiting for containers to be running/healthy"
  while true; do
    ready=1
    failed_name=""
    for name in "${containers[@]}"; do
      state="$(container_state "$name")"
      health="$(container_health "$name")"
      if [[ "$state" != "running" ]]; then
        ready=0
        failed_name="$name"
        break
      fi
      if [[ "$health" != "none" && "$health" != "healthy" ]]; then
        ready=0
        failed_name="$name"
        break
      fi
    done

    if [[ "$ready" -eq 1 ]]; then
      log "All containers are running/healthy"
      return 0
    fi

    if (( SECONDS >= deadline )); then
      warn "Container readiness timed out"
      for name in "${containers[@]}"; do
        state="$(container_state "$name")"
        health="$(container_health "$name")"
        warn "$name status=${state:-missing} health=${health:-missing}"
      done
      print_container_debug "${failed_name:-${containers[0]}}"
      return 1
    fi

    log "Waiting for ${failed_name:-unknown} status=${state:-missing} health=${health:-missing}; retrying in ${SERVICE_CHECK_INTERVAL}s"
    sleep "$SERVICE_CHECK_INTERVAL"
  done
}

tcp_port_open() {
  local host="$1"
  local port="$2"

  if command -v nc >/dev/null 2>&1; then
    nc -z -w 5 "$host" "$port" >/dev/null 2>&1
  elif command -v timeout >/dev/null 2>&1; then
    timeout 5 bash -c "</dev/tcp/$host/$port" >/dev/null 2>&1
  else
    bash -c "</dev/tcp/$host/$port" >/dev/null 2>&1
  fi
}

wait_for_check() {
  local label="$1"
  local container="$2"
  local check_func="$3"
  local deadline=$((SECONDS + SERVICE_CHECK_TIMEOUT))

  log "Checking $label"
  while true; do
    if "$check_func"; then
      log "$label OK"
      return 0
    fi

    if (( SECONDS >= deadline )); then
      warn "$label check timed out"
      print_container_debug "$container"
      return 1
    fi

    log "$label not ready; retrying in ${SERVICE_CHECK_INTERVAL}s"
    sleep "$SERVICE_CHECK_INTERVAL"
  done
}

check_simai_api() {
  curl -fsS --max-time "$CURL_TIMEOUT" \
    -X POST "http://127.0.0.1:27015/v1/generate-embeddings" \
    -H "Content-Type: application/json" \
    -d '{"features":["function calculate_sum(a, b) {\n  return a + b;\n}"]}' \
    >/dev/null
}

check_recopilot_api() {
  curl -fsS --max-time "$CURL_TIMEOUT" \
    -X POST "http://127.0.0.1:30856/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer none" \
    -d '{"model":"recopilot-v0.1-alpha","temperature":0.1,"stream":false,"max_tokens":16,"messages":[{"role":"user","content":"<pseudocode>function calculate_sum(a, b) {\\n  return a + b;\\n}</pseudocode><sourcecode>"}]}' \
    >/dev/null
}

check_crypto_api() {
  curl -fsS --max-time "$CURL_TIMEOUT" \
    -X POST "http://127.0.0.1:8001/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer none" \
    -d '{"model":"crypto-model","temperature":0.1,"stream":false,"max_tokens":16,"messages":[{"role":"user","content":"ASM:\\npush ebx\\nmov eax, 1\\nretn\\nPCODE:\\nint sample(){ return 1; }\\n<detect-crypto-feature>"}]}' \
    >/dev/null
}

check_qwen_api() {
  curl -fsS --max-time "$CURL_TIMEOUT" \
    -X POST "http://127.0.0.1:8000/v1/chat/completions" \
    -H "Content-Type: application/json" \
    -H "Authorization: Bearer none" \
    -d '{"model":"qwen3-32b","temperature":0.1,"top_p":0.8,"stream":false,"max_tokens":32,"messages":[{"role":"user","content":"Hello, introduce yourself briefly."}]}' \
    >/dev/null
}

check_milvus_port() {
  tcp_port_open "127.0.0.1" "19530"
}

run_service_checks() {
  [[ "$SKIP_SERVICE_CHECK" -eq 0 ]] || {
    log "Skipping post-deployment service checks"
    return
  }

  command -v curl >/dev/null 2>&1 || die "curl command not found; cannot run post-deployment API checks"

  wait_for_containers_ready
  wait_for_check "Milvus port 19530" "milvus-standalone" check_milvus_port
  wait_for_check "SimAi embedding API" "simai-model" check_simai_api
  wait_for_check "recopilot chat API" "recopilot" check_recopilot_api
  wait_for_check "crypto-model chat API" "crypto-model" check_crypto_api
  wait_for_check "Qwen3-32B chat API" "qwen3-32b" check_qwen_api
  show_stack_status
}

print_summary() {
  cat <<EOF

Deployment files:
  compose directory: $COMPOSE_DIR
  merged compose:    $STACK_FILE

GPU allocation:
  simai-model:       $SIMAI_GPU
  recopilot:         $RECOPILOT_GPU
  crypto-model:      $CRYPTO_GPU
  milvus-standalone: $MILVUS_GPU
  qwen3-32b:         $QWEN_GPUS

Endpoints:
  simai-model:       http://127.0.0.1:27015/v1/generate-embeddings
  recopilot:         http://127.0.0.1:30856/v1/chat/completions
  crypto-model:      http://127.0.0.1:8001/v1/chat/completions
  qwen3-32b:         http://127.0.0.1:8000/v1/chat/completions
  milvus:            127.0.0.1:19530
EOF
}

main() {
  parse_args "$@"
  validate_gpu_layout
  validate_runtime_options
  check_payload_files
  setup_docker

  if [[ "$CHECK_ONLY" -eq 1 ]]; then
    check_gpu_host
    check_ports
    log "Check-only passed. No images loaded and no containers started."
    print_summary
    return
  fi

  setup_compose
  install_compose_binary
  ensure_compose_dir_writable
  stop_existing_containers
  check_gpu_host
  check_ports
  load_images
  select_and_verify_images
  extract_runtime_data
  generate_compose_file
  if [[ "$PREPARE_ONLY" -eq 1 ]]; then
    log "Prepare-only completed. Containers were not started."
    print_summary
    return
  fi
  start_stack
  run_service_checks
  print_summary
}

main "$@"