#!/usr/bin/env python3
"""Extract per-DP-window speculative metrics from a vLLM engine log.

vLLM emits an ``Engine NNN`` throughput line followed by the corresponding
``SpecDecoding metrics`` line.  This script keeps that adjacency explicit and
does not infer request-level outcomes from the aggregate metrics.
"""

from __future__ import annotations

import argparse
import json
import math
import re
import statistics
from collections import defaultdict
from pathlib import Path
from typing import Any


ENGINE = re.compile(
    r"INFO (?P<clock>\d\d-\d\d \d\d:\d\d:\d\d\.\d+) .*?"
    r"Engine (?P<engine>\d+): .*?"
    r"Avg decode step duration: (?P<decode_ms>[0-9.]+) ms .*?"
    r"Running: (?P<running>\d+) reqs, Waiting: (?P<waiting>\d+) reqs"
)
SPEC = re.compile(
    r"INFO (?P<clock>\d\d-\d\d \d\d:\d\d:\d\d\.\d+) .*?"
    r"SpecDecoding metrics: Mean acceptance length: (?P<mean_accept>[0-9.]+), .*?"
    r"Avg Draft acceptance rate: (?P<accept_pct>[0-9.]+)%"
)


def percentile(values: list[float], q: float) -> float | None:
    if not values:
        return None
    values = sorted(values)
    index = (len(values) - 1) * q
    lower, upper = math.floor(index), math.ceil(index)
    if lower == upper:
        return values[lower]
    return values[lower] + (values[upper] - values[lower]) * (index - lower)


def describe(values: list[float]) -> dict[str, Any]:
    return {
        "n": len(values),
        "mean": statistics.fmean(values) if values else None,
        "min": min(values) if values else None,
        "p50": percentile(values, 0.5),
        "p95": percentile(values, 0.95),
        "max": max(values) if values else None,
        "distinct_value_count": len(set(values)),
    }


def extract(path: Path) -> dict[str, Any]:
    pending: tuple[str, dict[str, Any]] | None = None
    rows: list[dict[str, Any]] = []
    for line in path.read_text(errors="replace").splitlines():
        engine = ENGINE.search(line)
        if engine:
            pending = (
                engine.group("clock")[:17],
                {
                    "clock": engine.group("clock"),
                    "engine": int(engine.group("engine")),
                    "decode_ms": float(engine.group("decode_ms")),
                    "running": int(engine.group("running")),
                    "waiting": int(engine.group("waiting")),
                },
            )
            continue
        spec = SPEC.search(line)
        if not spec or pending is None:
            continue
        clock, row = pending
        if spec.group("clock")[:17] != clock:
            pending = None
            continue
        row["mean_accept_length"] = float(spec.group("mean_accept"))
        row["accept_pct"] = float(spec.group("accept_pct"))
        rows.append(row)
        pending = None

    by_engine: dict[int, list[dict[str, Any]]] = defaultdict(list)
    for row in rows:
        by_engine[row["engine"]].append(row)
    summaries: list[dict[str, Any]] = []
    for engine, items in sorted(by_engine.items()):
        summaries.append(
            {
                "engine": engine,
                "window_count": len(items),
                "accept_pct": describe([item["accept_pct"] for item in items]),
                "mean_accept_length": describe([item["mean_accept_length"] for item in items]),
                "decode_ms": describe([item["decode_ms"] for item in items]),
                "running_requests": describe([float(item["running"]) for item in items]),
            }
        )
    acceptance = [row["accept_pct"] for row in rows]
    return {
        "source": str(path),
        "records": rows,
        "per_engine": summaries,
        "data_sanity": {
            "accept_pct": {
                "n": len(acceptance),
                "min": min(acceptance) if acceptance else None,
                "max": max(acceptance) if acceptance else None,
                "distinct_value_count": len(set(acceptance)),
                "within_0_100": all(0 <= item <= 100 for item in acceptance),
            },
            "invariants": {
                "non_negative_decode_ms": all(row["decode_ms"] >= 0 for row in rows),
                "non_negative_running": all(row["running"] >= 0 for row in rows),
                "all_records_have_engine": all("engine" in row for row in rows),
            },
        },
    }


def main() -> int:
    parser = argparse.ArgumentParser(description=__doc__)
    parser.add_argument("--engine-log", type=Path, required=True)
    parser.add_argument("--output", type=Path, required=True)
    args = parser.parse_args()
    result = extract(args.engine_log)
    args.output.write_text(json.dumps(result, indent=2, sort_keys=True) + "\n")
    print(json.dumps({"per_engine": result["per_engine"], "data_sanity": result["data_sanity"]}, indent=2))
    return 0


if __name__ == "__main__":
    raise SystemExit(main())