#!/usr/bin/env python3
"""Cache-sweep patch v2.

Extends `microbench/connector_tax/patches/apply_step_timing.py` with:

  scheduler.py  → adds `cache_size`     field (block_pool dict len) to
                  the per-step JSONL.
  kv_connector_model_runner_mixin.py
                → wraps `kv_connector.start_load_kv(...)` and
                  `kv_connector.get_finished(...)` to record their
                  duration into per-process `worker_step.jsonl`
                  (path from env var `CT_WORKER_STEP_LOG_PATH`).

All inserts carry the marker `# CT_CACHE_SWEEP_PATCH` so revert is
"delete every line that contains the marker". The original
`CONNECTOR_TAX_PATCH` lines (step_duration_us, build_meta_us) are
left alone.

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

import argparse
import re
import sys
from pathlib import Path

MARKER = "# CT_CACHE_SWEEP_PATCH"
BLOCK_BEGIN = f"{MARKER}_BLOCK_BEGIN"
BLOCK_END = f"{MARKER}_BLOCK_END"
# Default vLLM is the in-tree third_party copy; on dash0 the .venv links to it.
DEFAULT_VLLM_ROOT = Path(
    "/home/admin/cpfs/wjh/agentic-kv/.venv/lib/python3.12/site-packages/vllm"
)


def already_patched(text: str) -> bool:
    return MARKER in text


def revert_text(text: str) -> str:
    """Two pass revert:
       1. drop everything between BLOCK_BEGIN / BLOCK_END (inclusive)
       2. drop any remaining single line that contains MARKER
    Order matters so single-line MARKERs inside a block don't shrink-wrap.
    """
    out_lines = []
    inside_block = False
    for line in text.splitlines():
        if BLOCK_BEGIN in line:
            inside_block = True
            continue
        if BLOCK_END in line:
            inside_block = False
            continue
        if inside_block:
            continue
        if MARKER in line:
            continue
        out_lines.append(line)
    return "\n".join(out_lines) + ("\n" if text.endswith("\n") else "")


# ── scheduler.py: add cache_size field to record dict ───────────────────────
def patch_scheduler(text: str) -> str:
    if already_patched(text):
        print("  scheduler.py already has CT_CACHE_SWEEP_PATCH, skipping")
        return text

    # Inject `cache_size` field into the record built by
    # `_agentic_emit_step_log`. We piggy-back on the existing
    # CONNECTOR_TAX_PATCH lines that already inject `step_duration_us`
    # and `build_meta_us`. Insert immediately after `"t_unix": _time.time(),`.
    pat = (
        r"(        record = \{\n"
        r"(?:[^\n]*CONNECTOR_TAX_PATCH\n)*"
        r"            \"t_unix\": _time\.time\(\),\n)"
    )

    def repl(m):
        return (
            m.group(1)
            + "            \"cache_size\": _ct_cache_size(self),  " + MARKER + "\n"
        )

    text, n = re.subn(pat, repl, text, count=1)
    if n == 0:
        raise RuntimeError(
            "Failed to inject cache_size — is the base CONNECTOR_TAX_PATCH "
            "applied?"
        )

    # Add a helper at module bottom (only once). Wrap in block sentinels
    # so revert removes the whole block, not just lines that happen to
    # contain MARKER.
    helper = (
        f"\n# {BLOCK_BEGIN}\n"
        f"def _ct_cache_size(scheduler):\n"
        f"    try:\n"
        f"        return len(scheduler.kv_cache_manager.block_pool.cached_block_hash_to_block._cache)\n"
        f"    except Exception:\n"
        f"        return -1\n"
        f"# {BLOCK_END}\n"
    )
    if BLOCK_BEGIN not in text:
        text = text.rstrip() + "\n" + helper

    return text


# ── worker mixin: time start_load_kv + get_finished ─────────────────────────
def patch_worker_mixin(text: str) -> str:
    if already_patched(text):
        print("  kv_connector_model_runner_mixin.py already patched, skipping")
        return text

    # Wrap the start_load_kv call.
    pat = (
        r"        kv_connector\.start_load_kv\(get_forward_context\(\)\)\n"
    )
    repl = (
        "        _ct_t0 = __import__('time').perf_counter_ns()  " + MARKER + "\n"
        "        kv_connector.start_load_kv(get_forward_context())\n"
        "        _ct_start_load_kv_us = (__import__('time').perf_counter_ns() - _ct_t0) // 1000  " + MARKER + "\n"
    )
    text, n = re.subn(pat, repl, text, count=1)
    if n == 0:
        raise RuntimeError("Failed to patch start_load_kv call")

    # Wrap the get_finished call. We assign into local _ct_get_finished_us
    # in the same `try/finally` block where get_finished is called.
    pat = (
        r"            output\.finished_sending, output\.finished_recving = \(\n"
        r"                kv_connector\.get_finished\(scheduler_output\.finished_req_ids\)\n"
        r"            \)\n"
    )
    repl = (
        "            _ct_t1 = __import__('time').perf_counter_ns()  " + MARKER + "\n"
        "            output.finished_sending, output.finished_recving = (\n"
        "                kv_connector.get_finished(scheduler_output.finished_req_ids)\n"
        "            )\n"
        "            _ct_get_finished_us = (__import__('time').perf_counter_ns() - _ct_t1) // 1000  " + MARKER + "\n"
        "            _ct_emit_worker_step(_ct_start_load_kv_us, _ct_get_finished_us)  " + MARKER + "\n"
    )
    text, n = re.subn(pat, repl, text, count=1)
    if n == 0:
        raise RuntimeError("Failed to patch get_finished call")

    # Add the emitter at module bottom. Lazy-opens worker_step.jsonl on
    # the first call. Wrapped in block sentinels so revert is clean.
    emitter = f"""
# {BLOCK_BEGIN}
_ct_worker_fh = None
_ct_worker_init_done = False
_ct_worker_step_id = 0


def _ct_emit_worker_step(start_load_kv_us, get_finished_us):
    global _ct_worker_fh, _ct_worker_init_done, _ct_worker_step_id
    import os, json, time
    if not _ct_worker_init_done:
        _ct_worker_init_done = True
        path = os.environ.get("CT_WORKER_STEP_LOG_PATH")
        if not path:
            return
        try:
            rank = os.environ.get("RANK", os.environ.get("LOCAL_RANK", "0"))
            base, ext = os.path.splitext(path)
            full = f"{{base}}.r{{rank}}{{ext or '.jsonl'}}"
            _ct_worker_fh = open(full, "a", buffering=1)
        except Exception:
            _ct_worker_fh = None
    fh = _ct_worker_fh
    if fh is None:
        return
    try:
        fh.write(json.dumps({{
            "t_unix": time.time(),
            "step_id": _ct_worker_step_id,
            "start_load_kv_us": int(start_load_kv_us),
            "get_finished_us": int(get_finished_us),
        }}) + "\\n")
        _ct_worker_step_id += 1
    except Exception:
        try:
            fh.close()
        except Exception:
            pass
        _ct_worker_fh = None
# {BLOCK_END}
"""
    if BLOCK_BEGIN not in text:
        text = text.rstrip() + "\n" + emitter

    return text


# ── driver ──────────────────────────────────────────────────────────────────
def apply_to_file(path: Path, fn) -> bool:
    if not path.exists():
        print(f"  SKIP {path} (not found)")
        return False
    orig = path.read_text()
    new = fn(orig)
    if new == orig:
        print(f"  unchanged: {path}")
        return False
    path.write_text(new)
    print(f"  patched ({new.count(MARKER)} marks): {path}")
    return True


def revert_file(path: Path) -> bool:
    if not path.exists():
        return False
    orig = path.read_text()
    new = revert_text(orig)
    if new == orig:
        print(f"  no marks: {path}")
        return False
    path.write_text(new)
    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")

    sched = args.vllm_root / "v1/core/sched/scheduler.py"
    mixin = args.vllm_root / "v1/worker/kv_connector_model_runner_mixin.py"

    if args.apply:
        print(f"Applying CT_CACHE_SWEEP_PATCH to {args.vllm_root}")
        apply_to_file(sched, patch_scheduler)
        apply_to_file(mixin, patch_worker_mixin)
    else:
        print(f"Reverting CT_CACHE_SWEEP_PATCH from {args.vllm_root}")
        revert_file(sched)
        revert_file(mixin)


if __name__ == "__main__":
    main()
