#!/usr/bin/env bash
# Run a single profile benchmark on P800 SGLang with PyTorch profiler.
# ISL=4096, OSL=1024, concurrency=16, TP=8, DP=1
set -Eeuo pipefail

SCRIPT_DIR="$(cd "$(dirname "${BASH_SOURCE[0]}")" && pwd)"
source "${SCRIPT_DIR}/../../../scripts/common/lib.sh"
source "${SCRIPT_DIR}/../../../scripts/common/platform.sh"
source "${SCRIPT_DIR}/config.env"

TIMESTAMP="$(date '+%Y%m%d-%H%M%S')"

# Override for this run
PROFILE_OUTPUT_DIR="${SCRIPT_DIR}/profile_output/${TIMESTAMP}"
BENCH_OUTPUT_DIR="${SCRIPT_DIR}/bench_output/${TIMESTAMP}"
mkdir -p "${PROFILE_OUTPUT_DIR}"
mkdir -p "${BENCH_OUTPUT_DIR}"

# ---------------------------------------------------------------------------
# Step 0: Kill any existing container
# ---------------------------------------------------------------------------
log "Stopping any existing container: ${CONTAINER_NAME}"
docker rm -f "${CONTAINER_NAME}" 2>/dev/null || true
sleep 2

# ---------------------------------------------------------------------------
# Step 1: Start the server (TP=8, DP=1)
# ---------------------------------------------------------------------------
log "Starting server (TP=8, DP=1)..."
bash "${SCRIPT_DIR}/start_sglang_docker.sh" 8 1
log "Server is healthy"

# ---------------------------------------------------------------------------
# Step 2: Start XPU monitoring in background
# ---------------------------------------------------------------------------
XPU_CSV="${PROFILE_OUTPUT_DIR}/xpu_monitor.csv"
log "Starting XPU monitor -> ${XPU_CSV}"
(
  echo "timestamp,index,memory.used [MiB],memory.total [MiB],utilization.gpu [%]"
  while true; do
    xpu-smi 2>/dev/null | python3 -c "
import sys, csv, io
data = sys.stdin.read()
lines = data.strip().split('\n')
ts = __import__('datetime').datetime.now().strftime('%Y/%m/%d %H:%M:%S')
for line in lines:
    if 'MiB' in line and '%' in line:
        parts = line.split('|')
        idx = parts[1].strip() if len(parts) > 1 else ''
        mem = parts[3].strip() if len(parts) > 3 else ''
        util = parts[4].strip() if len(parts) > 4 else ''
        if '/' in mem:
            u = mem.split('/')[0].strip().replace('MiB','').strip()
            t = mem.split('/')[1].strip().replace('MiB','').strip()
            p = util.replace('%','').strip()
            print(f'{ts}, {idx}, {u} MiB, {t} MiB, {p} %')
    "
    sleep 1
  done
) >> "${XPU_CSV}" 2>/dev/null &
XPU_MONITOR_PID=$!
log "XPU monitor PID=${XPU_MONITOR_PID}"

# ---------------------------------------------------------------------------
# Step 3: Run bench_serving with --profile
# ---------------------------------------------------------------------------
log "Starting bench_serving with profiler..."
log "  ISL=4096, OSL=1024, concurrency=16, num_prompts=80"
log "  Profile steps=30, output=${PROFILE_OUTPUT_DIR}"

BENCH_OUTPUT_FILE="/workspace/bench_output/results_${TIMESTAMP}.jsonl"
PROFILE_DIR_CONTAINER="/workspace/profile_output"

set +e
docker exec "${CONTAINER_NAME}" \
  env HF_HUB_OFFLINE=1 TRANSFORMERS_OFFLINE=1 HF_DATASETS_OFFLINE=1 \
  "${CONTAINER_PYTHON}" -m sglang.bench_serving \
    --backend sglang \
    --host 127.0.0.1 \
    --port "${SGLANG_PORT}" \
    --model "${MODEL_PATH}" \
    --dataset-name random \
    --dataset-path "${CONTAINER_DATASET_PATH}" \
    --random-input-len 4096 \
    --random-output-len 1024 \
    --random-range-ratio 1.0 \
    --num-prompts 80 \
    --max-concurrency 16 \
    --request-rate 10000 \
    --output-file "${BENCH_OUTPUT_FILE}" \
    --output-details \
    --profile \
    --profile-steps 30 \
    --profile-output-dir "${PROFILE_DIR_CONTAINER}" \
    --profile-prefix "p800_4k1k_c16" \
    2>&1 | tee "${PROFILE_OUTPUT_DIR}/bench_serving.log"
BENCH_EXIT_CODE=$?
set -e

log "bench_serving exit code: ${BENCH_EXIT_CODE}"

# ---------------------------------------------------------------------------
# Step 4: Stop XPU monitor
# ---------------------------------------------------------------------------
kill "${XPU_MONITOR_PID}" 2>/dev/null || true
wait "${XPU_MONITOR_PID}" 2>/dev/null || true
log "XPU monitor stopped"

# ---------------------------------------------------------------------------
# Step 5: Copy profile output files from container
# ---------------------------------------------------------------------------
log "Profile output files:"
docker exec "${CONTAINER_NAME}" find /tmp -name "*.json" -path "*profile*" 2>/dev/null \
  | while read -r f; do
    log "  Found in container: ${f}"
    docker cp "${CONTAINER_NAME}:${f}" "${PROFILE_OUTPUT_DIR}/" 2>/dev/null || true
  done

docker exec "${CONTAINER_NAME}" ls -la "/workspace/profile_output/" 2>/dev/null
docker exec "${CONTAINER_NAME}" ls -la "/workspace/bench_output/" 2>/dev/null

# Copy bench output
if docker exec "${CONTAINER_NAME}" test -f "${BENCH_OUTPUT_FILE}" 2>/dev/null; then
  docker cp "${CONTAINER_NAME}:${BENCH_OUTPUT_FILE}" "${BENCH_OUTPUT_DIR}/" 2>/dev/null
  log "Bench output copied"
fi

# ---------------------------------------------------------------------------
# Step 6: List collected files
# ---------------------------------------------------------------------------
echo ""
echo "============================================"
echo "Collected files:"
echo "============================================"
find "${PROFILE_OUTPUT_DIR}" -type f 2>/dev/null || echo "(no profile files)"
find "${BENCH_OUTPUT_DIR}" -type f 2>/dev/null || echo "(no bench files)"

# ---------------------------------------------------------------------------
# Step 7: Stop server
# ---------------------------------------------------------------------------
log "Stopping server container..."
docker rm -f "${CONTAINER_NAME}" 2>/dev/null || true
log "Done! Profile output: ${PROFILE_OUTPUT_DIR}"