"""Shared helpers for v2 GPU-hit-first experiments."""
from __future__ import annotations

import random
import time
import requests

# Qwen3-Coder geometry (from config.json): 48 layers, 4 KV heads, head_dim 128, bf16
KV_BYTES_PER_TOKEN = 98304  # 96 KiB
VOCAB = 151936
# Safe token-id range: avoid low special-ish ids and the high special tokens (>=151643)
TOK_LO, TOK_HI = 1000, 151000


def make_token_prompt(length: int, seed: int) -> list[int]:
    """Deterministic, content-addressed token-id prompt of exact `length`.

    Same (length, seed) -> same ids -> prefix-cache hit.
    Different seed -> fresh ids -> miss.
    """
    rng = random.Random(seed)
    return [rng.randint(TOK_LO, TOK_HI) for _ in range(length)]


def scrape_prefix_cache(endpoint: str) -> dict:
    """Return cumulative prefix-cache counters from vLLM /metrics.

    Keys: gpu_hits, gpu_queries, ext_hits, ext_queries (floats, cumulative).
    """
    out = {"gpu_hits": 0.0, "gpu_queries": 0.0, "ext_hits": 0.0, "ext_queries": 0.0}
    try:
        txt = requests.get(f"{endpoint}/metrics", timeout=10).text
    except Exception:
        return out
    for line in txt.splitlines():
        if line.startswith("#") or not line:
            continue
        try:
            name, val = line.rsplit(" ", 1)
            v = float(val)
        except ValueError:
            continue
        # strip prometheus labels and match only the cumulative _total counters
        # (exclude _created epoch-timestamp series, which would dominate the sum)
        metric = name.split("{", 1)[0]
        if metric == "vllm:external_prefix_cache_hits_total":
            out["ext_hits"] += v
        elif metric == "vllm:external_prefix_cache_queries_total":
            out["ext_queries"] += v
        elif metric == "vllm:prefix_cache_hits_total":
            out["gpu_hits"] += v
        elif metric == "vllm:prefix_cache_queries_total":
            out["gpu_queries"] += v
    return out


def measure_ttft(endpoint: str, model: str, prompt_ids: list[int],
                 max_tokens: int = 1, timeout: float = 600.0) -> dict:
    """Send one streaming /v1/completions request; return TTFT and e2e seconds.

    TTFT = time from send to first streamed token chunk (== prefill wall time).
    """
    url = f"{endpoint}/v1/completions"
    payload = {
        "model": model,
        "prompt": prompt_ids,
        "max_tokens": max_tokens,
        "temperature": 0.0,
        "stream": True,
        "stream_options": {"include_usage": True},
    }
    t0 = time.perf_counter()
    ttft = None
    usage = None
    with requests.post(url, json=payload, stream=True, timeout=timeout) as r:
        r.raise_for_status()
        for raw in r.iter_lines():
            if not raw:
                continue
            line = raw.decode("utf-8") if isinstance(raw, bytes) else raw
            if not line.startswith("data: "):
                continue
            data = line[6:]
            if data.strip() == "[DONE]":
                break
            import json as _json
            obj = _json.loads(data)
            if obj.get("usage"):
                usage = obj["usage"]
            choices = obj.get("choices") or []
            if ttft is None and choices and choices[0].get("text"):
                ttft = time.perf_counter() - t0
    e2e = time.perf_counter() - t0
    return {"ttft_s": ttft if ttft is not None else e2e, "e2e_s": e2e, "usage": usage}


def wait_healthy(endpoint: str, timeout: float = 900.0) -> bool:
    deadline = time.time() + timeout
    while time.time() < deadline:
        try:
            if requests.get(f"{endpoint}/health", timeout=5).status_code == 200:
                return True
        except Exception:
            pass
        time.sleep(3)
    return False