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

workdir="$(mktemp -d)"
trap 'rm -rf "$workdir"' EXIT

cat > "$workdir/sglangmux.yml" <<'YAML'
cuda_version: 126
tensor_parallel_size: 1
shenron_version: latest
engine: sglang
gateway_upstream_port: 3000
prometheus_port: 9090
prometheus_version: v2.51.2
scouter_version: latest
scouter_collector_instance: host.docker.internal
scouter_reporter_interval: 10
scouter_ingest_api_key: api-key
sglangmux_listen_port: 8100
sglangmux_host: 0.0.0.0
sglangmux_upstream_timeout_secs: 120
sglangmux_model_ready_timeout_secs: 600
sglangmux_model_switch_timeout_secs: 90
models:
- model_name: Qwen/Qwen3-0.6B
  api_key: sk-model-a
  engine_port: 8001
  engine_envs:
  - VLLM_ENABLE_RESPONSES_API_STORE
  - -1
  engine_args:
  - --tp
  - 1
- model_name: Qwen/Qwen3-30B-A3B
  api_key: sk-model-b
  engine_port: 8002
  engine_use_cuda_ipc_transport: true
  engine_args:
  - --tp
  - 2
YAML

shenron "$workdir/sglangmux.yml" --output-dir "$workdir"

required=(
  "$workdir/docker-compose.yml"
  "$workdir/.generated/Caddyfile"
  "$workdir/.generated/prometheus.yml"
  "$workdir/.generated/scouter_reporter.env"
  "$workdir/.generated/engine_start_1.sh"
  "$workdir/.generated/engine_start_2.sh"
  "$workdir/.generated/sglangmux_start.sh"
)

for f in "${required[@]}"; do
  if [ ! -f "$f" ]; then
    echo "expected generated file not found: $f" >&2
    exit 1
  fi
done

mux_script="$workdir/.generated/sglangmux_start.sh"
caddyfile="$workdir/.generated/Caddyfile"
compose="$workdir/docker-compose.yml"
engine1="$workdir/.generated/engine_start_1.sh"

if ! grep -q "sglangmux" "$mux_script"; then
  echo "sglangmux_start.sh missing sglangmux command" >&2
  exit 1
fi

if ! grep -q -- "--listen-port" "$mux_script"; then
  echo "sglangmux_start.sh missing --listen-port" >&2
  exit 1
fi

if ! grep -q "VLLM_ENABLE_RESPONSES_API_STORE=-1" "$engine1"; then
  echo "engine_start_1.sh missing custom env export" >&2
  exit 1
fi

for engine_script in /generated/engine_start_1.sh /generated/engine_start_2.sh; do
  if ! grep -q "$engine_script" "$mux_script"; then
    echo "sglangmux_start.sh missing script argument: $engine_script" >&2
    exit 1
  fi
done

if ! grep -q "reverse_proxy vllm:8100" "$caddyfile"; then
  echo "Caddyfile missing reverse_proxy to sglangmux upstream (vllm:8100)" >&2
  exit 1
fi

if ! grep -q "Bearer sk-" "$caddyfile"; then
  echo "Caddyfile missing auth check" >&2
  exit 1
fi

if ! grep -q "caddy:2-alpine" "$compose"; then
  echo "docker-compose.yml missing caddy:2-alpine" >&2
  exit 1
fi

if ! grep -q "/generated/sglangmux_start.sh" "$compose"; then
  echo "docker-compose.yml missing mux startup command" >&2
  exit 1
fi

echo "sglangmux generation produced expected files and wiring"
