#!/usr/bin/env bash
# v3 trace replay with dst-side migration breakdown instrumentation.
#
# Same trace + DR_FIX as `run_v3_replay.sh`, plus:
#   - instrument_dst_migration.py applied to vLLM scheduler
#   - DM_LOG_DIR exported to all 8 vLLM instances so per-PID
#     dst-migration logs land in <RUNDIR>/dst_mig_log/
#   - analyze_dst_migration.py runs on completion to print the
#     T_kv_pull vs queue-residual decomposition
#
# Usage: bash run_v3_dst_breakdown.sh

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/b3_v3_dstbreak_${DATE}}"
PYTHON="$PROJ_DIR/.venv/bin/python"
VLLM_ROOT="${VLLM_ROOT:-$PROJ_DIR/.venv/lib/python3.12/site-packages/vllm}"
DR_FIX_SCRIPT="$PROJ_DIR/microbench/connector_tax/cache_sweep/apply_direct_read_fix.py"
DM_INSTRUMENT="$PROJ_DIR/microbench/fresh_setup/instrument_dst_migration.py"
ANALYZE="$PROJ_DIR/microbench/connector_tax/cache_sweep/analyze_dst_migration.py"

mkdir -p "$OUTROOT"
DST_LOG_DIR="$OUTROOT/dst_mig_log"
mkdir -p "$DST_LOG_DIR"

echo "=== unified_v3 + dst-side migration breakdown ==="
echo "Trace      : $TRACE"
echo "Out        : $OUTROOT"
echo "DST logs   : $DST_LOG_DIR"
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
    "$PYTHON" "$DM_INSTRUMENT" --revert --venv "$PROJ_DIR/.venv" 2>/dev/null || true
}
trap cleanup_all EXIT
cleanup_all

echo "[stage 0a] applying CT_DR_FIX (env-gated)"
"$PYTHON" "$DR_FIX_SCRIPT" --apply --vllm-root "$VLLM_ROOT"

echo "[stage 0b] applying DST migration instrumentation"
"$PYTHON" "$DM_INSTRUMENT" --apply --venv "$PROJ_DIR/.venv"
"$PYTHON" "$DM_INSTRUMENT" --check --venv "$PROJ_DIR/.venv"

cfg_dir="$OUTROOT/unified_v3"
mkdir -p "$cfg_dir"

# Activate DR-fix env gate (consistent with run_v3_replay.sh)
export VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=1
# Export DM_LOG_DIR — every vLLM EngineCore inherits this env and writes
# its own dm_mig_pid<pid>.jsonl into it.
export DM_LOG_DIR="$DST_LOG_DIR"

echo ""
echo "====== unified_v3 ; DR_SYNC_DISABLED=1 ; DM_LOG_DIR=$DST_LOG_DIR ======"
bash "$PROJ_DIR/scripts/b3_isolated_policy.sh" "unified_v3" "$TRACE" "$cfg_dir" \
    2>&1 | tee "$cfg_dir/orchestrator.log" | tail -30

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

echo ""
echo "[stage Z] reverting DR_FIX + DM instrument"
"$PYTHON" "$DR_FIX_SCRIPT" --revert --vllm-root "$VLLM_ROOT"
"$PYTHON" "$DM_INSTRUMENT" --revert --venv "$PROJ_DIR/.venv"

echo ""
echo "[stage analyze] dst-side migration breakdown"
"$PYTHON" "$ANALYZE" \
    --proxy-breakdown "$cfg_dir/breakdown.json" \
    --dst-log-dir "$DST_LOG_DIR" \
    --output "$cfg_dir/dst_migration_breakdown.csv" \
    --plot "$cfg_dir/dst_migration_breakdown.png" \
    2>&1 | tee "$cfg_dir/dst_migration_breakdown.txt"

echo ""
echo "Done."
echo "  proxy breakdown : $cfg_dir/breakdown.json"
echo "  dst per-PID log : $DST_LOG_DIR/"
echo "  decomposition   : $cfg_dir/dst_migration_breakdown.{csv,png,txt}"
