from __future__ import annotations

from collections.abc import Mapping, Sequence
from typing import Final

REQUEST_FEATURE_NONE: Final[int]
REQUEST_FEATURE_XGRAMMAR: Final[int]
REQUEST_FEATURE_TOOLS: Final[int]
REQUEST_FEATURE_IMAGE: Final[int]
__version__: Final[str]

class RateMetric:
    def record_count(self, count: int = 1) -> None: ...

class DistributionMetric:
    def record_value(self, value: float) -> None: ...

class RatioMetric:
    def record_ratio(self, numerator: float, denominator: float) -> None: ...

class PerformanceMetricsRegistry:
    def __init__(self) -> None: ...
    def new_rate_metric(
        self,
        name: str,
        quantiles: Sequence[float] | None = None,
        sample_period_seconds: float = 1.0,
        window_seconds: float | None = None,
    ) -> RateMetric: ...
    def new_distribution_metric(
        self,
        name: str,
        quantiles: Sequence[float] | None = None,
        window_seconds: float | None = None,
    ) -> DistributionMetric: ...
    def new_ratio_metric(
        self,
        name: str,
        quantiles: Sequence[float] | None = None,
        window_seconds: float | None = None,
    ) -> RatioMetric: ...
    def snapshot_all(self) -> dict[str, MetricSnapshot]: ...
    def snapshot_all_series(self) -> list[MetricSeriesSnapshot]: ...
    def unregister_metric(self, name: str) -> None: ...

class MetricSnapshot:
    kind: str
    rate_per_second: float | None
    average: float | None
    quantiles: dict[str, float]
    numerator_sum: float | None
    denominator_sum: float | None
    ratio: float | None

class RequestMetricsTotalsSnapshot:
    total_requests: int
    total_successful_requests: int
    total_pre_first_token_cancellations: int
    total_pre_20_token_cancellations: int
    total_mid_stream_cancellations: int
    total_input_tokens: int
    total_output_tokens: int
    total_net_new_input_tokens: int
    total_kv_cache_hits_tokens: int
    total_kv_cache_lookup_tokens: int
    total_itl_samples_recorded: int
    kv_cache_hit_rate_lifetime: float | None

class MetricSeriesSnapshot:
    name: str
    labels: dict[str, str]
    def snapshot(self) -> MetricSnapshot: ...

class RequestMetricsSnapshot:
    def series(self) -> list[MetricSeriesSnapshot]: ...
    def totals(self) -> RequestMetricsTotalsSnapshot: ...

class PrometheusSample:
    name: str
    value: float
    labels: dict[str, str]
    metric_type: str

class CompactResult:
    source_objects: int
    rows_written: int
    output_path: str
    min_timestamp: int

class RequestMetricsFactory:
    def __init__(
        self,
        block_size: int = 64,
        request_log_enabled: bool = False,
        request_log_hash_salt: str | None = None,
        request_log_batch_size: int = 2048,
        request_log_queue_capacity: int = 4096,
        local_file: bool = True,
        request_log_output_dir: str = "/tmp/records",
        s3: bool = False,
        s3_bucket: str | None = None,
        s3_region: str = "us-east-1",
        s3_endpoint: str | None = None,
        s3_access_key_id: str | None = None,
        s3_secret_access_key: str | None = None,
        s3_prefix: str = "records",
        file_format: str = "json",
        org_id: str = "",
        metric_prefix: str = "request_metrics",
        metrics_window_seconds: float = 60.0,
        metrics_quantiles: Sequence[float] = (0.5, 0.9, 0.99),
        itl_sample_rate: float = 0.05,
    ) -> None: ...
    def new_request(
        self,
        input_token_ids: Sequence[int],
        features: int = 0,
        session_id: str | None = None,
    ) -> RequestMetrics: ...
    def snapshot(self) -> RequestMetricsSnapshot: ...
    def prometheus_samples(
        self,
        base_labels: Mapping[str, str] | None = None,
    ) -> list[PrometheusSample]: ...
    def prometheus_strfmt(
        self, base_labels: Mapping[str, str] | None = None
    ) -> str: ...

class RequestMetrics:
    def record_tokens(
        self,
        token_ids: Sequence[int],
        cached_tokens: int | None = None,
        is_diff: bool = True,
        speculation_ratio: float | None = None,
    ) -> None: ...
    def success(self) -> None: ...
    def cancel(self) -> None: ...
    def total_hashes(self) -> list[bytes]: ...
    def features(self) -> int: ...

def compact_s3_parquet_to_local(
    s3_bucket: str,
    output_path: str,
    s3_region: str = "us-east-1",
    s3_prefix: str = "",
    s3_endpoint: str | None = None,
    s3_access_key_id: str | None = None,
    s3_secret_access_key: str | None = None,
    max_concurrent_downloads: int = 128,
    anonymize: bool = False,
) -> CompactResult: ...
def anonymize_parquet(
    input_path: str,
    output_path: str,
    columns_to_pseudonymize: Sequence[str] = [
        "provided_session_id",
        "org_id",
        "total_hashes",
    ],
    columns_to_drop: Sequence[str] = ["org_id"],
    min_timestamp: int = 0,
) -> int: ...
