from __future__ import annotations

import json
from pathlib import Path
import time

import pytest
import numpy as np
from prometheus_client import CollectorRegistry, generate_latest

import llm_runtime_metrics
from llm_runtime_metrics import (
    REQUEST_FEATURE_IMAGE,
    REQUEST_FEATURE_TOOLS,
    RequestMetricsCollector,
    RequestMetricsFactory,
    prometheus_strfmt,
)


def test_root_api_hides_low_level_registry() -> None:
    assert not hasattr(llm_runtime_metrics, "PerformanceMetricsRegistry")


def test_request_metrics_factory_unified_e2e(tmp_path: Path) -> None:
    factory = RequestMetricsFactory(
        request_log_enabled=True,
        request_log_batch_size=1,
        local_file=True,
        request_log_output_dir=str(tmp_path),
        s3=False,
        request_log_hash_salt="fixed-salt",
    )

    req = factory.new_request(list(range(128)))
    req.record_tokens([128, 129, 130], cached_tokens=64, is_diff=True)
    req.success()

    hashes = req.total_hashes()
    assert isinstance(hashes, list)
    assert len(hashes) > 0

    snap = factory.snapshot()
    series = snap.series()
    totals = snap.totals()
    ttft = find_metric_snapshot(series, "request_metrics_ttft_ms", feature="all")
    kv = find_metric_snapshot(series, "request_metrics_kv_cache_hit_rate")
    itl = find_metric_snapshot(series, "request_metrics_itl_ms", feature="all")
    successful = find_metric_snapshot(series, "request_metrics_successful_request")
    assert ttft.average is not None
    assert kv.ratio == pytest.approx(0.5, rel=1e-9)
    assert itl.average is not None
    assert successful.rate_per_second is not None
    assert any(
        item.name == "request_metrics_ttft_ms" and item.labels["feature"] == "all"
        for item in series
    )

    samples = factory.prometheus_samples({"model": "demo"})
    names = {sample.name for sample in samples}
    assert "request_metrics_ttft_ms_average" in names
    assert "request_metrics_kv_cache_hit_rate_ratio" in names
    assert "total_requests" in names
    assert all(sample.labels.get("model") == "demo" for sample in samples)

    assert totals.total_requests == 1
    assert totals.total_successful_requests == 1
    assert totals.total_input_tokens == 128
    assert totals.total_output_tokens == 3
    assert totals.total_kv_cache_hits_tokens == 64
    assert totals.total_kv_cache_lookup_tokens == 128
    assert totals.kv_cache_hit_rate_lifetime == pytest.approx(0.5, rel=1e-9)

    files = []
    for _ in range(50):
        files = list(tmp_path.glob("*.json"))
        if files:
            break
        time.sleep(0.01)
    assert files, "expected at least one request log JSON file"


def test_request_metrics_factory_does_not_log_by_default(tmp_path: Path) -> None:
    factory = RequestMetricsFactory(request_log_output_dir=str(tmp_path))
    req = factory.new_request(list(range(16)))
    req.record_tokens([16, 17], cached_tokens=8, is_diff=True)
    req.success()

    time.sleep(0.05)
    assert not list(tmp_path.glob("*.json"))


def test_request_metrics_request_log_session_id_is_optional(tmp_path: Path) -> None:
    factory = RequestMetricsFactory(
        request_log_enabled=True,
        request_log_batch_size=1,
        local_file=True,
        request_log_output_dir=str(tmp_path),
        s3=False,
        request_log_hash_salt="fixed-salt",
    )

    with_session = factory.new_request(list(range(16)), session_id="session-123")
    with_session.record_tokens([16, 17], cached_tokens=8, is_diff=True)
    with_session.success()

    without_session = factory.new_request(list(range(16)))
    without_session.record_tokens([16, 17], cached_tokens=8, is_diff=True)
    without_session.success()

    files = []
    for _ in range(50):
        files = list(tmp_path.glob("*.json"))
        if len(files) >= 2:
            break
        time.sleep(0.01)
    assert len(files) >= 2, "expected request log JSON files for both requests"

    entries = []
    for path in files:
        entries.extend(json.loads(path.read_text()))

    with_session_entry = next(entry for entry in entries if entry["input_tokens"] == 16 and "session_id" in entry)
    without_session_entry = next(entry for entry in entries if entry["input_tokens"] == 16 and "session_id" not in entry)

    assert with_session_entry["session_id"] != "session-123"
    assert len(with_session_entry["session_id"]) == 64
    assert "session_id" not in without_session_entry


def test_prometheus_collector_and_strfmt() -> None:
    factory = RequestMetricsFactory(request_log_enabled=False)
    req = factory.new_request(list(range(128)))
    req.record_tokens([128, 129, 130], cached_tokens=64, is_diff=True)
    req.success()

    registry = CollectorRegistry()
    RequestMetricsCollector(factory, base_labels={"engine": "vllm"}, registry=registry)
    text = generate_latest(registry).decode("utf-8")
    assert "request_metrics_ttft_ms_average" in text
    assert 'engine="vllm"' in text
    assert "total_requests_total" in text

    strfmt = factory.prometheus_strfmt({"engine": "vllm"})
    assert "# TYPE request_metrics_ttft_ms_average gauge" in strfmt
    assert 'request_metrics_ttft_ms_average{engine="vllm",feature="all"}' in strfmt
    assert "# TYPE total_requests counter" in strfmt

    strfmt_from_helper = prometheus_strfmt(factory, base_labels={"engine": "vllm"})
    assert strfmt_from_helper == strfmt


def test_prometheus_strfmt_sanitizes_metric_names() -> None:
    factory = RequestMetricsFactory(metric_prefix="req-metrics", request_log_enabled=False)
    req = factory.new_request(list(range(16)))
    req.record_tokens([16, 17], cached_tokens=8, is_diff=True)
    req.success()

    strfmt = factory.prometheus_strfmt()
    assert "req_metrics_ttft_ms_average" in strfmt
    assert "req-metrics" not in strfmt


def test_prometheus_strfmt_escapes_label_values() -> None:
    factory = RequestMetricsFactory(request_log_enabled=False)
    req = factory.new_request(list(range(16)))
    req.record_tokens([16, 17], cached_tokens=8, is_diff=True)
    req.success()

    strfmt = factory.prometheus_strfmt({"tag": 'a"b\nc\\d'})
    assert 'tag="a\\\"b\\nc\\\\d"' in strfmt


def test_prometheus_strfmt_default_cardinality_budget_and_kv_percentiles() -> None:
    factory = RequestMetricsFactory(request_log_enabled=False)
    req = factory.new_request(list(range(16)))
    req.record_tokens([16, 17], cached_tokens=8, is_diff=True)
    req.success()

    strfmt = factory.prometheus_strfmt()
    sample_lines = [
        line for line in strfmt.splitlines() if line and not line.startswith("#")
    ]
    assert len(sample_lines) <= 70

    kv_quantile_lines = [
        line
        for line in sample_lines
        if line.startswith("request_metrics_kv_cache_hit_rate_quantile{")
    ]
    assert len(kv_quantile_lines) == 10
    for quantile in ["0.01", "0.05", "0.1", "0.25", "0.5", "0.75", "0.9", "0.95", "0.99", "0.999"]:
        assert any(f'quantile="{quantile}"' in line for line in kv_quantile_lines)


def test_request_metrics_feature_bitmask_records_presence_series() -> None:
    factory = RequestMetricsFactory(request_log_enabled=False)
    req = factory.new_request(
        list(range(16)),
        features=REQUEST_FEATURE_TOOLS | REQUEST_FEATURE_IMAGE,
    )
    req.record_tokens([16, 17], cached_tokens=8, is_diff=True)
    req.record_tokens([18], cached_tokens=None, is_diff=True)
    req.success()

    assert req.features() == REQUEST_FEATURE_TOOLS | REQUEST_FEATURE_IMAGE

    series = factory.snapshot().series()
    assert find_metric_snapshot(series, "request_metrics_ttft_ms", feature="tools").average is not None
    assert find_metric_snapshot(series, "request_metrics_ttft_ms", feature="image").average is not None
    assert find_metric_snapshot(series, "request_metrics_ttft_ms", feature="xgrammar").average == pytest.approx(0.0)
    assert find_metric_snapshot(series, "request_metrics_itl_ms", feature="tools").average is not None
    assert find_metric_snapshot(series, "request_metrics_itl_ms", feature="image").average is not None
    assert find_metric_snapshot(series, "request_metrics_itl_ms", feature="xgrammar").average == pytest.approx(0.0)


@pytest.mark.parametrize(
    "input_ids, output_ids",
    [
        (list(range(128)), [128, 129, 130]),
        (np.arange(128, dtype=np.uint32), np.array([128, 129, 130], dtype=np.uint32)),
    ],
)
def test_request_metrics_accepts_list_and_numpy_inputs(
    input_ids, output_ids
) -> None:
    factory = RequestMetricsFactory(request_log_enabled=False)
    req = factory.new_request(input_ids)
    req.record_tokens(output_ids, cached_tokens=64, is_diff=True)
    req.success()

    kv = find_metric_snapshot(factory.snapshot().series(), "request_metrics_kv_cache_hit_rate")
    assert kv.ratio == pytest.approx(0.5, rel=1e-9)


def find_metric_snapshot(series, name: str, *, feature: str | None = None):
    for item in series:
        if item.name != name:
            continue
        if feature is None and "feature" in item.labels:
            continue
        if feature is not None and item.labels.get("feature") != feature:
            continue
        return item.snapshot()
    raise AssertionError(f"missing metric series {name!r} feature={feature!r}")
