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

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

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

if [[ "$TIER" != "TR2" && "$TIER" != "TR3" ]]; then
  echo "tier must be TR2 or TR3" >&2
  exit 2
fi
SPECULATIVE_MODE="${LAGUNA_SPECULATIVE_MODE:-none}"
if [[ "$SPECULATIVE_MODE" != "dflash" && "$SPECULATIVE_MODE" != "none" ]]; then
  echo "LAGUNA_SPECULATIVE_MODE must be dflash or none" >&2
  exit 2
fi
for value in "$MASTER_PORT" "$API_PORT"; do
  case "$value" in
    ''|*[!0-9]*) echo "ports must be integers" >&2; exit 2 ;;
  esac
done

IMAGE="laguna-vllm-hybrid:20260724-a9-hot-remap"
MODEL_DIR="${LAGUNA_MODEL_DIR:-/home/sero/models/Laguna-S-2.1-NVFP4}"
DFLASH_DIR="${LAGUNA_DFLASH_DIR:-/home/sero/models/Laguna-S-2.1-DFlash-NVFP4}"
MAP_DIR="${LAGUNA_MAP_DIR:-/home/sero/laguna-reap-saliency-v1/final-20260723-a2/tier-maps}"
WORKER_HOST="10.0.0.2"
HEAD_IP="10.0.0.1"
SERVED_MODEL="laguna-s-2.1-${TIER,,}-hybrid"
MAP_FILE="${LAGUNA_MAP_FILE:-laguna-${TIER,,}-tier-map.json}"

for host in head worker; do
  if [[ "$host" == "head" ]]; then
    docker inspect "$ATTEMPT_NAME" >/dev/null 2>&1 \
      && { echo "head container already exists: $ATTEMPT_NAME" >&2; exit 3; }
  else
    ssh -o BatchMode=yes -o ConnectTimeout=10 "$WORKER_HOST" \
      "docker inspect '$ATTEMPT_NAME' >/dev/null 2>&1" \
      && { echo "worker container already exists: $ATTEMPT_NAME" >&2; exit 3; }
  fi
done
[[ -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; }
if [[ "$SPECULATIVE_MODE" == "dflash" ]]; then
  [[ -s "$DFLASH_DIR/model.safetensors" ]] \
    || { echo "missing pinned Laguna DFlash checkpoint" >&2; exit 3; }
  ssh -o BatchMode=yes -o ConnectTimeout=10 "$WORKER_HOST" \
    "test -s '$DFLASH_DIR/model.safetensors'" \
    || { echo "worker is missing pinned Laguna DFlash checkpoint" >&2; exit 3; }
fi
DFLASH_MOUNT=()
if [[ "$SPECULATIVE_MODE" == "dflash" ]]; then
  DFLASH_MOUNT=(-v "$DFLASH_DIR:/dflash:ro")
fi

cleanup_failed_launch() {
  docker stop "$ATTEMPT_NAME" >/dev/null 2>&1 || true
  ssh -o BatchMode=yes -o ConnectTimeout=10 "$WORKER_HOST" \
    "docker stop '$ATTEMPT_NAME' >/dev/null 2>&1 || true" || true
}
trap cleanup_failed_launch ERR

COMMON_ENV=(
  -e "NCCL_IB_DISABLE=0"
  -e "NCCL_IB_HCA=rocep1s0f1,roceP2p1s0f1"
  -e "NCCL_IB_GID_INDEX=0"
  -e "NCCL_CUMEM_ENABLE=0"
  -e "NCCL_DEBUG=WARN"
  -e "TORCH_CUDA_ARCH_LIST=12.1a"
  -e MASTER_ADDR="$HEAD_IP"
  -e MASTER_PORT="$MASTER_PORT"
  -e "GLOO_SOCKET_IFNAME=enp1s0f1np1"
  -e "NCCL_IB_ROCE_VERSION_NUM=2"
  -e "MN_IF_NAME=enp1s0f1np1"
  -e "PYTORCH_CUDA_ALLOC_CONF=expandable_segments:True"
  -e "MAX_JOBS=4"
  -e "NCCL_SOCKET_IFNAME=enp1s0f1np1"
  -e "NCCL_NET=IB"
  -e "NCCL_CROSS_NIC=1"
  -e "NCCL_IGNORE_CPU_AFFINITY=1"
  -e "OMPI_MCA_btl_tcp_if_include=enp1s0f1np1"
  -e "CUTE_DSL_ARCH=sm_121a"
  -e "NCCL_IB_ADDR_FAMILY=AF_INET"
  -e "NCCL_NVLS_ENABLE=0"
  -e "TP_SOCKET_IFNAME=enp1s0f1np1"
  -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"
)

COMMON_ARGS=(
  /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}'
  --override-generation-config '{"temperature":0.7,"top_p":0.95}'
  --max-num-seqs 4
  --max-num-batched-tokens 4096
  --max-model-len 200000
  --gpu-memory-utilization 0.86
  --tensor-parallel-size 2
  --distributed-executor-backend mp
  --moe-backend flashinfer_cutlass
  --attention-backend TRITON_ATTN
  --host 0.0.0.0
  --port "$API_PORT"
  --nnodes 2
  --master-addr "$HEAD_IP"
  --master-port "$MASTER_PORT"
)
if [[ "$SPECULATIVE_MODE" == "dflash" ]]; then
  COMMON_ARGS+=(
    --speculative-config '{"model":"/dflash","num_speculative_tokens":15,"method":"dflash"}'
  )
fi

printf -v WORKER_ENV ' %q' "${COMMON_ENV[@]}"
printf -v WORKER_ARGS ' %q' "${COMMON_ARGS[@]}"
WORKER_DFLASH_MOUNT=""
if ((${#DFLASH_MOUNT[@]})); then
  printf -v WORKER_DFLASH_MOUNT ' %q' "${DFLASH_MOUNT[@]}"
fi
ssh -o BatchMode=yes -o ConnectTimeout=10 "$WORKER_HOST" \
  "docker run -d --name '$ATTEMPT_NAME' --gpus all --network host --ipc host --shm-size 64g --device /dev/infiniband --ulimit memlock=-1 --ulimit stack=67108864 -v '$MODEL_DIR:/model:ro'$WORKER_DFLASH_MOUNT -v '$TAIL_DIR:/tail:ro' -v '$MAP_DIR:/maps:ro'$WORKER_ENV -e VLLM_HOST_IP=10.0.0.2 '$IMAGE'$WORKER_ARGS --node-rank 1 --headless"

docker run -d \
  --name "$ATTEMPT_NAME" \
  --gpus all \
  --network host \
  --ipc host \
  --shm-size 64g \
  --device /dev/infiniband \
  --ulimit memlock=-1 \
  --ulimit stack=67108864 \
  -v "$MODEL_DIR:/model:ro" \
  "${DFLASH_MOUNT[@]}" \
  -v "$TAIL_DIR:/tail:ro" \
  -v "$MAP_DIR:/maps:ro" \
  "${COMMON_ENV[@]}" \
  -e VLLM_HOST_IP=10.0.0.1 \
  "$IMAGE" \
  "${COMMON_ARGS[@]}" \
  --node-rank 0

trap - ERR
printf '%s\n' "$ATTEMPT_NAME"