"""Speculative-decoding counter parsing.

The parser reads a live Prometheus payload, so its failure modes are all "the server said
something slightly different than expected". Each of these is a shape a real /metrics endpoint
actually produces.

    python3 -m unittest discover -s tests -v
"""

from __future__ import annotations

import sys
import unittest
from pathlib import Path

sys.path.insert(0, str(Path(__file__).resolve().parent.parent))

from kbench.specdec import SpecCounters, metrics_url, parse_metrics  # noqa: E402

PAYLOAD = """\
# HELP vllm:spec_decode_num_draft_tokens_total Number of draft tokens.
# TYPE vllm:spec_decode_num_draft_tokens_total counter
vllm:spec_decode_num_draft_tokens_total{model_name="qwen"} 1000.0
# HELP vllm:spec_decode_num_accepted_tokens_total Number accepted.
# TYPE vllm:spec_decode_num_accepted_tokens_total counter
vllm:spec_decode_num_accepted_tokens_total{model_name="qwen"} 720.0
vllm:num_requests_running{model_name="qwen"} 4.0
"""


class Parsing(unittest.TestCase):
    def test_reads_both_counters(self):
        c = parse_metrics(PAYLOAD)
        self.assertEqual((c.draft, c.accepted), (1000.0, 720.0))
        self.assertEqual(c.acceptance_rate, 0.72)

    def test_ignores_comments_and_unrelated_series(self):
        c = parse_metrics(PAYLOAD)
        self.assertEqual(c.draft, 1000.0, "an unrelated metric leaked into the total")

    def test_sums_across_label_sets(self):
        # A server hosting two models exposes one series each. Summing is correct, and taking
        # the first would silently measure only one of them.
        two = PAYLOAD + (
            'vllm:spec_decode_num_draft_tokens_total{model_name="b"} 500.0\n'
            'vllm:spec_decode_num_accepted_tokens_total{model_name="b"} 250.0\n'
        )
        c = parse_metrics(two)
        self.assertEqual((c.draft, c.accepted), (1500.0, 970.0))

    def test_no_speculative_counters_means_none_not_zero(self):
        # An engine without speculative decoding must report "no data". Zero would render as a
        # 0% acceptance rate, which reads as a broken draft model rather than an absent one.
        self.assertIsNone(parse_metrics("vllm:num_requests_running 1.0\n"))

    def test_a_missing_half_is_not_usable(self):
        self.assertIsNone(
            parse_metrics('vllm:spec_decode_num_draft_tokens_total{m="x"} 10.0\n')
        )

    def test_malformed_values_do_not_raise(self):
        payload = (
            'vllm:spec_decode_num_draft_tokens_total{m="x"} not-a-number\n'
            'vllm:spec_decode_num_accepted_tokens_total{m="x"} 5.0\n'
        )
        self.assertIsNone(parse_metrics(payload))

    def test_unlabelled_series_are_read(self):
        payload = (
            "vllm:spec_decode_num_draft_tokens_total 100.0\n"
            "vllm:spec_decode_num_accepted_tokens_total 50.0\n"
        )
        self.assertEqual(parse_metrics(payload).acceptance_rate, 0.5)


class WindowedRate(unittest.TestCase):
    def test_the_delta_is_what_gets_reported(self):
        # Lifetime counters. Reporting the total would fold warm-up and every earlier
        # concurrency level into this point's number.
        before = SpecCounters(accepted=700.0, draft=1000.0)
        after = SpecCounters(accepted=1600.0, draft=2000.0)
        window = after - before
        self.assertEqual(window.acceptance_rate, 0.9)
        self.assertNotEqual(window.acceptance_rate, after.acceptance_rate)

    def test_no_drafting_in_the_window_is_none(self):
        same = SpecCounters(accepted=10.0, draft=10.0)
        self.assertIsNone((same - same).acceptance_rate)


class MetricsUrl(unittest.TestCase):
    def test_strips_the_openai_prefix(self):
        # base_url points at the OpenAI-compatible surface; /metrics is at the server root.
        self.assertEqual(metrics_url("http://host:8001/v1"), "http://host:8001/metrics")
        self.assertEqual(metrics_url("http://host:8001/v1/"), "http://host:8001/metrics")

    def test_leaves_a_bare_root_alone(self):
        self.assertEqual(metrics_url("http://host:8001"), "http://host:8001/metrics")


if __name__ == "__main__":
    unittest.main()