use crate::metrics::{
    DistributionMetricHandle, MetricLabels, MetricSeriesSnapshot, PerformanceMetricsRegistry,
    RateMetricHandle, RatioMetricHandle,
};
use blake3::Hasher;
use bytes::Bytes;
use object_store::aws::AmazonS3Builder;
use object_store::path::Path as ObjectPath;
use object_store::{ObjectStore, PutOptions, PutPayload};
use rand::rngs::SmallRng;
use rand::{Rng, SeedableRng};
use serde::{Deserialize, Serialize};
use std::fmt;
use std::fs;
use std::path::PathBuf;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, Mutex, OnceLock};
use std::time::{Instant, SystemTime, UNIX_EPOCH};
use tokio::runtime::Builder as RuntimeBuilder;
use tokio::sync::mpsc;

const REQUEST_TRACE_HASH_DOMAIN: &[u8] = b"pm-request-trace-v2";
const REQUEST_SESSION_HASH_DOMAIN: &[u8] = b"pm-request-session-v1";

#[derive(Debug, Clone)]
pub struct RequestTraceOptions {
    pub block_size: usize,
    pub request_log_enabled: bool,
    pub request_log_hash_salt: Option<String>,
    pub request_log_batch_size: usize,
    pub request_log_queue_capacity: usize,
    pub local_output_dir: Option<PathBuf>,
    pub s3_config: Option<S3SinkConfig>,
}

#[derive(Debug, Clone)]
pub struct S3SinkConfig {
    pub bucket: String,
    pub region: String,
    pub endpoint: Option<String>,
    pub access_key_id: Option<String>,
    pub secret_access_key: Option<String>,
    pub prefix: String,
    pub allow_http: bool,
}

impl Default for RequestTraceOptions {
    fn default() -> Self {
        Self {
            block_size: 128,
            request_log_enabled: false,
            request_log_hash_salt: None,
            request_log_batch_size: 256,
            request_log_queue_capacity: 4096,
            local_output_dir: Some(PathBuf::from("/tmp/records")),
            s3_config: None,
        }
    }
}

#[derive(Debug, Clone)]
pub struct RequestMetricsOptions {
    pub trace: RequestTraceOptions,
    pub metric_prefix: String,
    pub request_quantiles: Vec<f64>,
    pub request_sample_period_seconds: Option<f64>,
    pub request_window_seconds: Option<f64>,
    pub input_tokens_quantiles: Vec<f64>,
    pub input_tokens_sample_period_seconds: Option<f64>,
    pub input_tokens_window_seconds: Option<f64>,
    pub ttft_quantiles: Vec<f64>,
    pub ttft_window_seconds: Option<f64>,
    pub ttft_per_input_token_quantiles: Vec<f64>,
    pub ttft_per_input_token_window_seconds: Option<f64>,
    pub itl_quantiles: Vec<f64>,
    pub itl_window_seconds: Option<f64>,
    pub pre_first_token_cancellation_quantiles: Vec<f64>,
    pub pre_first_token_cancellation_sample_period_seconds: Option<f64>,
    pub pre_first_token_cancellation_window_seconds: Option<f64>,
    pub mid_stream_cancellation_quantiles: Vec<f64>,
    pub mid_stream_cancellation_sample_period_seconds: Option<f64>,
    pub mid_stream_cancellation_window_seconds: Option<f64>,
    pub successful_request_quantiles: Vec<f64>,
    pub successful_request_sample_period_seconds: Option<f64>,
    pub successful_request_window_seconds: Option<f64>,
    pub kv_cache_hit_rate_quantiles: Vec<f64>,
    pub kv_cache_hit_rate_window_seconds: Option<f64>,
    pub itl_sample_rate: f64,
}

impl Default for RequestMetricsOptions {
    fn default() -> Self {
        Self {
            trace: RequestTraceOptions::default(),
            metric_prefix: "request_metrics".to_string(),
            request_quantiles: vec![0.1, 0.5, 0.9],
            request_sample_period_seconds: Some(1.0),
            request_window_seconds: None,
            input_tokens_quantiles: vec![0.01, 0.1, 0.2, 0.35, 0.5, 0.65, 0.8, 0.9, 0.99],
            input_tokens_sample_period_seconds: Some(1.0),
            input_tokens_window_seconds: None,
            ttft_quantiles: vec![0.01, 0.1, 0.5, 0.9, 0.99],
            ttft_window_seconds: None,
            ttft_per_input_token_quantiles: vec![0.01, 0.1, 0.5, 0.8, 0.9, 0.99],
            ttft_per_input_token_window_seconds: None,
            itl_quantiles: vec![0.01, 0.05, 0.1, 0.2, 0.5, 0.8, 0.9, 0.95, 0.99, 0.999],
            itl_window_seconds: None,
            pre_first_token_cancellation_quantiles: vec![],
            pre_first_token_cancellation_sample_period_seconds: Some(1.0),
            pre_first_token_cancellation_window_seconds: None,
            mid_stream_cancellation_quantiles: vec![],
            mid_stream_cancellation_sample_period_seconds: Some(1.0),
            mid_stream_cancellation_window_seconds: None,
            successful_request_quantiles: vec![],
            successful_request_sample_period_seconds: Some(1.0),
            successful_request_window_seconds: None,
            kv_cache_hit_rate_quantiles: vec![
                0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99, 0.999,
            ],
            kv_cache_hit_rate_window_seconds: None,
            itl_sample_rate: 0.05,
        }
    }
}

#[derive(Debug, Clone, Serialize, Deserialize)]
pub struct RequestLogEntry {
    pub timestamp_start_unix_ms: u64,
    pub timestamp_end_unix_ms: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    pub total_hashes: Vec<String>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub session_id: Option<String>,
    pub request_canceled: bool,
}

#[derive(Clone)]
pub struct RequestTraceFactory {
    recorder: Option<Arc<EventRecorder>>,
    hash_key: Option<[u8; 32]>,
    session_hash_key: Option<[u8; 32]>,
    block_size: usize,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash, Serialize, Deserialize, Default)]
pub struct RequestFeatures(u8);

#[derive(Debug, Clone, Copy, PartialEq, Eq, Hash)]
enum RequestFeatureKind {
    XGrammar,
    Tools,
    Image,
}

impl RequestFeatureKind {
    const ALL: [Self; 3] = [Self::XGrammar, Self::Tools, Self::Image];

    const fn bit(self) -> u8 {
        match self {
            Self::XGrammar => 1 << 0,
            Self::Tools => 1 << 1,
            Self::Image => 1 << 2,
        }
    }

    const fn metric_suffix(self) -> &'static str {
        match self {
            Self::XGrammar => "xgrammar",
            Self::Tools => "tools",
            Self::Image => "image",
        }
    }
}

impl RequestFeatures {
    pub const NONE: Self = Self(0);
    pub const XGRAMMAR: Self = Self(RequestFeatureKind::XGrammar.bit());
    pub const TOOLS: Self = Self(RequestFeatureKind::Tools.bit());
    pub const IMAGE: Self = Self(RequestFeatureKind::Image.bit());

    pub const fn empty() -> Self {
        Self::NONE
    }

    pub const fn bits(self) -> u8 {
        self.0
    }

    pub const fn is_empty(self) -> bool {
        self.0 == 0
    }

    pub const fn contains(self, other: Self) -> bool {
        (self.0 & other.0) == other.0
    }

    pub const fn from_bools(xgrammar: bool, tools: bool, image: bool) -> Self {
        let mut bits = 0;
        if xgrammar {
            bits |= Self::XGRAMMAR.0;
        }
        if tools {
            bits |= Self::TOOLS.0;
        }
        if image {
            bits |= Self::IMAGE.0;
        }
        Self(bits)
    }

    fn includes(self, feature: RequestFeatureKind) -> bool {
        self.0 & feature.bit() != 0
    }
}

impl fmt::Display for RequestFeatures {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        if self.is_empty() {
            return f.write_str("none");
        }
        let mut first = true;
        for feature in RequestFeatureKind::ALL {
            if !self.includes(feature) {
                continue;
            }
            if !first {
                f.write_str(",")?;
            }
            f.write_str(feature.metric_suffix())?;
            first = false;
        }
        Ok(())
    }
}

impl std::ops::BitOr for RequestFeatures {
    type Output = Self;

    fn bitor(self, rhs: Self) -> Self::Output {
        Self(self.0 | rhs.0)
    }
}

impl std::ops::BitOrAssign for RequestFeatures {
    fn bitor_assign(&mut self, rhs: Self) {
        self.0 |= rhs.0;
    }
}

impl From<u8> for RequestFeatures {
    fn from(bits: u8) -> Self {
        Self(bits & (Self::XGRAMMAR.0 | Self::TOOLS.0 | Self::IMAGE.0))
    }
}

#[derive(Clone)]
struct RequestFeatureMetricConfig {
    metric_prefix: String,
    ttft_quantiles: Vec<f64>,
    ttft_window_seconds: Option<f64>,
    itl_quantiles: Vec<f64>,
    itl_window_seconds: Option<f64>,
}

#[derive(Clone)]
struct RequestFeatureMetricHandles {
    ttft_by_feature: Vec<(RequestFeatureKind, DistributionMetricHandle)>,
    itl_by_feature: Vec<(RequestFeatureKind, DistributionMetricHandle)>,
}

#[derive(Clone)]
pub struct RequestMetricsFactory {
    trace_factory: RequestTraceFactory,
    metrics: PerformanceMetricsRegistry,
    request: RateMetricHandle,
    input_tokens: RateMetricHandle,
    net_new_input_tokens: RateMetricHandle,
    output_tokens: RateMetricHandle,
    ttft_ms: DistributionMetricHandle,
    ttft_ms_per_input_token: DistributionMetricHandle,
    ttft_ms_per_net_new_input_token: DistributionMetricHandle,
    itl_ms: DistributionMetricHandle,
    pre_first_token_cancellation: RateMetricHandle,
    mid_stream_cancellation: RateMetricHandle,
    successful_request: RateMetricHandle,
    kv_cache_hit_rate: RatioMetricHandle,
    itl_sample_rate: f64,
    totals: Arc<RequestMetricsTotals>,
    feature_metric_config: RequestFeatureMetricConfig,
    feature_metrics: Arc<OnceLock<RequestFeatureMetricHandles>>,
}

pub struct RequestTrace {
    factory: RequestTraceFactory,
    started_at_wallclock: SystemTime,
    input_tokens: usize,
    output_tokens_len: usize,
    total_hashes: Vec<String>,
    session_id: Option<String>,
    previous_hash: Option<[u8; 32]>,
    pending: Vec<u32>,
    terminal: bool,
}

pub struct RequestMetricsRequest {
    trace: RequestTrace,
    started_at: Instant,
    input_tokens: u64,
    features: RequestFeatures,
    last_token_at: Option<Instant>,
    last_total_tokens: u64,
    cached_tokens_hint: Option<u64>,
    rng: SmallRng,
    input_tokens_rate: RateMetricHandle,
    net_new_input_tokens_rate: RateMetricHandle,
    output_tokens_rate: RateMetricHandle,
    ttft_ms: DistributionMetricHandle,
    ttft_ms_per_input_token: DistributionMetricHandle,
    ttft_ms_per_net_new_input_token: DistributionMetricHandle,
    itl_ms: DistributionMetricHandle,
    pre_first_token_cancellation: RateMetricHandle,
    mid_stream_cancellation: RateMetricHandle,
    successful_request: RateMetricHandle,
    kv_cache_hit_rate: RatioMetricHandle,
    itl_sample_rate: f64,
    totals: Arc<RequestMetricsTotals>,
    ttft_feature_metrics: Vec<DistributionMetricHandle>,
    itl_feature_metrics: Vec<DistributionMetricHandle>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct RequestMetricsTotalsSnapshot {
    pub total_requests: u64,
    pub total_successful_requests: u64,
    pub total_pre_first_token_cancellations: u64,
    pub total_mid_stream_cancellations: u64,
    pub total_input_tokens: u64,
    pub total_output_tokens: u64,
    pub total_net_new_input_tokens: u64,
    pub total_kv_cache_hits_tokens: u64,
    pub total_kv_cache_lookup_tokens: u64,
    pub total_itl_samples_recorded: u64,
    pub kv_cache_hit_rate_lifetime: Option<f64>,
}

#[derive(Debug, Serialize, Deserialize, Clone)]
pub struct RequestMetricsSnapshot {
    pub series: Vec<MetricSeriesSnapshot>,
    pub totals: RequestMetricsTotalsSnapshot,
}

#[derive(Debug, Default)]
struct RequestMetricsTotals {
    total_requests: AtomicU64,
    total_successful_requests: AtomicU64,
    total_pre_first_token_cancellations: AtomicU64,
    total_mid_stream_cancellations: AtomicU64,
    total_input_tokens: AtomicU64,
    total_output_tokens: AtomicU64,
    total_net_new_input_tokens: AtomicU64,
    total_kv_cache_hits_tokens: AtomicU64,
    total_kv_cache_lookup_tokens: AtomicU64,
    total_itl_samples_recorded: AtomicU64,
}

impl RequestTraceFactory {
    pub fn new(options: RequestTraceOptions) -> anyhow::Result<Self> {
        let recorder = if options.request_log_enabled {
            let sink = build_event_recorder_sink(
                options.local_output_dir.clone(),
                options.s3_config.clone(),
            )?;
            Some(Arc::new(EventRecorder::new(
                sink,
                options.request_log_batch_size,
                options.request_log_queue_capacity,
            )))
        } else {
            None
        };

        Ok(Self {
            recorder,
            hash_key: options
                .request_log_hash_salt
                .as_deref()
                .map(|salt| *blake3::hash(salt.as_bytes()).as_bytes()),
            session_hash_key: options.request_log_hash_salt.as_deref().map(|salt| {
                let base_key = blake3::hash(salt.as_bytes());
                *blake3::hash(base_key.as_bytes()).as_bytes()
            }),
            block_size: options.block_size.max(1),
        })
    }

    pub fn new_trace(&self, input_token_ids: &[u32], session_id: Option<&str>) -> RequestTrace {
        let mut trace = RequestTrace {
            factory: self.clone(),
            started_at_wallclock: SystemTime::now(),
            input_tokens: input_token_ids.len(),
            output_tokens_len: 0,
            total_hashes: Vec::new(),
            session_id: session_id.map(|id| hash_session_id(id, self.session_hash_key.as_ref())),
            previous_hash: None,
            pending: Vec::with_capacity(self.block_size),
            terminal: false,
        };
        trace.append_tokens(input_token_ids);
        trace
    }
}

impl RequestMetricsFactory {
    pub fn new(mut options: RequestMetricsOptions) -> anyhow::Result<Self> {
        options.itl_sample_rate = options.itl_sample_rate.clamp(0.0, 1.0);
        let trace_factory = RequestTraceFactory::new(options.trace)?;
        let metrics = PerformanceMetricsRegistry::default();
        let prefix = options.metric_prefix.as_str();

        let request = metrics.new_rate_metric(
            format!("{prefix}_request"),
            options.request_quantiles.clone(),
            options.request_sample_period_seconds,
            options.request_window_seconds,
        )?;
        let input_tokens = metrics.new_rate_metric(
            format!("{prefix}_input_tokens"),
            options.input_tokens_quantiles.clone(),
            options.input_tokens_sample_period_seconds,
            options.input_tokens_window_seconds,
        )?;
        let net_new_input_tokens = metrics.new_rate_metric(
            format!("{prefix}_net_new_input_tokens"),
            options.input_tokens_quantiles.clone(),
            options.input_tokens_sample_period_seconds,
            options.input_tokens_window_seconds,
        )?;
        let output_tokens = metrics.new_rate_metric(
            format!("{prefix}_output_tokens"),
            options.input_tokens_quantiles.clone(),
            options.input_tokens_sample_period_seconds,
            options.input_tokens_window_seconds,
        )?;
        let ttft_ms = metrics.new_distribution_metric_with_labels(
            format!("{prefix}_ttft_ms"),
            feature_metric_labels("all"),
            options.ttft_quantiles.clone(),
            options.ttft_window_seconds,
        )?;
        let ttft_ms_per_input_token = metrics.new_distribution_metric(
            format!("{prefix}_ttft_ms_per_input_token"),
            options.ttft_per_input_token_quantiles.clone(),
            options.ttft_per_input_token_window_seconds,
        )?;
        let ttft_ms_per_net_new_input_token = metrics.new_distribution_metric(
            format!("{prefix}_ttft_ms_per_net_new_input_token"),
            options.ttft_per_input_token_quantiles.clone(),
            options.ttft_per_input_token_window_seconds,
        )?;
        let itl_ms = metrics.new_distribution_metric_with_labels(
            format!("{prefix}_itl_ms"),
            feature_metric_labels("all"),
            options.itl_quantiles.clone(),
            options.itl_window_seconds,
        )?;
        let pre_first_token_cancellation = metrics.new_rate_metric(
            format!("{prefix}_pre_first_token_cancellation"),
            options.pre_first_token_cancellation_quantiles.clone(),
            options.pre_first_token_cancellation_sample_period_seconds,
            options.pre_first_token_cancellation_window_seconds,
        )?;
        let mid_stream_cancellation = metrics.new_rate_metric(
            format!("{prefix}_mid_stream_cancellation"),
            options.mid_stream_cancellation_quantiles.clone(),
            options.mid_stream_cancellation_sample_period_seconds,
            options.mid_stream_cancellation_window_seconds,
        )?;
        let successful_request = metrics.new_rate_metric(
            format!("{prefix}_successful_request"),
            options.successful_request_quantiles.clone(),
            options.successful_request_sample_period_seconds,
            options.successful_request_window_seconds,
        )?;
        let kv_cache_hit_rate = metrics.new_ratio_metric(
            format!("{prefix}_kv_cache_hit_rate"),
            options.kv_cache_hit_rate_quantiles,
            options.kv_cache_hit_rate_window_seconds,
        )?;

        Ok(Self {
            trace_factory,
            metrics,
            request,
            input_tokens,
            net_new_input_tokens,
            output_tokens,
            ttft_ms,
            ttft_ms_per_input_token,
            ttft_ms_per_net_new_input_token,
            itl_ms,
            pre_first_token_cancellation,
            mid_stream_cancellation,
            successful_request,
            kv_cache_hit_rate,
            itl_sample_rate: options.itl_sample_rate,
            totals: Arc::new(RequestMetricsTotals::default()),
            feature_metric_config: RequestFeatureMetricConfig {
                metric_prefix: options.metric_prefix,
                ttft_quantiles: options.ttft_quantiles,
                ttft_window_seconds: options.ttft_window_seconds,
                itl_quantiles: options.itl_quantiles,
                itl_window_seconds: options.itl_window_seconds,
            },
            feature_metrics: Arc::new(OnceLock::new()),
        })
    }

    pub fn new_request(
        &self,
        input_token_ids: &[u32],
        features: RequestFeatures,
        session_id: Option<&str>,
    ) -> RequestMetricsRequest {
        let _ = self.request.record_count(1);
        self.totals.total_requests.fetch_add(1, Ordering::Relaxed);
        self.totals
            .total_input_tokens
            .fetch_add(input_token_ids.len() as u64, Ordering::Relaxed);
        let (ttft_feature_metrics, itl_feature_metrics) = self.feature_metric_handles(features);
        RequestMetricsRequest {
            trace: self.trace_factory.new_trace(input_token_ids, session_id),
            started_at: Instant::now(),
            input_tokens: input_token_ids.len() as u64,
            features,
            last_token_at: None,
            last_total_tokens: 0,
            cached_tokens_hint: None,
            rng: SmallRng::seed_from_u64(unix_ms(SystemTime::now())),
            input_tokens_rate: self.input_tokens.clone(),
            net_new_input_tokens_rate: self.net_new_input_tokens.clone(),
            output_tokens_rate: self.output_tokens.clone(),
            ttft_ms: self.ttft_ms.clone(),
            ttft_ms_per_input_token: self.ttft_ms_per_input_token.clone(),
            ttft_ms_per_net_new_input_token: self.ttft_ms_per_net_new_input_token.clone(),
            itl_ms: self.itl_ms.clone(),
            pre_first_token_cancellation: self.pre_first_token_cancellation.clone(),
            mid_stream_cancellation: self.mid_stream_cancellation.clone(),
            successful_request: self.successful_request.clone(),
            kv_cache_hit_rate: self.kv_cache_hit_rate.clone(),
            itl_sample_rate: self.itl_sample_rate,
            totals: Arc::clone(&self.totals),
            ttft_feature_metrics,
            itl_feature_metrics,
        }
    }

    fn feature_metric_handles(
        &self,
        features: RequestFeatures,
    ) -> (Vec<DistributionMetricHandle>, Vec<DistributionMetricHandle>) {
        if features.is_empty() {
            return (Vec::new(), Vec::new());
        }

        let handles = self.feature_metrics.get_or_init(|| {
            RequestFeatureMetricHandles::new(&self.metrics, &self.feature_metric_config)
                .expect("failed to register request feature metrics")
        });
        handles.for_features(features)
    }

    pub fn snapshot(&self) -> RequestMetricsSnapshot {
        RequestMetricsSnapshot {
            series: self.metrics.snapshot_all_series(),
            totals: self.totals.snapshot(),
        }
    }
}

impl RequestMetricsRequest {
    pub fn record_tokens(&mut self, token_ids: &[u32], cached_tokens: Option<u64>, is_diff: bool) {
        if let Some(cached) = cached_tokens {
            self.cached_tokens_hint.get_or_insert(cached);
        }

        let appended = self.trace.record_output_tokens(token_ids, is_diff);
        if appended == 0 {
            return;
        }
        let _ = self.output_tokens_rate.record_count(appended as u64);
        self.totals
            .total_output_tokens
            .fetch_add(appended as u64, Ordering::Relaxed);

        let total_tokens = self.trace.output_tokens() as u64;
        if total_tokens <= self.last_total_tokens {
            return;
        }
        let now = Instant::now();

        if self.last_token_at.is_none() {
            if self.input_tokens > 0 {
                let _ = self.input_tokens_rate.record_count(self.input_tokens);
            }

            let ttft_ms = now.duration_since(self.started_at).as_secs_f64() * 1000.0;
            let _ = self.ttft_ms.record_value(ttft_ms);
            for metric in &self.ttft_feature_metrics {
                let _ = metric.record_value(ttft_ms);
            }
            if self.input_tokens > 0 {
                let _ = self
                    .ttft_ms_per_input_token
                    .record_value(ttft_ms / self.input_tokens as f64);
            }
            if let Some(cached) = self.cached_tokens_hint.take() {
                let total = self.input_tokens as f64;
                if total > 0.0 {
                    let hits = (cached as f64).min(total);
                    let _ = self.kv_cache_hit_rate.record_ratio(hits, total);
                    self.totals
                        .total_kv_cache_hits_tokens
                        .fetch_add(hits as u64, Ordering::Relaxed);
                    self.totals
                        .total_kv_cache_lookup_tokens
                        .fetch_add(total as u64, Ordering::Relaxed);
                }
                let net_new = self.input_tokens.saturating_sub(cached);
                if net_new > 0 {
                    let _ = self.net_new_input_tokens_rate.record_count(net_new);
                    self.totals
                        .total_net_new_input_tokens
                        .fetch_add(net_new, Ordering::Relaxed);
                    let _ = self
                        .ttft_ms_per_net_new_input_token
                        .record_value(ttft_ms / net_new as f64);
                }
            }
            self.last_token_at = Some(now);
            self.last_total_tokens = total_tokens;
            return;
        }

        let delta_tokens = total_tokens.saturating_sub(self.last_total_tokens);
        if delta_tokens > 0 {
            let elapsed_ms = now
                .duration_since(self.last_token_at.unwrap_or(now))
                .as_secs_f64()
                * 1000.0;
            if elapsed_ms > 0.0 {
                let itl_ms = elapsed_ms / delta_tokens as f64;
                let samples = sampled_count(delta_tokens, self.itl_sample_rate, &mut self.rng);
                for _ in 0..samples {
                    let _ = self.itl_ms.record_value(itl_ms);
                    for metric in &self.itl_feature_metrics {
                        let _ = metric.record_value(itl_ms);
                    }
                }
                self.totals
                    .total_itl_samples_recorded
                    .fetch_add(samples, Ordering::Relaxed);
            }
        }
        self.last_token_at = Some(now);
        self.last_total_tokens = total_tokens;
    }

    pub fn success(&mut self) {
        if self.trace.is_terminal() {
            return;
        }
        let _ = self.successful_request.record_count(1);
        self.totals
            .total_successful_requests
            .fetch_add(1, Ordering::Relaxed);
        self.trace.success();
    }

    pub fn cancel(&mut self) {
        if self.trace.is_terminal() {
            return;
        }
        if self.last_token_at.is_some() {
            let _ = self.mid_stream_cancellation.record_count(1);
            self.totals
                .total_mid_stream_cancellations
                .fetch_add(1, Ordering::Relaxed);
        } else {
            let _ = self.pre_first_token_cancellation.record_count(1);
            self.totals
                .total_pre_first_token_cancellations
                .fetch_add(1, Ordering::Relaxed);
        }
        self.trace.cancel();
    }

    pub fn total_hashes(&self) -> &[String] {
        self.trace.total_hashes()
    }

    pub fn features(&self) -> RequestFeatures {
        self.features
    }
}

impl Drop for RequestMetricsRequest {
    fn drop(&mut self) {
        self.cancel();
    }
}

impl RequestFeatureMetricHandles {
    fn new(
        metrics: &PerformanceMetricsRegistry,
        config: &RequestFeatureMetricConfig,
    ) -> anyhow::Result<Self> {
        let mut ttft_by_feature = Vec::with_capacity(RequestFeatureKind::ALL.len());
        let mut itl_by_feature = Vec::with_capacity(RequestFeatureKind::ALL.len());

        for feature in RequestFeatureKind::ALL {
            ttft_by_feature.push((
                feature,
                metrics.new_distribution_metric_with_labels(
                    format!("{}_ttft_ms", config.metric_prefix),
                    feature_metric_labels(feature.metric_suffix()),
                    config.ttft_quantiles.clone(),
                    config.ttft_window_seconds,
                )?,
            ));
            itl_by_feature.push((
                feature,
                metrics.new_distribution_metric_with_labels(
                    format!("{}_itl_ms", config.metric_prefix),
                    feature_metric_labels(feature.metric_suffix()),
                    config.itl_quantiles.clone(),
                    config.itl_window_seconds,
                )?,
            ));
        }

        Ok(Self {
            ttft_by_feature,
            itl_by_feature,
        })
    }

    fn for_features(
        &self,
        features: RequestFeatures,
    ) -> (Vec<DistributionMetricHandle>, Vec<DistributionMetricHandle>) {
        let ttft = self
            .ttft_by_feature
            .iter()
            .filter(|(feature, _)| features.includes(*feature))
            .map(|(_, handle)| handle.clone())
            .collect();
        let itl = self
            .itl_by_feature
            .iter()
            .filter(|(feature, _)| features.includes(*feature))
            .map(|(_, handle)| handle.clone())
            .collect();
        (ttft, itl)
    }
}

fn feature_metric_labels(feature: &str) -> MetricLabels {
    let mut labels = MetricLabels::new();
    labels.insert("feature".to_string(), feature.to_string());
    labels
}

impl RequestTrace {
    pub fn record_output_tokens(&mut self, output_token_ids: &[u32], is_diff: bool) -> usize {
        if self.terminal {
            return 0;
        }
        let delta = if is_diff {
            output_token_ids
        } else {
            if output_token_ids.len() <= self.output_tokens_len {
                return 0;
            }
            &output_token_ids[self.output_tokens_len..]
        };
        self.append_tokens(delta);
        self.output_tokens_len += delta.len();
        delta.len()
    }

    pub fn success(&mut self) {
        if self.terminal {
            return;
        }
        self.flush_full_block();
        self.emit(false);
        self.terminal = true;
    }

    pub fn cancel(&mut self) {
        if self.terminal {
            return;
        }
        self.flush_full_block();
        self.emit(true);
        self.terminal = true;
    }

    pub fn total_hashes(&self) -> &[String] {
        self.total_hashes.as_slice()
    }

    pub fn input_tokens(&self) -> usize {
        self.input_tokens
    }

    pub fn output_tokens(&self) -> usize {
        self.output_tokens_len
    }

    pub fn is_terminal(&self) -> bool {
        self.terminal
    }

    fn append_tokens(&mut self, token_ids: &[u32]) {
        let mut remaining = token_ids;
        while !remaining.is_empty() {
            let needed = self.factory.block_size - self.pending.len();
            let take = needed.min(remaining.len());
            self.pending.extend_from_slice(&remaining[..take]);
            remaining = &remaining[take..];
            if self.pending.len() == self.factory.block_size {
                self.flush_full_block();
            }
        }
    }

    fn flush_full_block(&mut self) {
        if self.pending.is_empty() {
            return;
        }
        let digest = hash_token_block(
            self.pending.as_slice(),
            self.previous_hash.as_ref(),
            self.factory.hash_key.as_ref(),
        );
        self.previous_hash = Some(*digest.as_bytes());
        self.total_hashes.push(digest.to_hex().to_string());
        self.pending.clear();
    }

    fn emit(&self, canceled: bool) {
        let Some(recorder) = &self.factory.recorder else {
            return;
        };
        recorder.record(self.build_log_entry(canceled));
    }

    fn build_log_entry(&self, canceled: bool) -> RequestLogEntry {
        RequestLogEntry {
            timestamp_start_unix_ms: unix_ms(self.started_at_wallclock),
            timestamp_end_unix_ms: unix_ms(SystemTime::now()),
            input_tokens: self.input_tokens as u64,
            output_tokens: self.output_tokens_len as u64,
            total_hashes: self.total_hashes.clone(),
            session_id: self.session_id.clone(),
            request_canceled: canceled,
        }
    }
}

impl RequestMetricsTotals {
    fn snapshot(&self) -> RequestMetricsTotalsSnapshot {
        let total_requests = self.total_requests.load(Ordering::Relaxed);
        let total_successful_requests = self.total_successful_requests.load(Ordering::Relaxed);
        let total_pre_first_token_cancellations = self
            .total_pre_first_token_cancellations
            .load(Ordering::Relaxed);
        let total_mid_stream_cancellations =
            self.total_mid_stream_cancellations.load(Ordering::Relaxed);
        let total_input_tokens = self.total_input_tokens.load(Ordering::Relaxed);
        let total_output_tokens = self.total_output_tokens.load(Ordering::Relaxed);
        let total_net_new_input_tokens = self.total_net_new_input_tokens.load(Ordering::Relaxed);
        let total_kv_cache_hits_tokens = self.total_kv_cache_hits_tokens.load(Ordering::Relaxed);
        let total_kv_cache_lookup_tokens =
            self.total_kv_cache_lookup_tokens.load(Ordering::Relaxed);
        let total_itl_samples_recorded = self.total_itl_samples_recorded.load(Ordering::Relaxed);
        let kv_cache_hit_rate_lifetime = if total_kv_cache_lookup_tokens > 0 {
            Some(total_kv_cache_hits_tokens as f64 / total_kv_cache_lookup_tokens as f64)
        } else {
            None
        };
        RequestMetricsTotalsSnapshot {
            total_requests,
            total_successful_requests,
            total_pre_first_token_cancellations,
            total_mid_stream_cancellations,
            total_input_tokens,
            total_output_tokens,
            total_net_new_input_tokens,
            total_kv_cache_hits_tokens,
            total_kv_cache_lookup_tokens,
            total_itl_samples_recorded,
            kv_cache_hit_rate_lifetime,
        }
    }
}

impl Drop for RequestTrace {
    fn drop(&mut self) {
        if !self.terminal {
            self.cancel();
        }
    }
}

#[derive(Debug)]
struct EventRecorder {
    tx: mpsc::Sender<RecorderMessage>,
    worker: Mutex<Option<std::thread::JoinHandle<()>>>,
}

#[derive(Debug)]
struct EventRecorderSink {
    local_output_dir: Option<PathBuf>,
    s3: Option<S3Writer>,
}

#[derive(Debug)]
struct S3Writer {
    store: Arc<dyn ObjectStore>,
    prefix: String,
}

#[derive(Debug)]
struct EventRecorderInner {
    entries: Vec<RequestLogEntry>,
    sink: EventRecorderSink,
    batch_size: usize,
    file_seq: u64,
}

enum RecorderMessage {
    Record(RequestLogEntry),
    Shutdown,
}

impl EventRecorder {
    fn new(sink: EventRecorderSink, batch_size: usize, queue_capacity: usize) -> Self {
        let (tx, mut rx) = mpsc::channel::<RecorderMessage>(queue_capacity.max(1));
        let worker = std::thread::spawn(move || {
            let rt = RuntimeBuilder::new_current_thread()
                .enable_all()
                .build()
                .expect("failed to build recorder runtime");
            rt.block_on(async move {
                let mut inner = EventRecorderInner {
                    entries: Vec::new(),
                    sink,
                    batch_size: batch_size.max(1),
                    file_seq: 0,
                };
                while let Some(msg) = rx.recv().await {
                    match msg {
                        RecorderMessage::Record(entry) => {
                            inner.entries.push(entry);
                            if inner.entries.len() >= inner.batch_size
                                && let Err(err) = inner.flush_next_batch().await
                            {
                                eprintln!(
                                    "performance-metrics: dropping request-log batch after sink failure: {err}"
                                );
                            }
                        }
                        RecorderMessage::Shutdown => {
                            break;
                        }
                    }
                }
                while !inner.entries.is_empty() {
                    if let Err(err) = inner.flush_next_batch().await {
                        eprintln!(
                            "performance-metrics: dropping request-log batch after sink failure: {err}"
                        );
                    }
                }
            });
        });

        Self {
            tx,
            worker: Mutex::new(Some(worker)),
        }
    }

    fn record(&self, entry: RequestLogEntry) {
        match self.tx.try_send(RecorderMessage::Record(entry)) {
            Ok(()) => {}
            Err(tokio::sync::mpsc::error::TrySendError::Full(msg)) => {
                let _ = self.tx.blocking_send(msg);
            }
            Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {}
        }
    }
}

impl Drop for EventRecorder {
    fn drop(&mut self) {
        let _ = self.tx.blocking_send(RecorderMessage::Shutdown);
        if let Some(worker) = self
            .worker
            .lock()
            .expect("event recorder worker mutex poisoned")
            .take()
        {
            let _ = worker.join();
        }
    }
}

impl EventRecorderInner {
    async fn flush_next_batch(&mut self) -> anyhow::Result<()> {
        let chunk_len = self.batch_size.min(self.entries.len());
        if chunk_len == 0 {
            return Ok(());
        }
        let filename = format!("{}-{:06}.json", unix_ms(SystemTime::now()), self.file_seq);
        self.file_seq = self.file_seq.saturating_add(1);
        let batch = self.entries.drain(..chunk_len).collect::<Vec<_>>();
        let payload = serde_json::to_vec(&batch)?;

        self.sink
            .write_batch(filename.as_str(), payload.as_slice())
            .await
    }
}

impl EventRecorderSink {
    async fn write_batch(&self, filename: &str, payload: &[u8]) -> anyhow::Result<()> {
        let mut attempted = 0usize;
        let mut succeeded = 0usize;
        let mut errors = Vec::new();

        if let Some(output_dir) = &self.local_output_dir {
            attempted += 1;
            let path = output_dir.join(filename);
            match fs::create_dir_all(output_dir.as_path()).and_then(|()| fs::write(path, payload)) {
                Ok(()) => {
                    succeeded += 1;
                }
                Err(err) => {
                    errors.push(format!("local sink error: {err}"));
                }
            }
        }

        if let Some(s3) = &self.s3 {
            attempted += 1;
            let key = if s3.prefix.is_empty() {
                filename.to_string()
            } else {
                format!("{}/{}", s3.prefix.trim_matches('/'), filename)
            };
            match upload_bytes_to_object_store(&s3.store, key.as_str(), payload).await {
                Ok(()) => {
                    succeeded += 1;
                }
                Err(err) => {
                    errors.push(format!("s3 sink error: {err}"));
                }
            }
        }

        if attempted == succeeded {
            return Ok(());
        }
        if succeeded > 0 {
            eprintln!(
                "performance-metrics: request-log batch only persisted to {succeeded}/{attempted} sinks: {}",
                errors.join("; ")
            );
            return Ok(());
        }
        anyhow::bail!("{}", errors.join("; "))
    }
}

fn build_event_recorder_sink(
    local_output_dir: Option<PathBuf>,
    s3_config: Option<S3SinkConfig>,
) -> anyhow::Result<EventRecorderSink> {
    let s3 = if let Some(config) = s3_config {
        let mut builder = AmazonS3Builder::new()
            .with_bucket_name(config.bucket)
            .with_region(config.region)
            .with_allow_http(config.allow_http);
        if let Some(endpoint) = config.endpoint {
            builder = builder.with_endpoint(endpoint);
        }
        if let Some(access_key_id) = config.access_key_id {
            builder = builder.with_access_key_id(access_key_id);
        }
        if let Some(secret_access_key) = config.secret_access_key {
            builder = builder.with_secret_access_key(secret_access_key);
        }
        let store = builder
            .build()
            .map_err(|e| anyhow::anyhow!("failed to build s3 store: {e}"))?;
        Some(S3Writer {
            store: Arc::new(store),
            prefix: config.prefix.trim_matches('/').to_string(),
        })
    } else {
        None
    };

    if local_output_dir.is_none() && s3.is_none() {
        anyhow::bail!("request_log_enabled requires local_output_dir and/or s3_config");
    }

    Ok(EventRecorderSink {
        local_output_dir,
        s3,
    })
}

async fn upload_bytes_to_object_store(
    store: &Arc<dyn ObjectStore>,
    key: &str,
    payload: &[u8],
) -> anyhow::Result<()> {
    let path = ObjectPath::from(key);
    store
        .put_opts(
            &path,
            PutPayload::from(Bytes::copy_from_slice(payload)),
            PutOptions::default(),
        )
        .await
        .map_err(|e| anyhow::anyhow!("failed to upload request-log chunk to object store: {e}"))?;
    Ok(())
}

fn unix_ms(ts: SystemTime) -> u64 {
    ts.duration_since(UNIX_EPOCH)
        .ok()
        .map(|d| d.as_millis() as u64)
        .unwrap_or(0)
}

fn hash_token_block(
    token_ids: &[u32],
    previous_hash: Option<&[u8; 32]>,
    key: Option<&[u8; 32]>,
) -> blake3::Hash {
    let mut hasher = if let Some(key) = key {
        Hasher::new_keyed(key)
    } else {
        Hasher::new()
    };
    hasher.update(REQUEST_TRACE_HASH_DOMAIN);
    hasher.update(&[u8::from(previous_hash.is_some())]);
    if let Some(previous_hash) = previous_hash {
        hasher.update(previous_hash);
    }
    for token_id in token_ids {
        hasher.update(&token_id.to_le_bytes());
    }
    hasher.finalize()
}

fn hash_session_id(session_id: &str, key: Option<&[u8; 32]>) -> String {
    let mut hasher = if let Some(key) = key {
        Hasher::new_keyed(key)
    } else {
        Hasher::new()
    };
    hasher.update(REQUEST_SESSION_HASH_DOMAIN);
    hasher.update(session_id.as_bytes());
    hasher.finalize().to_hex().to_string()
}

fn sampled_count(tokens: u64, sample_rate: f64, rng: &mut SmallRng) -> u64 {
    if tokens == 0 || sample_rate <= 0.0 {
        return 0;
    }
    let expected = (tokens as f64 * sample_rate).min(50.0);
    let base = expected.floor() as u64;
    let frac = expected - base as f64;
    if frac > 0.0 && rng.random::<f64>() < frac {
        base + 1
    } else {
        base
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use std::sync::mpsc::channel;
    use std::thread;
    use std::time::Duration;

    #[test]
    fn request_trace_hashes_are_consecutive_for_cumulative_output() {
        let factory = RequestTraceFactory::new(RequestTraceOptions::default()).unwrap();
        let mut trace = factory.new_trace(&[], None);
        trace.record_output_tokens(&(0..140).collect::<Vec<u32>>(), false);
        trace.record_output_tokens(&(0..260).collect::<Vec<u32>>(), false);
        trace.success();
        assert_eq!(trace.total_hashes.len(), 3);
    }

    #[test]
    fn request_trace_hash_salt_changes_hash() {
        let unsalted = RequestTraceFactory::new(RequestTraceOptions::default()).unwrap();
        let salted = RequestTraceFactory::new(RequestTraceOptions {
            request_log_hash_salt: Some("salt".to_string()),
            ..Default::default()
        })
        .unwrap();

        let mut t1 = unsalted.new_trace(&[], None);
        let mut t2 = salted.new_trace(&[], None);
        t1.record_output_tokens(&[1, 2, 3, 4, 5], true);
        t2.record_output_tokens(&[1, 2, 3, 4, 5], true);
        t1.success();
        t2.success();
        assert_ne!(t1.total_hashes(), t2.total_hashes());
    }

    #[test]
    fn request_trace_same_ids_same_salt_same_hashes() {
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            request_log_hash_salt: Some("fixed-salt".to_string()),
            ..Default::default()
        })
        .unwrap();

        let input_ids = vec![10_u32, 11, 12, 13];
        let output_ids = vec![101_u32, 102, 103, 104, 105];

        let mut a = factory.new_trace(input_ids.as_slice(), None);
        a.record_output_tokens(output_ids.as_slice(), true);
        a.success();

        let mut b = factory.new_trace(input_ids.as_slice(), None);
        b.record_output_tokens(output_ids.as_slice(), true);
        b.success();

        assert_eq!(a.total_hashes(), b.total_hashes());
    }

    #[test]
    fn request_trace_hashes_depend_on_previous_block_history() {
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            block_size: 2,
            ..Default::default()
        })
        .unwrap();

        let mut a = factory.new_trace(&[1, 2], None);
        a.record_output_tokens(&[7, 8], true);
        a.success();

        let mut b = factory.new_trace(&[3, 4], None);
        b.record_output_tokens(&[7, 8], true);
        b.success();

        assert_eq!(a.total_hashes().len(), 2);
        assert_eq!(b.total_hashes().len(), 2);
        assert_ne!(a.total_hashes()[1], b.total_hashes()[1]);
    }

    #[test]
    fn request_trace_recorder_writes_batch() {
        let tmp = std::env::temp_dir().join(format!("pm-core-test-{}", unix_ms(SystemTime::now())));
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            request_log_enabled: true,
            local_output_dir: Some(tmp.clone()),
            request_log_batch_size: 2,
            ..Default::default()
        })
        .unwrap();

        let mut a = factory.new_trace(&[], None);
        a.record_output_tokens(&[1], true);
        a.success();

        let mut b = factory.new_trace(&[], None);
        b.record_output_tokens(&[2], true);
        b.cancel();

        thread::sleep(std::time::Duration::from_millis(20));
        let file_count = fs::read_dir(tmp.as_path()).unwrap().count();
        assert!(file_count >= 1);
    }

    #[test]
    fn request_trace_diff_and_cumulative_modes_produce_same_hashes() {
        let factory = RequestTraceFactory::new(RequestTraceOptions::default()).unwrap();

        let mut cumulative = factory.new_trace(&[], None);
        cumulative.record_output_tokens(&[1, 2, 3], false);
        cumulative.record_output_tokens(&[1, 2, 3, 4, 5], false);
        cumulative.record_output_tokens(&[1, 2, 3, 4, 5, 6], false);
        cumulative.success();

        let mut diff = factory.new_trace(&[], None);
        diff.record_output_tokens(&[1, 2, 3], true);
        diff.record_output_tokens(&[4, 5], true);
        diff.record_output_tokens(&[6], true);
        diff.success();

        assert_eq!(cumulative.total_hashes(), diff.total_hashes());
    }

    #[test]
    fn request_trace_drop_flushes_partial_batch() {
        let tmp =
            std::env::temp_dir().join(format!("pm-core-drop-test-{}", unix_ms(SystemTime::now())));
        {
            let factory = RequestTraceFactory::new(RequestTraceOptions {
                request_log_enabled: true,
                local_output_dir: Some(tmp.clone()),
                request_log_batch_size: 10,
                ..Default::default()
            })
            .unwrap();
            let mut trace = factory.new_trace(&[1, 2, 3], None);
            trace.record_output_tokens(&[4, 5], true);
            // no success/cancel, should flush on drop
        }

        let file_count = fs::read_dir(tmp.as_path()).unwrap().count();
        assert!(file_count >= 1);
    }

    #[test]
    fn request_trace_sink_failure_does_not_block_shutdown() {
        let tmp = std::env::temp_dir().join(format!(
            "pm-core-invalid-sink-{}",
            unix_ms(SystemTime::now())
        ));
        fs::write(tmp.as_path(), b"not-a-directory").unwrap();

        let (done_tx, done_rx) = channel();
        thread::spawn(move || {
            {
                let factory = RequestTraceFactory::new(RequestTraceOptions {
                    request_log_enabled: true,
                    local_output_dir: Some(tmp),
                    request_log_batch_size: 1,
                    ..Default::default()
                })
                .unwrap();
                let mut trace = factory.new_trace(&[1, 2, 3], None);
                trace.record_output_tokens(&[4], true);
                trace.success();
            }
            let _ = done_tx.send(());
        });

        assert!(done_rx.recv_timeout(Duration::from_secs(1)).is_ok());
    }

    #[test]
    fn request_trace_s3_config_builds_without_network_io() {
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            request_log_enabled: true,
            local_output_dir: None,
            s3_config: Some(S3SinkConfig {
                bucket: "example-bucket".to_string(),
                region: "us-east-1".to_string(),
                endpoint: Some("http://localhost:4566".to_string()),
                access_key_id: Some("x".to_string()),
                secret_access_key: Some("x".to_string()),
                prefix: "records".to_string(),
                allow_http: true,
            }),
            ..Default::default()
        });
        assert!(factory.is_ok());
    }

    #[test]
    fn request_metrics_factory_tracks_ttft_and_kv_ratio() {
        let factory = RequestMetricsFactory::new(RequestMetricsOptions::default()).unwrap();
        let mut req = factory.new_request(&[1, 2, 3, 4], RequestFeatures::NONE, None);
        req.record_tokens(&[100, 101], Some(2), true);
        req.success();

        let snapshot = factory.snapshot();
        let ttft = find_metric_snapshot(&snapshot, "request_metrics_ttft_ms", Some("all"));
        let kv = find_metric_snapshot(&snapshot, "request_metrics_kv_cache_hit_rate", None);
        let itl = find_metric_snapshot(&snapshot, "request_metrics_itl_ms", Some("all"));
        let successful =
            find_metric_snapshot(&snapshot, "request_metrics_successful_request", None);
        assert!(ttft.average.unwrap_or(0.0) >= 0.0);
        assert!((kv.ratio.unwrap_or(0.0) - 0.5).abs() < 1e-9);
        assert!(itl.average.unwrap_or(0.0) >= 0.0);
        assert!(successful.rate_per_second.unwrap_or(0.0) >= 0.0);
    }

    #[test]
    fn request_metrics_feature_presence_series_overlap() {
        let factory = RequestMetricsFactory::new(RequestMetricsOptions::default()).unwrap();
        let mut req = factory.new_request(
            &[1, 2, 3, 4],
            RequestFeatures::TOOLS | RequestFeatures::IMAGE,
            None,
        );
        req.record_tokens(&[100, 101, 102], Some(2), true);
        req.record_tokens(&[103, 104], None, true);
        req.success();

        let snapshot = factory.snapshot();
        let ttft_tools = find_metric_snapshot(&snapshot, "request_metrics_ttft_ms", Some("tools"));
        let ttft_image = find_metric_snapshot(&snapshot, "request_metrics_ttft_ms", Some("image"));
        let ttft_xgrammar =
            find_metric_snapshot(&snapshot, "request_metrics_ttft_ms", Some("xgrammar"));
        let itl_tools = find_metric_snapshot(&snapshot, "request_metrics_itl_ms", Some("tools"));
        let itl_image = find_metric_snapshot(&snapshot, "request_metrics_itl_ms", Some("image"));
        let itl_xgrammar =
            find_metric_snapshot(&snapshot, "request_metrics_itl_ms", Some("xgrammar"));

        assert!(ttft_tools.average.unwrap_or(0.0) >= 0.0);
        assert!(ttft_image.average.unwrap_or(0.0) >= 0.0);
        assert_eq!(ttft_xgrammar.average.unwrap_or(0.0), 0.0);
        assert!(itl_tools.average.unwrap_or(0.0) >= 0.0);
        assert!(itl_image.average.unwrap_or(0.0) >= 0.0);
        assert_eq!(itl_xgrammar.average.unwrap_or(0.0), 0.0);
    }

    #[test]
    fn request_metrics_factory_tracks_lifetime_totals() {
        let factory = RequestMetricsFactory::new(RequestMetricsOptions::default()).unwrap();
        let mut req = factory.new_request(&[1, 2, 3, 4], RequestFeatures::NONE, None);
        req.record_tokens(&[10, 11], Some(2), true);
        req.record_tokens(&[12], None, true);
        req.success();

        let snapshot = factory.snapshot();
        let totals = &snapshot.totals;
        assert_eq!(totals.total_requests, 1);
        assert_eq!(totals.total_successful_requests, 1);
        assert_eq!(totals.total_pre_first_token_cancellations, 0);
        assert_eq!(totals.total_mid_stream_cancellations, 0);
        assert_eq!(totals.total_input_tokens, 4);
        assert_eq!(totals.total_output_tokens, 3);
        assert_eq!(totals.total_kv_cache_hits_tokens, 2);
        assert_eq!(totals.total_kv_cache_lookup_tokens, 4);
        assert_eq!(totals.kv_cache_hit_rate_lifetime.unwrap_or(0.0), 0.5);
    }

    #[test]
    fn dropped_request_counts_as_cancellation() {
        let factory = RequestMetricsFactory::new(RequestMetricsOptions::default()).unwrap();
        {
            let _req = factory.new_request(&[1, 2, 3, 4], RequestFeatures::NONE, None);
        }

        let snapshot = factory.snapshot();
        let totals = &snapshot.totals;
        assert_eq!(totals.total_requests, 1);
        assert_eq!(totals.total_pre_first_token_cancellations, 1);
        assert_eq!(totals.total_mid_stream_cancellations, 0);

        let cancellation = find_metric_snapshot(
            &snapshot,
            "request_metrics_pre_first_token_cancellation",
            None,
        );
        assert!(cancellation.rate_per_second.unwrap_or(0.0) >= 0.0);
    }

    #[test]
    fn request_metrics_default_has_dense_kv_cache_quantiles() {
        let options = RequestMetricsOptions::default();
        assert_eq!(
            options.kv_cache_hit_rate_quantiles,
            vec![0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99, 0.999]
        );
    }

    #[test]
    fn request_log_entry_omits_session_id_when_absent() {
        let entry = RequestLogEntry {
            timestamp_start_unix_ms: 1,
            timestamp_end_unix_ms: 2,
            input_tokens: 3,
            output_tokens: 4,
            total_hashes: vec!["abc".to_string()],
            session_id: None,
            request_canceled: false,
        };

        let value = serde_json::to_value(entry).unwrap();
        assert!(value.get("session_id").is_none());
    }

    #[test]
    fn request_trace_hashes_session_id_when_present() {
        let factory = RequestTraceFactory::new(RequestTraceOptions::default()).unwrap();
        let trace = factory.new_trace(&[1, 2, 3], Some("session-123"));
        let entry = trace.build_log_entry(false);
        let expected = hash_session_id("session-123", None);

        assert_eq!(entry.session_id.as_deref(), Some(expected.as_str()));
    }

    #[test]
    fn request_trace_hashes_session_id_with_derived_secret() {
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            request_log_hash_salt: Some("fixed-salt".to_string()),
            ..Default::default()
        })
        .unwrap();
        let trace = factory.new_trace(&[1, 2, 3], Some("session-123"));
        let entry = trace.build_log_entry(false);
        let derived_key = blake3::hash(blake3::hash(b"fixed-salt").as_bytes());
        let expected = hash_session_id("session-123", Some(derived_key.as_bytes()));

        assert_eq!(entry.session_id.as_deref(), Some(expected.as_str()));
    }

    fn find_metric_snapshot<'a>(
        snapshot: &'a RequestMetricsSnapshot,
        name: &str,
        feature: Option<&str>,
    ) -> &'a crate::MetricSnapshot {
        &snapshot
            .series
            .iter()
            .find(|series| {
                series.name == name
                    && match feature {
                        Some(feature) => {
                            series.labels.get("feature").map(String::as_str) == Some(feature)
                        }
                        None => !series.labels.contains_key("feature"),
                    }
            })
            .unwrap()
            .snapshot
    }
}
