#!/usr/bin/env python3
"""Env-gate the direct-RDMA-read hash sync in MooncakeConnectorScheduler.

The `set(self._block_pool.cached_block_hash_to_block._cache.keys())` walk
at `mooncake_connector.py:434-450` was introduced in our own commit
`a7df84b` ("Direct RDMA read: D reads cached KV from C's GPU without
C's scheduler") and runs unconditionally every scheduler step regardless
of whether any direct-read consumer is actually attached.

This patch lets us turn it off at runtime via:

    VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC=1

so we can measure the slope contribution of that block in isolation. With
the env var set, the entire hash-table sync block is skipped and the
bootstrap server's `_hash_table` stays empty (which is fine as long as
no direct_read consumer ever queries — and our synthetic bench never
does).

Usage:
  python apply_direct_read_fix.py --apply  [--vllm-root PATH]
  python apply_direct_read_fix.py --revert [--vllm-root PATH]
"""

import argparse
import re
from pathlib import Path

MARKER = "# CT_DR_FIX"
DEFAULT_VLLM_ROOT = Path(
    "/home/admin/cpfs/wjh/agentic-kv/.venv/lib/python3.12/site-packages/vllm"
)

CONN_REL = "distributed/kv_transfer/kv_connector/v1/mooncake/mooncake_connector.py"

# The exact line we replace
ORIG_LINE = "        if self._block_pool is not None:"
PATCHED_LINE = (
    "        if self._block_pool is not None and "
    "__import__('os').environ.get('VLLM_MOONCAKE_DISABLE_DIRECT_READ_SYNC', '0') != '1':"
    "  " + MARKER
)


def apply_to_file(path: Path) -> bool:
    if not path.exists():
        print(f"  SKIP {path} (not found)")
        return False
    text = path.read_text()
    if MARKER in text:
        print(f"  already patched: {path}")
        return False
    if ORIG_LINE not in text:
        raise RuntimeError(
            f"Cannot find target line `{ORIG_LINE.strip()}` in {path}. "
            "File may have been refactored."
        )
    new_text = text.replace(ORIG_LINE, PATCHED_LINE, 1)
    path.write_text(new_text)
    print(f"  patched: {path}")
    return True


def revert_file(path: Path) -> bool:
    if not path.exists():
        return False
    text = path.read_text()
    if MARKER not in text:
        print(f"  no marker: {path}")
        return False
    new_text = text.replace(PATCHED_LINE, ORIG_LINE, 1)
    path.write_text(new_text)
    print(f"  reverted: {path}")
    return True


def main():
    ap = argparse.ArgumentParser()
    ap.add_argument("--apply", action="store_true")
    ap.add_argument("--revert", action="store_true")
    ap.add_argument("--vllm-root", type=Path, default=DEFAULT_VLLM_ROOT)
    args = ap.parse_args()
    if not (args.apply ^ args.revert):
        ap.error("Specify exactly one of --apply / --revert")
    conn = args.vllm_root / CONN_REL
    if args.apply:
        print(f"Applying CT_DR_FIX to {args.vllm_root}")
        apply_to_file(conn)
    else:
        print(f"Reverting CT_DR_FIX from {args.vllm_root}")
        revert_file(conn)


if __name__ == "__main__":
    main()
