#!/usr/bin/env python3
"""Smoke test: does remote-prefill on dst leave its prefix cache discoverable
to a follow-up turn on the same instance?

Reproducer for the v3 rotation bug observed in unified_v3 (next-turn at
decode_target sees cached_tokens=0 despite migration's `cache_blocks`
supposedly running).

Expects 2 vLLM instances running on 127.0.0.1:8000 and 8001 with
--kv-transfer-config '{"kv_connector":"MooncakeConnector","kv_role":"kv_both"}'
and Mooncake bootstrap servers on 8998 and 8999.

Run flow:
  1. Query both bootstrap servers for engine_ids.
  2. Send migration: src=8000 do_remote_decode (max_tokens=1), then
     dst=8001 do_remote_prefill (pulls KV via Mooncake) with same prompt.
  3. Send follow-up: same session prompt + tiny extension, hit 8001
     directly (no kv_transfer_params), check cached_tokens.

A working migration with prefix-cache visibility would see ~100% cached
on the follow-up (full prefix hit). The v3 bug shows cached=0.
"""
import asyncio
import argparse
import json
import sys
import uuid
import httpx


async def get_engine_id(client: httpx.AsyncClient, port: int) -> str:
    url = f"http://127.0.0.1:{port}/query"
    r = await client.get(url)
    r.raise_for_status()
    data = r.json()
    # data = {"0": {"engine_id": "..."}, ...}
    return data["0"]["engine_id"]


async def send_completion(
    client: httpx.AsyncClient,
    host_port: int,
    prompt: list[int],
    max_tokens: int,
    kv_transfer_params: dict | None = None,
    stream: bool = False,
) -> dict:
    payload = {
        "model": "/home/admin/cpfs/wjh/models/Qwen/Qwen3-Coder-30B-A3B-Instruct",
        "prompt": prompt,
        "max_tokens": max_tokens,
        "min_tokens": max_tokens if max_tokens == 1 else 1,
        "temperature": 0.0,
        "stream": stream,
    }
    if stream:
        payload["stream_options"] = {"include_usage": True}
    if kv_transfer_params:
        payload["kv_transfer_params"] = kv_transfer_params
    url = f"http://127.0.0.1:{host_port}/v1/completions"

    if not stream:
        r = await client.post(url, json=payload, timeout=300.0)
        r.raise_for_status()
        return r.json()
    else:
        # Stream and collect chunks; keep the LAST chunk that contains
        # `usage` (with include_usage, the very last data: chunk has it).
        last_with_usage = None
        last_any = None
        async with client.stream("POST", url, json=payload, timeout=300.0) as resp:
            resp.raise_for_status()
            buffer = ""
            async for chunk in resp.aiter_bytes():
                buffer += chunk.decode("utf-8", errors="replace")
                while "\n\n" in buffer:
                    line, buffer = buffer.split("\n\n", 1)
                    if line.startswith("data: "):
                        data_str = line[6:].strip()
                        if data_str == "[DONE]":
                            continue
                        try:
                            d = json.loads(data_str)
                            last_any = d
                            if d.get("usage"):
                                last_with_usage = d
                        except Exception:
                            pass
        return last_with_usage or last_any or {}


def short(d: dict) -> str:
    """Pull cached_tokens out of the response usage section."""
    if not d:
        return "no_resp"
    usage = d.get("usage") or {}
    details = usage.get("prompt_tokens_details") or {}
    cached = details.get("cached_tokens", 0) or usage.get("cached_tokens", 0)
    return (
        f"cached={cached}/{usage.get('prompt_tokens', 0)} "
        f"completion={usage.get('completion_tokens', 0)} "
        f"id={d.get('id', '?')[:24]}"
    )


async def main():
    p = argparse.ArgumentParser()
    p.add_argument("--src-port", type=int, default=8000)
    p.add_argument("--dst-port", type=int, default=8001)
    p.add_argument("--src-bp", type=int, default=8998)
    p.add_argument("--dst-bp", type=int, default=8999)
    p.add_argument("--n-prefix-tokens", type=int, default=8192,
                   help="Length of synthetic prompt prefix (tokens)")
    p.add_argument("--n-extension", type=int, default=32,
                   help="Tokens added in the follow-up request")
    p.add_argument("--noise-reqs", type=int, default=0,
                   help="Number of unrelated requests to send to dst between "
                        "migration and follow-up (eviction-pressure test)")
    p.add_argument("--noise-tokens", type=int, default=16384,
                   help="Tokens per noise request")
    p.add_argument("--noise-parallel", type=int, default=4,
                   help="How many noise requests in parallel")
    p.add_argument("--prefix-base", type=int, default=100,
                   help="Token-id base for the prompt prefix (use distinct value "
                        "across iterations to avoid prefix-cache pollution).")
    p.add_argument("--decode-tokens", type=int, default=16,
                   help="max_tokens on dst decode request")
    args = p.parse_args()

    async with httpx.AsyncClient() as client:
        src_id = await get_engine_id(client, args.src_bp)
        dst_id = await get_engine_id(client, args.dst_bp)
        print(f"src engine_id = {src_id}")
        print(f"dst engine_id = {dst_id}")
        print()

        # Build a deterministic prompt: a long enough token sequence to be
        # multiple blocks. Use simple range to avoid tokenizer dependence.
        # Build deterministic prompt using values in safe vocab range [1024, 100000)
        # so high prefix-base values don't overflow the tokenizer vocab.
        import random as _r
        rng = _r.Random(f"prefix-{args.prefix_base}")
        prompt = [rng.randint(1024, 99_999) for _ in range(args.n_prefix_tokens)]

        # ----- Step 1: Migration. src does prefill (max_tokens=1, no stream),
        # then dst pulls KV and decodes. -----
        transfer_id = f"smoke-xfer-{uuid.uuid4().hex[:8]}"
        print(f"[1] migration: transfer_id={transfer_id}")

        # First: cold-prefill on src (no Mooncake yet, just establish src has the KV)
        # The proxy convention: src sees do_remote_decode=True so it will SEND its KV
        # via Mooncake later. We do this as a single-shot.
        import time as _t
        t_src_start = _t.monotonic()
        src_resp_task = asyncio.create_task(
            send_completion(
                client, args.src_port, prompt, max_tokens=1,
                kv_transfer_params={
                    "do_remote_decode": True,
                    "do_remote_prefill": False,
                    "transfer_id": transfer_id,
                },
                stream=False,
            )
        )

        # Slight stagger: dst needs to be ready to pull. In production the
        # proxy waits for src to finish before sending dst. We'll do the
        # same — await src then send dst.
        src_resp = await src_resp_task
        t_src_done = _t.monotonic()
        print(f"    src prefill resp ({(t_src_done-t_src_start)*1000:.0f}ms): {short(src_resp)}")

        bootstrap_addr = f"http://127.0.0.1:{args.src_bp}"
        t_dst_start = _t.monotonic()
        dst_resp = await send_completion(
            client, args.dst_port, prompt, max_tokens=args.decode_tokens,
            kv_transfer_params={
                "do_remote_decode": False,
                "do_remote_prefill": True,
                "remote_bootstrap_addr": bootstrap_addr,
                "remote_engine_id": src_id,
                "transfer_id": transfer_id,
            },
            stream=True,
        )
        t_dst_done = _t.monotonic()
        # dst time = KV transfer + decode of N tokens. Subtract approx decode time
        # to isolate transfer cost.
        usage_d = dst_resp.get("usage") or {}
        n_completion = usage_d.get("completion_tokens", 0)
        dst_total_ms = (t_dst_done - t_dst_start) * 1000
        print(f"    dst decode resp ({dst_total_ms:.0f}ms, {n_completion} completion tokens): {short(dst_resp)}")
        print(f"    [TIMING] proto={args.n_prefix_tokens}p src_prefill={int((t_src_done-t_src_start)*1000)}ms "
              f"dst_total={int(dst_total_ms)}ms (KV xfer + {n_completion}-token decode)")
        print()

        # ----- Step 1.5: Noise. Send unrelated requests to dst to test eviction. -----
        if args.noise_reqs > 0:
            print(f"[1.5] sending {args.noise_reqs} noise requests "
                  f"(tokens={args.noise_tokens}, parallel={args.noise_parallel}) to dst")
            async def noise(idx):
                rng_n = _r.Random(f"noise-{args.prefix_base}-{idx}")
                p_n = [rng_n.randint(1024, 99_999) for _ in range(args.noise_tokens)]
                return await send_completion(
                    client, args.dst_port, p_n, max_tokens=1,
                    kv_transfer_params=None, stream=False,
                )
            sem = asyncio.Semaphore(args.noise_parallel)
            async def gated(idx):
                async with sem:
                    return await noise(idx)
            results = await asyncio.gather(*[gated(i) for i in range(args.noise_reqs)])
            done = sum(1 for r in results if r and r.get("usage"))
            print(f"    noise: {done}/{args.noise_reqs} completed")
            print()

        # ----- Step 2: Follow-up. Same session, extended prompt, hit dst directly. -----
        rng_ext = _r.Random(f"ext-{args.prefix_base}")
        follow_prompt = prompt + [rng_ext.randint(1024, 99_999) for _ in range(args.n_extension)]
        print(f"[2] follow-up direct to dst (no kv_transfer_params): "
              f"prefix_len={args.n_prefix_tokens}, extended_len={len(follow_prompt)}")
        fu = await send_completion(
            client, args.dst_port, follow_prompt, max_tokens=4,
            kv_transfer_params=None,
            stream=False,
        )
        print(f"    follow-up resp: {short(fu)}")
        print()

        # ----- Step 3: Same prompt twice on dst (sanity) -----
        print(f"[3] sanity: same prompt again to dst (should see local hit "
              f"from step 2's just-cached blocks)")
        sanity = await send_completion(
            client, args.dst_port, follow_prompt, max_tokens=4,
            kv_transfer_params=None,
            stream=False,
        )
        print(f"    sanity resp: {short(sanity)}")
        print()

        # ----- Verdict -----
        usage_fu = fu.get("usage") or {}
        details_fu = usage_fu.get("prompt_tokens_details") or {}
        cached_fu = details_fu.get("cached_tokens", 0) or usage_fu.get("cached_tokens", 0)
        # Expect ~n_prefix_tokens (minus the last token + alignment)
        expected_min = int(args.n_prefix_tokens * 0.95)
        verdict = "PASS" if cached_fu >= expected_min else "FAIL"
        print(f"=== verdict: {verdict} (follow-up cached={cached_fu}, "
              f"expected >= {expected_min} of {args.n_prefix_tokens}) ===")
        sys.exit(0 if verdict == "PASS" else 1)


if __name__ == "__main__":
    asyncio.run(main())
