"""Speculative-decoding metrics, read from the engine rather than inferred.

Throughput alone cannot tell you whether speculative decoding is working. A draft model that
is rejected almost every time still produces tokens — just slowly, having paid for the draft
pass as well as the verify pass — and on a synthetic benchmark that can even look like a
modest win. The number that says whether it is working is the **acceptance rate**: of the
tokens the draft proposed, how many survived verification.

vLLM exposes this on its Prometheus endpoint as counters, so a rate for one sweep point is a
delta across that point rather than the process-lifetime total. Reading the total instead would
mix in every request since the server started, including warm-up.

Parsing is pure standard library and `httpx` is only needed to fetch — the import is deferred
so the parser can be tested, and reasoned about, without the HTTP client installed. Same rule as
`kbench/schema.py`: the logic that decides what a number means must not depend on the machinery
that fetches it.

This is deliberately best-effort. An engine without these counters, or with metrics disabled,
gets `None` and the sweep continues — an absent metric must never fail a benchmark run.

Why it matters for this bench specifically: `perf.py` sends a random nonce prefix and
`ignore_eos`, both correct for defeating prefix caching and pinning output length, and both
adversarial to speculative decoding. Forced continuation past a natural stop is degenerate
text, and a draft model's acceptance on degenerate text is not its acceptance on real work. So
a speculative target measured with the synthetic prompt needs its acceptance rate reported
beside the throughput, or the throughput will be read as if it transfers.
"""

from __future__ import annotations

import re
from dataclasses import dataclass
from typing import TYPE_CHECKING

if TYPE_CHECKING:  # pragma: no cover
    import httpx

# vLLM V1 names. The older V0 engine used vllm:spec_decode_* with the same meaning; both are
# accepted so a mixed fleet does not silently report nothing.
ACCEPTED = ("vllm:spec_decode_num_accepted_tokens_total",)
DRAFT = ("vllm:spec_decode_num_draft_tokens_total",)

_SAMPLE = re.compile(r"^(?P<name>[a-zA-Z_:][\w:]*)(?P<labels>\{[^}]*\})?\s+(?P<value>[-+0-9.eE]+)$")


@dataclass(frozen=True)
class SpecCounters:
    accepted: float
    draft: float

    def __sub__(self, other: "SpecCounters") -> "SpecCounters":
        return SpecCounters(self.accepted - other.accepted, self.draft - other.draft)

    @property
    def acceptance_rate(self) -> float | None:
        """Accepted / proposed over this window. None when nothing was proposed."""
        if self.draft <= 0:
            return None
        return round(self.accepted / self.draft, 4)


def parse_metrics(text: str) -> SpecCounters | None:
    """Pull the speculative counters out of a Prometheus exposition payload."""
    totals: dict[str, float] = {}
    for line in text.splitlines():
        line = line.strip()
        if not line or line.startswith("#"):
            continue
        m = _SAMPLE.match(line)
        if not m:
            continue
        name = m.group("name")
        if name in ACCEPTED or name in DRAFT:
            try:
                # Counters are per-label-set (one per model); sum them. A single-model server
                # has one series, and summing is still correct there.
                totals[name] = totals.get(name, 0.0) + float(m.group("value"))
            except ValueError:
                continue

    accepted = next((totals[n] for n in ACCEPTED if n in totals), None)
    draft = next((totals[n] for n in DRAFT if n in totals), None)
    if accepted is None or draft is None:
        return None
    return SpecCounters(accepted=accepted, draft=draft)


def metrics_url(base_url: str) -> str:
    """/metrics sits at the server root, not under the OpenAI /v1 prefix."""
    root = base_url.rstrip("/")
    for suffix in ("/v1", "/openai/v1"):
        if root.endswith(suffix):
            root = root[: -len(suffix)]
            break
    return f"{root}/metrics"


async def read_counters(client: "httpx.AsyncClient", base_url: str) -> SpecCounters | None:
    """Best-effort read. Any failure means "no data", never an exception into the sweep."""
    try:
        resp = await client.get(metrics_url(base_url), timeout=5.0)
        if resp.status_code != 200:
            return None
        return parse_metrics(resp.text)
    except Exception:  # noqa: BLE001 - a missing metrics endpoint is an ordinary outcome
        return None