#!/usr/bin/env python3
"""Apply / revert step-timing additions to vLLM's existing agentic step log.

vLLM already has a per-step JSONL emitter triggered by AGENTIC_STEP_LOG_PATH.
This patch enriches it with three duration fields:

  step_duration_us   : full schedule() wall time
  build_meta_us      : duration of self.connector.build_connector_meta(...)
  emit_overhead_us   : duration of _agentic_emit_step_log itself
                       (lets us subtract the patch's own cost)

We also add a worker-side patch in kv_connector_model_runner_mixin.py to
record start_load_kv() duration into a per-process JSONL file pointed to
by VLLM_PD_PROFILE_LOG=$RUN_DIR/worker_step.jsonl.

All inserts are marked with `# CONNECTOR_TAX_PATCH` so revert is just
"delete those lines".

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

import argparse
import re
import sys
from pathlib import Path

MARKER = "# CONNECTOR_TAX_PATCH"
DEFAULT_VLLM_ROOT = (
    Path.home()
    / "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:
    out = [l for l in text.splitlines() if MARKER not in l]
    return "\n".join(out) + ("\n" if text.endswith("\n") else "")


# ── scheduler.py patches ────────────────────────────────────────────────────
def patch_scheduler(text: str) -> str:
    if already_patched(text):
        print("  scheduler.py already patched, skipping")
        return text

    # 1. At top of schedule(): record _step_t0
    pat = (
        r"(    def schedule\(self\) -> SchedulerOutput:\n"
        r"        # NOTE\(woosuk\) on the scheduling algorithm:)"
    )
    repl = (
        r"    def schedule(self) -> SchedulerOutput:\n"
        r"        import time as _t  " + MARKER + "\n"
        r"        _step_t0 = _t.perf_counter_ns()  " + MARKER + "\n"
        r"        self._ct_build_meta_ns = 0  " + MARKER + "\n"
        r"        # NOTE(woosuk) on the scheduling algorithm:"
    )
    text, n = re.subn(pat, repl, text, count=1)
    if n == 0:
        raise RuntimeError("Failed to patch schedule() entry")

    # 2. Time the build_connector_meta call
    pat = (
        r"        if self\.connector is not None:\n"
        r"            meta: KVConnectorMetadata = self\.connector\.build_connector_meta\(\n"
        r"                scheduler_output\n"
        r"            \)\n"
        r"            scheduler_output\.kv_connector_metadata = meta"
    )
    repl = (
        "        if self.connector is not None:\n"
        f"            _bm_t0 = _t.perf_counter_ns()  {MARKER}\n"
        "            meta: KVConnectorMetadata = self.connector.build_connector_meta(\n"
        "                scheduler_output\n"
        "            )\n"
        f"            self._ct_build_meta_ns = _t.perf_counter_ns() - _bm_t0  {MARKER}\n"
        "            scheduler_output.kv_connector_metadata = meta"
    )
    text, n = re.subn(pat, repl, text, count=1)
    if n == 0:
        raise RuntimeError("Failed to patch build_connector_meta")

    # 3. Pass step duration into _agentic_emit_step_log via attribute
    #    (cleaner than threading kwargs through). Then in the emit
    #    function inject the new fields into `record`.
    pat = (
        r"        if self\._agentic_step_log_fh is not None:\n"
        r"            self\._agentic_emit_step_log\("
    )
    repl = (
        f"        if self._agentic_step_log_fh is not None:\n"
        f"            self._ct_step_duration_ns = _t.perf_counter_ns() - _step_t0  {MARKER}\n"
        f"            self._agentic_emit_step_log("
    )
    text, n = re.subn(pat, repl, text, count=1)
    if n == 0:
        raise RuntimeError("Failed to patch step_duration insertion")

    # 4. Inject extra fields into the `record` dict in _agentic_emit_step_log
    pat = (
        r"        record = \{\n"
        r"            \"t_unix\": _time\.time\(\),"
    )
    repl = (
        "        record = {\n"
        f"            \"step_duration_us\": getattr(self, '_ct_step_duration_ns', 0) // 1000,  {MARKER}\n"
        f"            \"build_meta_us\":    getattr(self, '_ct_build_meta_ns', 0) // 1000,  {MARKER}\n"
        "            \"t_unix\": _time.time(),"
    )
    text, n = re.subn(pat, repl, text, count=1)
    if n == 0:
        raise RuntimeError("Failed to patch record dict")

    return text


# ── apply / revert 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"

    if args.apply:
        print(f"Applying connector-tax step-timing patch to {args.vllm_root}")
        apply_to_file(sched, patch_scheduler)
    else:
        print(f"Reverting connector-tax step-timing patch from {args.vllm_root}")
        revert_file(sched)


if __name__ == "__main__":
    main()