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

if [[ $# -ne 5 ]]; then
  echo "usage: $0 TR2_OR_TR3 ATTEMPT_NAME API_PORT TAIL_DIR MAP_DIR" >&2
  exit 2
fi

TIER="${1^^}"
ATTEMPT_NAME="$2"
API_PORT="$3"
TAIL_DIR="$4"
MAP_DIR="$5"

if [[ "$TIER" != "TR2" && "$TIER" != "TR3" ]]; then
  echo "tier must be TR2 or TR3" >&2
  exit 2
fi
case "$API_PORT" in
  ''|*[!0-9]*) echo "API port must be an integer" >&2; exit 2 ;;
esac

IMAGE="${LAGUNA_IMAGE:-laguna-vllm-hybrid:20260725-a10-tp1}"
MODEL_DIR="${LAGUNA_MODEL_DIR:-/home/sero/models/Laguna-S-2.1-NVFP4}"
MAP_FILE="${LAGUNA_MAP_FILE:-laguna-${TIER,,}-tier-map.json}"
SERVED_MODEL="${LAGUNA_SERVED_MODEL:-laguna-s-2.1-${TIER,,}-hybrid}"
MAX_NUM_SEQS="${LAGUNA_MAX_NUM_SEQS:-4}"
MAX_BATCHED_TOKENS="${LAGUNA_MAX_BATCHED_TOKENS:-4096}"
GPU_MEMORY_UTILIZATION="${LAGUNA_GPU_MEMORY_UTILIZATION:-0.86}"

for value in "$MAX_NUM_SEQS" "$MAX_BATCHED_TOKENS"; do
  case "$value" in
    ''|*[!0-9]*) echo "sequence and batching settings must be integers" >&2; exit 2 ;;
    0) echo "sequence and batching settings must be positive" >&2; exit 2 ;;
  esac
done

docker inspect "$ATTEMPT_NAME" >/dev/null 2>&1 \
  && { echo "container already exists: $ATTEMPT_NAME" >&2; exit 3; }
docker image inspect "$IMAGE" >/dev/null
[[ -d "$MODEL_DIR" ]] || { echo "missing source model: $MODEL_DIR" >&2; exit 3; }
[[ -d "$TAIL_DIR/${TIER,,}" ]] \
  || { echo "missing tail tier: $TAIL_DIR/${TIER,,}" >&2; exit 3; }
[[ -s "$MAP_DIR/$MAP_FILE" ]] \
  || { echo "missing tier map: $MAP_DIR/$MAP_FILE" >&2; exit 3; }

docker run -d \
  --name "$ATTEMPT_NAME" \
  --gpus all \
  --network host \
  --ipc host \
  --shm-size 64g \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -v "$MODEL_DIR:/model:ro" \
  -v "$TAIL_DIR:/tail:ro" \
  -v "$MAP_DIR:/maps:ro" \
  -e "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True" \
  -e "TORCH_CUDA_ARCH_LIST=12.1a" \
  -e "CUTE_DSL_ARCH=sm_121a" \
  -e "MAX_JOBS=4" \
  -e "VLLM_NO_USAGE_STATS=1" \
  -e "LAGUNA_HYBRID_TIER=$TIER" \
  -e "LAGUNA_HYBRID_MODEL_DIR=/model" \
  -e "LAGUNA_HYBRID_TAIL_DIR=/tail" \
  -e "LAGUNA_HYBRID_TIER_MAP=/maps/$MAP_FILE" \
  -e "LAGUNA_HYBRID_TRELLIS_CHUNK=128" \
  -e "PYTHONPATH=/usr/local/lib/python3.12/dist-packages" \
  "$IMAGE" \
  /model \
  --served-model-name "$SERVED_MODEL" \
  --enable-auto-tool-choice \
  --tool-call-parser poolside_v1 \
  --reasoning-parser poolside_v1 \
  --structured-outputs-config '{"enable_in_reasoning":true}' \
  --compilation-config '{"inductor_compile_config":{"benchmark_combo_kernel":false}}' \
  --override-generation-config '{"temperature":0.7,"top_p":0.95}' \
  --max-num-seqs "$MAX_NUM_SEQS" \
  --max-num-batched-tokens "$MAX_BATCHED_TOKENS" \
  --max-model-len 200000 \
  --gpu-memory-utilization "$GPU_MEMORY_UTILIZATION" \
  --tensor-parallel-size 1 \
  --moe-backend flashinfer_cutlass \
  --attention-backend TRITON_ATTN \
  --host 0.0.0.0 \
  --port "$API_PORT"

printf '%s\n' "$ATTEMPT_NAME"