#!/bin/bash
# Common environment for all connector_tax launch scripts.
#
# Usage: source common.sh
# Sets: MODEL_PATH, PYTHON, PORT, LOG_DIR, COMMON_VLLM_ARGS

set -euo pipefail

# Resolve project root (microbench/connector_tax/launch/common.sh → ../../..)
SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
CT_DIR="$(cd "$SCRIPT_DIR/.." && pwd)"          # microbench/connector_tax
MB_DIR="$(cd "$CT_DIR/.."   && pwd)"            # microbench
PROJ_DIR="$(cd "$MB_DIR/.." && pwd)"            # repo root

VENV="$PROJ_DIR/.venv/bin"
PYTHON="${VENV}/python"
[[ -x "$PYTHON" ]] || PYTHON="$(command -v python3)"

MODEL_PATH="${MODEL_PATH:-$HOME/models/Qwen/Qwen3-Coder-30B-A3B-Instruct}"
PORT="${PORT:-8000}"
GPU_ID="${GPU_ID:-0}"

# Per-run log dir (caller may override RUN_DIR)
RUN_DIR="${RUN_DIR:-$CT_DIR/results/_default}"
mkdir -p "$RUN_DIR"
LOG_DIR="$RUN_DIR"

# PYTHONPATH so `microbench.connector_tax.tools.noop_connector:NoOpConnector`
# resolves when launched from anywhere.
export PYTHONPATH="${PROJ_DIR}:${PYTHONPATH:-}"

# Step-timing log path (vLLM's existing scheduler emits when this is set;
# our patch enriches the record with timing fields)
export AGENTIC_STEP_LOG_PATH="${AGENTIC_STEP_LOG_PATH:-$LOG_DIR/engine_step.jsonl}"

# Common vLLM flags shared by all configs
COMMON_VLLM_ARGS=(
    --model "$MODEL_PATH"
    --host 0.0.0.0
    --port "$PORT"
    --tensor-parallel-size 1
    --trust-remote-code
    --enable-prefix-caching
    --dtype auto
    --gpu-memory-utilization 0.9
    --max-model-len 200000
    --no-enable-log-requests
)

# Wait until /v1/models returns 200, then return.
# Times out at $1 seconds (default 240).
wait_for_ready() {
    local timeout="${1:-240}"
    local i
    for i in $(seq 1 "$timeout"); do
        if curl -sf "http://127.0.0.1:$PORT/v1/models" >/dev/null 2>&1; then
            echo "Server ready after ${i}s"
            return 0
        fi
        sleep 1
    done
    echo "ERROR: Server did not become ready within ${timeout}s" >&2
    return 1
}