#!/usr/bin/env bash
# Re-run the elastic_migration_v2 trace replay (8x TP1 + unified policy)
# with the DR-fix applied, to test whether the +45% TTFT p90 gap closes.
#
# Three policies in sequence on the SAME trace and SAME machine:
#   unified              — plain (no connector)               control
#   unified_kv_both      — Mooncake kv_both, DR sync ON       baseline
#   unified_kv_both_drfix — Mooncake kv_both, DR sync OFF      fix under test
#
# The DR-fix patch is applied at the start and reverted at the end.
# When env VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=1 is set, the patched
# build_connector_meta skips the O(|cache|) hash sync.

set -uo pipefail

PROJ_DIR="${PROJ_DIR:-/home/admin/cpfs/wjh/agentic-kv}"
TRACE="${TRACE:-$PROJ_DIR/traces/w600_r0.0015_st30.jsonl}"
DATE="$(date +%Y%m%d_%H%M)"
OUTROOT="${OUTROOT:-$PROJ_DIR/outputs/trace_replay_drfix_${DATE}}"
PYTHON="$PROJ_DIR/.venv/bin/python"
DR_FIX_SCRIPT="$PROJ_DIR/microbench/connector_tax/cache_sweep/apply_direct_read_fix.py"
VLLM_ROOT="${VLLM_ROOT:-$PROJ_DIR/.venv/lib/python3.12/site-packages/vllm}"

mkdir -p "$OUTROOT"
echo "=== Trace-replay DR-fix A/B/C ==="
echo "Trace : $TRACE"
echo "Out   : $OUTROOT"
echo "Policies in order: unified → unified_kv_both → unified_kv_both_drfix"
echo ""

cleanup_all() {
    pkill -9 -f cache_aware_proxy 2>/dev/null || true
    pkill -9 -f "vllm serve" 2>/dev/null || true
    pkill -9 -f "EngineCore" 2>/dev/null || true
    sleep 5
    "$PYTHON" "$DR_FIX_SCRIPT" --revert --vllm-root "$VLLM_ROOT" 2>/dev/null || true
}
trap cleanup_all EXIT
cleanup_all

echo "[stage 0] applying CT_DR_FIX patch (skip flag still respects env var)"
"$PYTHON" "$DR_FIX_SCRIPT" --apply --vllm-root "$VLLM_ROOT"

run_policy() {
    local policy="$1"
    local rundir="$2"
    local skip_dr="$3"
    echo ""
    echo "====== $policy → $(basename $rundir) ; VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=$skip_dr ======"
    mkdir -p "$rundir"
    # Pass-through to the existing isolated runner; it spawns 8 vLLMs and
    # the replayer. The env var is inherited by all vLLM child processes.
    if [ "$skip_dr" = "1" ]; then
        export VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=1
    else
        unset VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC
    fi
    bash "$PROJ_DIR/scripts/b3_isolated_policy.sh" "$policy" "$TRACE" "$rundir" \
        2>&1 | tee "$rundir/orchestrator.log" | tail -30
    rc="${PIPESTATUS[0]}"
    if [ "$rc" != "0" ]; then
        echo "[FAIL] policy $policy rc=$rc"
    fi
    # Hard cleanup between policies (b3_isolated_policy.sh's trap should
    # have already done it but belt-and-braces)
    pkill -9 -f cache_aware_proxy 2>/dev/null || true
    pkill -9 -f "vllm serve" 2>/dev/null || true
    pkill -9 -f "EngineCore" 2>/dev/null || true
    sleep 10
    return 0
}

run_policy "unified"          "$OUTROOT/unified"                  "0"
run_policy "unified_kv_both"  "$OUTROOT/unified_kv_both"          "0"
# Same policy + launcher; the env-gated skip is what makes this differ.
run_policy "unified_kv_both"  "$OUTROOT/unified_kv_both_drfix"    "1"

echo ""
echo "[stage Z] reverting CT_DR_FIX patch"
"$PYTHON" "$DR_FIX_SCRIPT" --revert --vllm-root "$VLLM_ROOT"

echo ""
echo "Done. Artifacts: $OUTROOT"
echo "  unified              : $OUTROOT/unified/metrics.jsonl"
echo "  unified_kv_both      : $OUTROOT/unified_kv_both/metrics.jsonl"
echo "  unified_kv_both_drfix: $OUTROOT/unified_kv_both_drfix/metrics.jsonl"
