#!/usr/bin/env python3
"""Scrape vLLM's own Prometheus /metrics on each locally-served model and append
one JSON snapshot line per reachable model to vllm_snapshots.jsonl.

vLLM counters (prompt_tokens_total, generation_tokens_total, request_success_total)
reset to zero on every process restart -- this is what turns them into a durable
history. A model that's down is skipped for this tick, not an error: run this
every few minutes from cron and it just accumulates whatever was actually up.
"""
import json
import os
import re
import time
import urllib.request

USAGE_DIR = os.path.expanduser("~/lumbridge/compute/.compute/usage")
SNAPSHOT_FILE = os.path.join(USAGE_DIR, "vllm_snapshots.jsonl")

# model_name -> port, per registry/models.yaml
MODELS = {
    "brain": 8001,
    "embed": 8012,
    "ocr": 8013,
}

# Prometheus exposition line: metric{labels} value
# vLLM metric names contain a colon (vllm:prompt_tokens_total), which \w does not match.
LINE_RE = re.compile(r'^([\w:]+)(\{[^}]*\})?\s+([0-9eE+\-.]+)\s*$')

COUNTERS = (
    "vllm:prompt_tokens_total",
    "vllm:generation_tokens_total",
    "vllm:num_requests_running",
)
# request_success_total is split by finished_reason -- sum all reasons.
SUCCESS_METRIC = "vllm:request_success_total"


def _fetch(port):
    url = f"http://127.0.0.1:{port}/metrics"
    with urllib.request.urlopen(url, timeout=3) as resp:
        return resp.read().decode()


def _parse(text):
    values = {name: 0.0 for name in COUNTERS}
    success_total = 0.0
    for line in text.splitlines():
        if not line or line.startswith("#"):
            continue
        m = LINE_RE.match(line)
        if not m:
            continue
        name, value = m.group(1), m.group(3)
        try:
            value = float(value)
        except ValueError:
            continue
        if name in values:
            values[name] = value
        elif name == SUCCESS_METRIC:
            success_total += value
    values[SUCCESS_METRIC] = success_total
    return values


def main():
    os.makedirs(USAGE_DIR, exist_ok=True)
    ts = time.time()
    lines = []
    for model, port in MODELS.items():
        try:
            text = _fetch(port)
        except Exception:
            continue  # model is down -- skip silently, not an error
        parsed = _parse(text)
        record = {
            "ts": ts,
            "model": model,
            "prompt_tokens_total": parsed["vllm:prompt_tokens_total"],
            "generation_tokens_total": parsed["vllm:generation_tokens_total"],
            "request_success_total": parsed[SUCCESS_METRIC],
            "num_requests_running": parsed["vllm:num_requests_running"],
        }
        lines.append(json.dumps(record))
    if lines:
        with open(SNAPSHOT_FILE, "a") as f:
            for line in lines:
                f.write(line + "\n")
    print(f"scraped {len(lines)}/{len(MODELS)} models reachable")


if __name__ == "__main__":
    main()