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";

/// Blake3 hash output size in bytes (256 bits = 32 bytes)
const HASH_SIZE: usize = 32;
pub const DEFAULT_REQUEST_LOG_VERSION: &str = env!("CARGO_PKG_VERSION");
pub const DEFAULT_BLOCK_SIZE: usize = 64;
pub const DEFAULT_REQUEST_LOG_BATCH_SIZE: usize = 2048;
pub const DEFAULT_REQUEST_LOG_QUEUE_CAPACITY: usize = 4096;
pub const DEFAULT_LOCAL_OUTPUT_DIR: &str = "/tmp/records";

/// Fixed-size array for a 32-byte Blake3 hash.
pub type HashBytes = [u8; 32];

#[derive(Deserialize)]
#[serde(untagged)]
enum JsonHashBytes {
    Hex(String),
    Bytes(HashBytes),
}

fn encode_hash_hex(hash: &HashBytes) -> String {
    let mut out = String::with_capacity(HASH_SIZE * 2);
    for byte in hash {
        use std::fmt::Write as _;
        let _ = write!(out, "{byte:02x}");
    }
    out
}

fn decode_hash_hex(value: &str) -> Result<HashBytes, String> {
    if value.len() != HASH_SIZE * 2 {
        return Err(format!(
            "expected {} hex chars, got {}",
            HASH_SIZE * 2,
            value.len()
        ));
    }

    let mut out = [0u8; HASH_SIZE];
    for (i, chunk) in value.as_bytes().chunks_exact(2).enumerate() {
        let piece = std::str::from_utf8(chunk).map_err(|_| "hash contains non-utf8 bytes")?;
        out[i] = u8::from_str_radix(piece, 16)
            .map_err(|_| format!("invalid hex at byte offset {}", i * 2))?;
    }
    Ok(out)
}

fn serialize_hash_bytes_as_hex<S>(hashes: &[HashBytes], serializer: S) -> Result<S::Ok, S::Error>
where
    S: serde::Serializer,
{
    let hex_hashes = hashes.iter().map(encode_hash_hex).collect::<Vec<_>>();
    hex_hashes.serialize(serializer)
}

fn deserialize_hash_bytes_from_json<'de, D>(deserializer: D) -> Result<Vec<HashBytes>, D::Error>
where
    D: serde::Deserializer<'de>,
{
    let hashes = Vec::<JsonHashBytes>::deserialize(deserializer)?;
    hashes
        .into_iter()
        .map(|hash| match hash {
            JsonHashBytes::Hex(value) => decode_hash_hex(&value).map_err(serde::de::Error::custom),
            JsonHashBytes::Bytes(value) => Ok(value),
        })
        .collect()
}

#[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>,
    pub model_name: String,
    pub org_id: String,
    pub file_format: FileFormat,
}

#[derive(Debug, Clone, Copy, PartialEq, Eq, Default, Serialize, Deserialize)]
#[serde(rename_all = "lowercase")]
pub enum FileFormat {
    #[default]
    Json,
    Parquet,
}

#[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,
}

impl Default for RequestTraceOptions {
    fn default() -> Self {
        Self {
            block_size: DEFAULT_BLOCK_SIZE,
            request_log_enabled: false,
            request_log_hash_salt: None,
            request_log_batch_size: DEFAULT_REQUEST_LOG_BATCH_SIZE,
            request_log_queue_capacity: DEFAULT_REQUEST_LOG_QUEUE_CAPACITY,
            local_output_dir: Some(PathBuf::from(DEFAULT_LOCAL_OUTPUT_DIR)),
            s3_config: None,
            model_name: "unset_model_name".to_string(),
            org_id: String::new(),
            file_format: FileFormat::Json,
        }
    }
}

#[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 pre_20_token_cancellation_quantiles: Vec<f64>,
    pub pre_20_token_cancellation_sample_period_seconds: Option<f64>,
    pub pre_20_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,
            pre_20_token_cancellation_quantiles: vec![],
            pre_20_token_cancellation_sample_period_seconds: Some(1.0),
            pre_20_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 duration_e2e_ms: u64,
    pub duration_ttft_ms: u64,
    pub input_tokens: u64,
    pub output_tokens: u64,
    #[serde(
        serialize_with = "serialize_hash_bytes_as_hex",
        deserialize_with = "deserialize_hash_bytes_from_json"
    )]
    pub total_hashes: Vec<HashBytes>,
    pub provided_session_id: Option<String>,
    pub request_canceled: bool,
    pub cached_tokens_reference: Option<u64>,
    pub model_name: String,
    pub org_id: String,
    pub block_size: usize,
    pub features: Vec<String>,
    pub speculation_ratio: Option<f64>,
    pub __version__: String,
}

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

#[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
    }

    pub fn to_vec(self) -> Vec<String> {
        RequestFeatureKind::ALL
            .iter()
            .filter(|&&f| self.includes(f))
            .map(|f| f.metric_suffix().to_string())
            .collect()
    }
}

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,
    request_quantiles: Vec<f64>,
    request_sample_period_seconds: Option<f64>,
    request_window_seconds: Option<f64>,
    ttft_quantiles: Vec<f64>,
    ttft_window_seconds: Option<f64>,
    itl_quantiles: Vec<f64>,
    itl_window_seconds: Option<f64>,
}

#[derive(Clone)]
struct RequestFeatureMetricHandles {
    request_all: RateMetricHandle,
    request_none: RateMetricHandle,
    request_by_feature: Vec<(RequestFeatureKind, RateMetricHandle)>,
    ttft_none: DistributionMetricHandle,
    ttft_by_feature: Vec<(RequestFeatureKind, DistributionMetricHandle)>,
    itl_none: DistributionMetricHandle,
    itl_by_feature: Vec<(RequestFeatureKind, DistributionMetricHandle)>,
}

#[derive(Clone)]
pub struct RequestMetricsFactory {
    trace_factory: RequestTraceFactory,
    metrics: PerformanceMetricsRegistry,
    request: RateMetricHandle,
    input_tokens: RateMetricHandle,
    input_length_tokens: DistributionMetricHandle,
    net_new_input_tokens: RateMetricHandle,
    output_tokens: RateMetricHandle,
    output_length_tokens: DistributionMetricHandle,
    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,
    pre_20_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,
    first_token_at: Option<SystemTime>,
    input_tokens: usize,
    output_tokens_len: usize,
    total_hashes: Vec<HashBytes>,
    provided_session_id: Option<String>,
    previous_hash: Option<[u8; 32]>,
    pending: Vec<u32>,
    terminal: bool,
    cached_tokens: Option<u64>,
    model_name: String,
    org_id: String,
    block_size: usize,
    features: Vec<String>,
    speculation_ratio: Option<f64>,
}

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,
    output_length_tokens: DistributionMetricHandle,
    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,
    pre_20_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_pre_20_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_pre_20_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,
}

const PRE_20_TOKEN_CANCELLATION_EXCLUSIVE_MAX_OUTPUT_TOKENS: usize = 20;

/// Hash a string using unkeyed blake3 to avoid exposing semi sensitive information.
fn hash_unkeyed_string(value: &str) -> String {
    blake3::hash(value.as_bytes()).to_hex().to_string()
}

impl RequestTraceFactory {
    pub fn new(options: RequestTraceOptions) -> anyhow::Result<Self> {
        // Hash the model_name using unkeyed blake3 to avoid exposing sensitive model names
        let model_name_hash = hash_unkeyed_string(&options.model_name);

        // Hash the org_id using unkeyed blake3 to avoid exposing sensitive organization identifiers
        let org_id_hash = hash_unkeyed_string(&options.org_id);

        println!(
            "Initializing RequestTraceFactory, request_log_enabled={}, s3={:?}",
            options.request_log_enabled,
            options.s3_config.is_some(),
        );

        let recorder = if options.request_log_enabled {
            let sink = build_event_recorder_sink(
                options.local_output_dir.clone(),
                options.s3_config.clone(),
                options.file_format,
            )?;
            Some(Arc::new(EventRecorder::new(
                sink,
                options.request_log_batch_size,
                options.request_log_queue_capacity,
                model_name_hash.clone(),
            )))
        } 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),
            model_name: model_name_hash,
            org_id: org_id_hash,
        })
    }

    /// Create a trace with token array (includes hashing if block_size permits).
    pub fn new_trace(
        &self,
        input_token_ids: &[u32],
        provided_session_id: Option<&str>,
        features: Vec<String>,
    ) -> RequestTrace {
        let mut trace =
            self.create_trace_base(input_token_ids.len(), provided_session_id, features);
        trace.append_tokens(input_token_ids);
        trace
    }

    /// Create a trace with just the token count (no hashing, no token array needed).
    pub fn new_trace_with_count(
        &self,
        input_token_count: usize,
        provided_session_id: Option<&str>,
        features: Vec<String>,
    ) -> RequestTrace {
        // Skip append_tokens - no hashing when using count-only
        self.create_trace_base(input_token_count, provided_session_id, features)
    }

    /// Common initialization for both trace creation methods.
    fn create_trace_base(
        &self,
        input_token_count: usize,
        provided_session_id: Option<&str>,
        features: Vec<String>,
    ) -> RequestTrace {
        RequestTrace {
            factory: self.clone(),
            started_at_wallclock: SystemTime::now(),
            first_token_at: None,
            input_tokens: input_token_count,
            output_tokens_len: 0,
            total_hashes: Vec::new(),
            provided_session_id: provided_session_id
                .map(|id| hash_session_id(id, self.session_hash_key.as_ref())),
            previous_hash: None,
            pending: Vec::with_capacity(self.block_size.min(input_token_count)),
            terminal: false,
            cached_tokens: None,
            model_name: self.model_name.clone(),
            org_id: self.org_id.clone(),
            block_size: self.block_size,
            features,
            speculation_ratio: None,
        }
    }
}

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 input_length_tokens = metrics.new_distribution_metric(
            format!("{prefix}_input_length_tokens"),
            options.input_tokens_quantiles.clone(),
            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 output_length_tokens = metrics.new_distribution_metric(
            format!("{prefix}_output_length_tokens"),
            options.input_tokens_quantiles.clone(),
            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 pre_20_token_cancellation = metrics.new_rate_metric(
            format!("{prefix}_pre_20_token_cancellation"),
            options.pre_20_token_cancellation_quantiles.clone(),
            options.pre_20_token_cancellation_sample_period_seconds,
            options.pre_20_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,
        )?;

        println!(
            "Registered RequestMetricsFactory with prefix '{}'",
            options.metric_prefix
        );

        Ok(Self {
            trace_factory,
            metrics,
            request,
            input_tokens,
            input_length_tokens,
            net_new_input_tokens,
            output_tokens,
            output_length_tokens,
            ttft_ms,
            ttft_ms_per_input_token,
            ttft_ms_per_net_new_input_token,
            itl_ms,
            pre_first_token_cancellation,
            pre_20_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,
                request_quantiles: options.request_quantiles,
                request_sample_period_seconds: options.request_sample_period_seconds,
                request_window_seconds: options.request_window_seconds,
                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()),
        })
    }

    /// Create a new request with the given token array.
    pub fn new_request(
        &self,
        input_token_ids: &[u32],
        features: RequestFeatures,
        provided_session_id: Option<&str>,
    ) -> RequestMetricsRequest {
        self.new_request_internal(
            input_token_ids.len(),
            features,
            provided_session_id,
            |_count, sid, feats| self.trace_factory.new_trace(input_token_ids, sid, feats),
        )
    }

    /// Create a new request with just the token count (no hashing, no token array needed).
    pub fn new_request_with_count(
        &self,
        input_token_count: usize,
        features: RequestFeatures,
        provided_session_id: Option<&str>,
    ) -> RequestMetricsRequest {
        self.new_request_internal(
            input_token_count,
            features,
            provided_session_id,
            |count, sid, feats| self.trace_factory.new_trace_with_count(count, sid, feats),
        )
    }

    /// Common implementation for both request creation methods.
    fn new_request_internal<F>(
        &self,
        input_token_count: usize,
        features: RequestFeatures,
        provided_session_id: Option<&str>,
        create_trace: F,
    ) -> RequestMetricsRequest
    where
        F: FnOnce(usize, Option<&str>, Vec<String>) -> RequestTrace,
    {
        let _ = self.request.record_count(1);
        let _ = self
            .input_length_tokens
            .record_value(input_token_count as f64);
        self.totals.total_requests.fetch_add(1, Ordering::Relaxed);
        self.totals
            .total_input_tokens
            .fetch_add(input_token_count as u64, Ordering::Relaxed);
        let (ttft_feature_metrics, itl_feature_metrics) = self.feature_metric_handles(features);
        let features_vec = features.to_vec();

        RequestMetricsRequest {
            trace: create_trace(input_token_count, provided_session_id, features_vec),
            started_at: Instant::now(),
            input_tokens: input_token_count 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(),
            output_length_tokens: self.output_length_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(),
            pre_20_token_cancellation: self.pre_20_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>) {
        let handles = self.feature_metrics.get_or_init(|| {
            RequestFeatureMetricHandles::new(&self.metrics, &self.feature_metric_config)
                .expect("failed to register request feature metrics")
        });
        handles.record_request(features);
        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,
        speculation_ratio: Option<f64>,
    ) {
        if let Some(cached) = cached_tokens {
            self.cached_tokens_hint.get_or_insert(cached);
            self.trace.set_cached_tokens(cached);
        }
        if let Some(ratio) = speculation_ratio {
            self.trace.set_speculation_ratio(ratio);
        }

        let appended = self
            .trace
            .record_output_tokens(token_ids, is_diff, speculation_ratio);
        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
            .output_length_tokens
            .record_value(self.trace.output_tokens() as f64);
        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;
        }
        let _ = self
            .output_length_tokens
            .record_value(self.trace.output_tokens() as f64);
        let output_tokens = self.trace.output_tokens();
        if output_tokens == 0 {
            let _ = self.pre_first_token_cancellation.record_count(1);
            self.totals
                .total_pre_first_token_cancellations
                .fetch_add(1, Ordering::Relaxed);
        } else if output_tokens < PRE_20_TOKEN_CANCELLATION_EXCLUSIVE_MAX_OUTPUT_TOKENS {
            let _ = self.pre_20_token_cancellation.record_count(1);
            self.totals
                .total_pre_20_token_cancellations
                .fetch_add(1, Ordering::Relaxed);
        } else {
            let _ = self.mid_stream_cancellation.record_count(1);
            self.totals
                .total_mid_stream_cancellations
                .fetch_add(1, Ordering::Relaxed);
        }
        self.trace.cancel();
    }

    pub fn total_hashes(&self) -> &[HashBytes] {
        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 request_all = metrics.new_rate_metric_with_labels(
            format!("{}_request", config.metric_prefix),
            feature_metric_labels("all"),
            config.request_quantiles.clone(),
            config.request_sample_period_seconds,
            config.request_window_seconds,
        )?;
        let request_none = metrics.new_rate_metric_with_labels(
            format!("{}_request", config.metric_prefix),
            feature_metric_labels("none"),
            config.request_quantiles.clone(),
            config.request_sample_period_seconds,
            config.request_window_seconds,
        )?;
        let mut request_by_feature = Vec::with_capacity(RequestFeatureKind::ALL.len());
        for feature in RequestFeatureKind::ALL {
            request_by_feature.push((
                feature,
                metrics.new_rate_metric_with_labels(
                    format!("{}_request", config.metric_prefix),
                    feature_metric_labels(feature.metric_suffix()),
                    config.request_quantiles.clone(),
                    config.request_sample_period_seconds,
                    config.request_window_seconds,
                )?,
            ));
        }

        let ttft_none = metrics.new_distribution_metric_with_labels(
            format!("{}_ttft_ms", config.metric_prefix),
            feature_metric_labels("none"),
            config.ttft_quantiles.clone(),
            config.ttft_window_seconds,
        )?;
        let mut ttft_by_feature = Vec::with_capacity(RequestFeatureKind::ALL.len());
        let itl_none = metrics.new_distribution_metric_with_labels(
            format!("{}_itl_ms", config.metric_prefix),
            feature_metric_labels("none"),
            config.itl_quantiles.clone(),
            config.itl_window_seconds,
        )?;
        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 {
            request_all,
            request_none,
            request_by_feature,
            ttft_none,
            ttft_by_feature,
            itl_none,
            itl_by_feature,
        })
    }

    fn record_request(&self, features: RequestFeatures) {
        let _ = self.request_all.record_count(1);
        if features.is_empty() {
            let _ = self.request_none.record_count(1);
            return;
        }
        for (_, handle) in self
            .request_by_feature
            .iter()
            .filter(|(feature, _)| features.includes(*feature))
        {
            let _ = handle.record_count(1);
        }
    }

    fn for_features(
        &self,
        features: RequestFeatures,
    ) -> (Vec<DistributionMetricHandle>, Vec<DistributionMetricHandle>) {
        let ttft = if features.is_empty() {
            vec![self.ttft_none.clone()]
        } else {
            self.ttft_by_feature
                .iter()
                .filter(|(feature, _)| features.includes(*feature))
                .map(|(_, handle)| handle.clone())
                .collect()
        };
        let itl = if features.is_empty() {
            vec![self.itl_none.clone()]
        } else {
            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,
        speculation_ratio: Option<f64>,
    ) -> usize {
        if self.terminal {
            return 0;
        }
        // Store the last speculation_ratio if provided
        if let Some(ratio) = speculation_ratio {
            self.speculation_ratio = Some(ratio);
        }
        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..]
        };
        // Track first token time if this is the first output
        if self.output_tokens_len == 0 && !delta.is_empty() {
            self.first_token_at = Some(SystemTime::now());
        }
        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) -> &[HashBytes] {
        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
    }

    pub fn set_cached_tokens(&mut self, cached_tokens: u64) {
        self.cached_tokens = Some(cached_tokens);
    }

    pub fn set_speculation_ratio(&mut self, speculation_ratio: f64) {
        self.speculation_ratio = Some(speculation_ratio);
    }

    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.as_bytes());
        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 {
        let now = SystemTime::now();
        let start_ms = unix_ms(self.started_at_wallclock);
        let now_ms = unix_ms(now);
        let duration_e2e_ms = now_ms.saturating_sub(start_ms);
        let duration_ttft_ms = self
            .first_token_at
            .map(|t| unix_ms(t).saturating_sub(start_ms))
            .unwrap_or(0);
        RequestLogEntry {
            timestamp_start_unix_ms: start_ms,
            duration_e2e_ms,
            duration_ttft_ms,
            input_tokens: self.input_tokens as u64,
            output_tokens: self.output_tokens_len as u64,
            total_hashes: self.total_hashes.clone(),
            provided_session_id: self.provided_session_id.clone(),
            request_canceled: canceled,
            cached_tokens_reference: self.cached_tokens,
            model_name: self.model_name.clone(),
            org_id: self.org_id.clone(),
            block_size: self.block_size,
            features: self.features.clone(),
            speculation_ratio: self.speculation_ratio,
            __version__: DEFAULT_REQUEST_LOG_VERSION.to_string(),
        }
    }
}

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_pre_20_token_cancellations = self
            .total_pre_20_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_pre_20_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>,
    file_format: FileFormat,
}

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

/// Time-based flush configuration - flush every 15 minutes if there's at least 1 entry
const TIME_BASED_FLUSH_INTERVAL: std::time::Duration = std::time::Duration::from_secs(15 * 60); // 15 minutes

#[derive(Debug)]
struct EventRecorderInner {
    entries: Vec<RequestLogEntry>,
    sink: EventRecorderSink,
    batch_size: usize,
    model_name_hash: String,
    last_flush_time: std::time::Instant,
}

enum RecorderMessage {
    Record(Box<RequestLogEntry>),
    Shutdown,
}

impl EventRecorder {
    fn new(
        sink: EventRecorderSink,
        batch_size: usize,
        queue_capacity: usize,
        model_name_hash: String,
    ) -> 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),
                    model_name_hash: model_name_hash.chars().take(8).collect(),
                    last_flush_time: std::time::Instant::now(),
                };

                loop {
                    // Calculate next flush deadline based on last flush time
                    let next_flush_deadline = tokio::time::Instant::from_std(
                        inner.last_flush_time + TIME_BASED_FLUSH_INTERVAL,
                    );

                    // Use tokio::select! to wait for either a message or timeout
                    let msg = tokio::select! {
                        msg = rx.recv() => msg,
                        _ = tokio::time::sleep_until(next_flush_deadline) => {
                            // Always advance the timer cursor, even when idle, so we don't spin on a past deadline.
                            inner.last_flush_time = std::time::Instant::now();

                            // Timeout: flush if we have entries.
                            if !inner.entries.is_empty() {
                                println!("performance-metrics: time-based flush triggered, {} events since {} minutes ago.", inner.entries.len(), TIME_BASED_FLUSH_INTERVAL.as_secs() / 60);
                                if let Err(err) = inner.flush_next_batch().await {
                                    eprintln!(
                                        "performance-metrics: dropping request-log batch after time-based flush failure: {err}"
                                    );
                                }
                            }
                            continue;
                        }
                    };

                    match msg {
                        Some(RecorderMessage::Record(entry)) => {
                            inner.entries.push(*entry);

                            // Size-based flush
                            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}"
                                    );
                                }
                        }
                        Some(RecorderMessage::Shutdown) => {
                            break;
                        }
                        None => {
                            // Channel closed
                            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) {
        // Non-blocking: try to send, drop entry if channel is full or closed
        // This ensures telemetry never blocks the main request path
        match self.tx.try_send(RecorderMessage::Record(Box::new(entry))) {
            Ok(()) => {}
            Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
                // Channel full, drop entry to avoid blocking main thread
                eprintln!("performance-metrics: dropping request-log entry, channel full");
            }
            Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
                // Channel closed (shutting down), drop entry
            }
        }
    }
}

impl Drop for EventRecorder {
    fn drop(&mut self) {
        println!("Shutting down EventRecorder...");
        // Try to send shutdown message with a timeout to avoid hanging if channel is full
        // Use try_send in a loop with sleep to implement a timeout
        let shutdown_timeout = std::time::Duration::from_millis(100);
        let start = std::time::Instant::now();
        while start.elapsed() < shutdown_timeout {
            match self.tx.try_send(RecorderMessage::Shutdown) {
                Ok(()) => break,
                Err(tokio::sync::mpsc::error::TrySendError::Full(_)) => {
                    // Channel full, wait a bit and retry
                    std::thread::sleep(std::time::Duration::from_millis(5));
                }
                Err(tokio::sync::mpsc::error::TrySendError::Closed(_)) => {
                    // Channel already closed, that's fine
                    break;
                }
            }
        }

        if let Some(worker) = self
            .worker
            .lock()
            .expect("event recorder worker mutex poisoned")
            .take()
        {
            // Hail mary: wait up to 2s for worker to finish flushing remaining entries
            let timeout = std::time::Duration::from_secs(2);
            let start = std::time::Instant::now();
            while start.elapsed() < timeout {
                if worker.is_finished() {
                    let _ = worker.join();
                    return;
                }
                std::thread::sleep(std::time::Duration::from_millis(10));
            }
            // Timeout reached, abandon the worker
            eprintln!(
                "performance-metrics: recorder drop timed out after 2s, abandoning remaining entries"
            );
        }
    }
}

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(());
        }

        // Update last flush time before flushing
        self.last_flush_time = std::time::Instant::now();

        let batch = self.entries.drain(..chunk_len).collect::<Vec<_>>();

        let (payload, extension) = match self.sink.file_format {
            FileFormat::Json => {
                let payload = serde_json::to_vec(&batch)?;
                (payload, "json")
            }
            FileFormat::Parquet => {
                let payload = serialize_to_parquet(&batch)?;
                (payload, "parquet")
            }
        };

        // Generate random 4-character hex suffix for distributed uniqueness
        let random_suffix: String = (0..4)
            .map(|_| format!("{:02x}", rand::random::<u8>()))
            .collect();
        let filename = format!(
            "{}-{}-{}.{}",
            self.model_name_hash,
            unix_ms(SystemTime::now()),
            random_suffix,
            extension
        );

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

fn serialize_to_parquet(entries: &[RequestLogEntry]) -> anyhow::Result<Vec<u8>> {
    use arrow::array::{
        FixedSizeBinaryArray, Float64Array, ListArray, StringArray, UInt8Array, UInt64Array,
    };
    use arrow::buffer::OffsetBuffer;
    use arrow::datatypes::{DataType, Field, Schema};
    use arrow::record_batch::RecordBatch;
    use parquet::arrow::ArrowWriter;
    use parquet::basic::Compression;
    use parquet::file::properties::WriterProperties;
    use std::sync::Arc;

    let schema = Schema::new(vec![
        Field::new("timestamp_start_unix_ms", DataType::UInt64, false),
        Field::new("duration_e2e_ms", DataType::UInt64, false),
        Field::new("duration_ttft_ms", DataType::UInt64, false),
        Field::new("input_tokens", DataType::UInt64, false),
        Field::new("output_tokens", DataType::UInt64, false),
        Field::new(
            "total_hashes",
            DataType::List(Arc::new(Field::new(
                "item",
                DataType::FixedSizeBinary(32),
                false,
            ))),
            true,
        ),
        Field::new("provided_session_id", DataType::Utf8, true),
        Field::new("request_canceled", DataType::UInt8, false),
        Field::new("cached_tokens_reference", DataType::UInt64, true),
        Field::new("model_name", DataType::Utf8, false),
        Field::new("org_id", DataType::Utf8, false),
        Field::new("block_size", DataType::UInt64, false),
        Field::new("features", DataType::Utf8, true),
        Field::new("speculation_ratio", DataType::Float64, true),
        Field::new("__version__", DataType::Utf8, false),
    ]);

    let timestamp_start: Vec<u64> = entries.iter().map(|e| e.timestamp_start_unix_ms).collect();
    let duration_e2e: Vec<u64> = entries.iter().map(|e| e.duration_e2e_ms).collect();
    let duration_ttft: Vec<u64> = entries.iter().map(|e| e.duration_ttft_ms).collect();
    let input_tokens: Vec<u64> = entries.iter().map(|e| e.input_tokens).collect();
    let output_tokens: Vec<u64> = entries.iter().map(|e| e.output_tokens).collect();

    // Build ListArray for total_hashes (each entry is a list of FixedSizeBinary(HASH_SIZE))
    // Pre-allocate buffer: sum all hash lengths (each hash is HASH_SIZE bytes)
    let total_hashes_count: usize = entries.iter().map(|e| e.total_hashes.len()).sum();
    let mut total_hashes_values: Vec<u8> = Vec::with_capacity(total_hashes_count * HASH_SIZE);
    for entry in entries {
        for hash in &entry.total_hashes {
            // Hash is already raw bytes (32 bytes)
            total_hashes_values.extend_from_slice(hash);
        }
    }

    // Build offsets: [0, len1, len1+len2, len1+len2+len3, ...]
    // Pre-allocate: entries.len() + 1 offsets
    let mut total_hashes_offsets: Vec<i32> = Vec::with_capacity(entries.len() + 1);
    total_hashes_offsets.push(0);
    for entry in entries {
        let next_offset = total_hashes_offsets.last().unwrap() + entry.total_hashes.len() as i32;
        total_hashes_offsets.push(next_offset);
    }

    let total_hashes_offsets = OffsetBuffer::new(total_hashes_offsets.into());

    // Build FixedSizeBinaryArray from the raw bytes
    // Each hash is HASH_SIZE bytes, use chunks for cleaner slicing
    let hash_refs: Vec<&[u8]> = total_hashes_values.chunks(HASH_SIZE).collect();
    let fixed_size_binary = FixedSizeBinaryArray::from(hash_refs);

    let total_hashes = ListArray::new(
        Arc::new(Field::new(
            "item",
            DataType::FixedSizeBinary(HASH_SIZE as i32),
            false,
        )),
        total_hashes_offsets,
        Arc::new(fixed_size_binary),
        None,
    );
    let provided_session_id: Vec<Option<String>> = entries
        .iter()
        .map(|e| e.provided_session_id.clone())
        .collect();
    let request_canceled: Vec<u8> = entries
        .iter()
        .map(|e| if e.request_canceled { 1 } else { 0 })
        .collect();
    let cached_tokens_ref: Vec<Option<u64>> =
        entries.iter().map(|e| e.cached_tokens_reference).collect();
    let model_name: Vec<String> = entries.iter().map(|e| e.model_name.clone()).collect();
    let org_id: Vec<String> = entries.iter().map(|e| e.org_id.clone()).collect();
    let block_size: Vec<u64> = entries.iter().map(|e| e.block_size as u64).collect();
    let features: Vec<String> = entries.iter().map(|e| e.features.join(",")).collect();
    let speculation_ratio: Vec<Option<f64>> = entries.iter().map(|e| e.speculation_ratio).collect();
    let version: Vec<String> = entries.iter().map(|e| e.__version__.clone()).collect();

    let batch = RecordBatch::try_new(
        Arc::new(schema),
        vec![
            Arc::new(UInt64Array::from(timestamp_start)),
            Arc::new(UInt64Array::from(duration_e2e)),
            Arc::new(UInt64Array::from(duration_ttft)),
            Arc::new(UInt64Array::from(input_tokens)),
            Arc::new(UInt64Array::from(output_tokens)),
            Arc::new(total_hashes),
            Arc::new(StringArray::from(provided_session_id)),
            Arc::new(UInt8Array::from(request_canceled)),
            Arc::new(UInt64Array::from(cached_tokens_ref)),
            Arc::new(StringArray::from(model_name)),
            Arc::new(StringArray::from(org_id)),
            Arc::new(UInt64Array::from(block_size)),
            Arc::new(StringArray::from(features)),
            Arc::new(Float64Array::from(speculation_ratio)),
            Arc::new(StringArray::from(version)),
        ],
    )?;

    let mut buf = Vec::new();
    let props = WriterProperties::builder()
        .set_compression(Compression::SNAPPY)
        .build();
    let mut writer = ArrowWriter::try_new(&mut buf, batch.schema(), Some(props))?;
    writer.write(&batch)?;
    writer.close()?;

    Ok(buf)
}

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 read_credential_from_path(path: &str) -> anyhow::Result<String> {
    let content = std::fs::read_to_string(path)
        .map_err(|e| anyhow::anyhow!("failed to read credential from {path}: {e}"))?;
    Ok(content.trim().to_string())
}

fn resolve_credential(value: &str) -> String {
    // If value starts with / and exists as a file, read the credential from the file
    if value.starts_with('/') && std::path::Path::new(value).exists() {
        match read_credential_from_path(value) {
            Ok(credential) => credential,
            Err(e) => {
                eprintln!("performance-metrics: warning: {e}");
                value.to_string()
            }
        }
    } else {
        value.to_string()
    }
}

fn build_event_recorder_sink(
    local_output_dir: Option<PathBuf>,
    s3_config: Option<S3SinkConfig>,
    file_format: FileFormat,
) -> 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);
        if let Some(endpoint) = config.endpoint {
            builder = builder.with_endpoint(endpoint);
        }
        if let Some(access_key_id) = config.access_key_id {
            let resolved = resolve_credential(&access_key_id);
            builder = builder.with_access_key_id(resolved);
        }
        if let Some(secret_access_key) = config.secret_access_key {
            let resolved = resolve_credential(&secret_access_key);
            builder = builder.with_secret_access_key(resolved);
        }
        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,
        file_format,
    })
}

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, vec![]);
        // With DEFAULT_BLOCK_SIZE=64, we expect:
        // - 140 tokens / 64 = 2.18 -> 2 hashes (first call)
        // - 260 tokens / 64 = 4.06 -> 4 hashes (second call)
        // Total: 6 hashes
        trace.record_output_tokens(&(0..140).collect::<Vec<u32>>(), false, None);
        trace.record_output_tokens(&(0..260).collect::<Vec<u32>>(), false, None);
        trace.success();
        assert_eq!(trace.total_hashes.len(), 5);
    }

    #[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, vec![]);
        let mut t2 = salted.new_trace(&[], None, vec![]);
        t1.record_output_tokens(&[1, 2, 3, 4, 5], true, None);
        t2.record_output_tokens(&[1, 2, 3, 4, 5], true, None);
        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, vec![]);
        a.record_output_tokens(output_ids.as_slice(), true, None);
        a.success();

        let mut b = factory.new_trace(input_ids.as_slice(), None, vec![]);
        b.record_output_tokens(output_ids.as_slice(), true, None);
        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, vec![]);
        a.record_output_tokens(&[7, 8], true, None);
        a.success();

        let mut b = factory.new_trace(&[3, 4], None, vec![]);
        b.record_output_tokens(&[7, 8], true, None);
        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, vec![]);
        a.record_output_tokens(&[1], true, None);
        a.success();

        let mut b = factory.new_trace(&[], None, vec![]);
        b.record_output_tokens(&[2], true, None);
        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, vec![]);
        cumulative.record_output_tokens(&[1, 2, 3], false, None);
        cumulative.record_output_tokens(&[1, 2, 3, 4, 5], false, None);
        cumulative.record_output_tokens(&[1, 2, 3, 4, 5, 6], false, None);
        cumulative.success();

        let mut diff = factory.new_trace(&[], None, vec![]);
        diff.record_output_tokens(&[1, 2, 3], true, None);
        diff.record_output_tokens(&[4, 5], true, None);
        diff.record_output_tokens(&[6], true, None);
        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, vec![]);
            trace.record_output_tokens(&[4, 5], true, None);
            // 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, vec![]);
                trace.record_output_tokens(&[4], true, None);
                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(),
            }),
            ..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, None);
        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 input_len =
            find_metric_snapshot(&snapshot, "request_metrics_input_length_tokens", None);
        let output_len =
            find_metric_snapshot(&snapshot, "request_metrics_output_length_tokens", None);
        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!(input_len.average.unwrap_or(0.0) >= 0.0);
        assert!(output_len.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, None);
        req.record_tokens(&[103, 104], None, true, None);
        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_none = find_metric_snapshot(&snapshot, "request_metrics_ttft_ms", Some("none"));
        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_none = find_metric_snapshot(&snapshot, "request_metrics_itl_ms", Some("none"));
        let itl_xgrammar =
            find_metric_snapshot(&snapshot, "request_metrics_itl_ms", Some("xgrammar"));
        let request_all = find_metric_snapshot(&snapshot, "request_metrics_request", Some("all"));
        let request_tools =
            find_metric_snapshot(&snapshot, "request_metrics_request", Some("tools"));
        let request_image =
            find_metric_snapshot(&snapshot, "request_metrics_request", Some("image"));
        let request_none = find_metric_snapshot(&snapshot, "request_metrics_request", Some("none"));
        let request_xgrammar =
            find_metric_snapshot(&snapshot, "request_metrics_request", 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_none.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_none.average.unwrap_or(0.0), 0.0);
        assert_eq!(itl_xgrammar.average.unwrap_or(0.0), 0.0);
        assert!(request_all.rate_per_second.unwrap_or(0.0) >= 0.0);
        assert!(request_tools.rate_per_second.unwrap_or(0.0) >= 0.0);
        assert!(request_image.rate_per_second.unwrap_or(0.0) >= 0.0);
        assert_eq!(request_none.rate_per_second.unwrap_or(0.0), 0.0);
        assert_eq!(request_xgrammar.rate_per_second.unwrap_or(0.0), 0.0);
    }

    #[test]
    fn request_metrics_none_feature_series_records_for_plain_requests() {
        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, None);
        req.success();

        let snapshot = factory.snapshot();
        let ttft_none = find_metric_snapshot(&snapshot, "request_metrics_ttft_ms", Some("none"));
        let itl_none = find_metric_snapshot(&snapshot, "request_metrics_itl_ms", Some("none"));
        let request_none = find_metric_snapshot(&snapshot, "request_metrics_request", Some("none"));
        let request_all = find_metric_snapshot(&snapshot, "request_metrics_request", Some("all"));

        assert!(ttft_none.average.unwrap_or(0.0) >= 0.0);
        assert!(itl_none.average.unwrap_or(0.0) >= 0.0);
        assert!(request_none.rate_per_second.unwrap_or(0.0) >= 0.0);
        assert!(request_all.rate_per_second.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, None);
        req.record_tokens(&[12], None, true, None);
        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_pre_20_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_pre_20_token_cancellations, 0);
        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 cancellation_before_20_tokens_tracks_pre_20_bucket() {
        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, 12], Some(2), true, None);
        req.cancel();

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

        let cancellation =
            find_metric_snapshot(&snapshot, "request_metrics_pre_20_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_includes_null_session_id_when_absent() {
        let entry = RequestLogEntry {
            timestamp_start_unix_ms: 1,
            duration_e2e_ms: 100,
            duration_ttft_ms: 50,
            input_tokens: 3,
            output_tokens: 4,
            total_hashes: vec![[
                0xab, 0xcd, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00, 0x00,
                0x00, 0x00, 0x00, 0x00,
            ]],
            provided_session_id: None,
            request_canceled: false,
            cached_tokens_reference: None,
            model_name: "unset_model_name".to_string(),
            org_id: String::new(),
            block_size: DEFAULT_BLOCK_SIZE,
            features: vec![],
            speculation_ratio: None,
            __version__: DEFAULT_REQUEST_LOG_VERSION.to_string(),
        };

        let value = serde_json::to_value(entry).unwrap();
        // Now that skip_serializing_if is removed, null values should be included
        assert!(value.get("provided_session_id").is_some());
        assert!(value.get("provided_session_id").unwrap().is_null());
        assert_eq!(
            value
                .get("total_hashes")
                .and_then(|v| v.as_array())
                .and_then(|v| v.first())
                .and_then(|v| v.as_str()),
            Some("abcd000000000000000000000000000000000000000000000000000000000000")
        );
    }

    #[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"), vec![]);
        let entry = trace.build_log_entry(false);
        let expected = hash_session_id("session-123", None);

        assert_eq!(
            entry.provided_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"), vec![]);
        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.provided_session_id.as_deref(),
            Some(expected.as_str())
        );
    }

    #[test]
    fn request_trace_hashes_model_name() {
        let model_name = "my-sensitive-model-v1";
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            model_name: model_name.to_string(),
            ..Default::default()
        })
        .unwrap();

        // The model_name should be hashed (64 hex characters for blake3)
        assert_eq!(factory.model_name.len(), 64);
        assert_ne!(factory.model_name, model_name);

        // Verify it's a valid hex string
        assert!(factory.model_name.chars().all(|c| c.is_ascii_hexdigit()));

        // Verify the hash is correct
        let expected_hash = blake3::hash(model_name.as_bytes()).to_hex().to_string();
        assert_eq!(factory.model_name, expected_hash);
    }

    #[test]
    fn request_trace_hashes_org_id() {
        let org_id = "my-sensitive-org-123";
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            model_name: "test-model".to_string(),
            org_id: org_id.to_string(),
            ..Default::default()
        })
        .unwrap();

        // The org_id should be hashed (64 hex characters for blake3)
        assert_eq!(factory.org_id.len(), 64);
        assert_ne!(factory.org_id, org_id);

        // Verify it's a valid hex string
        assert!(factory.org_id.chars().all(|c| c.is_ascii_hexdigit()));

        // Verify the hash is correct
        let expected_hash = blake3::hash(org_id.as_bytes()).to_hex().to_string();
        assert_eq!(factory.org_id, expected_hash);
    }

    #[test]
    fn request_trace_org_id_is_empty_when_not_provided() {
        let factory = RequestTraceFactory::new(RequestTraceOptions {
            model_name: "test-model".to_string(),
            org_id: String::new(),
            ..Default::default()
        })
        .unwrap();

        // Empty string should be hashed (not empty)
        assert_eq!(factory.org_id.len(), 64);
        let expected_hash = blake3::hash(b"").to_hex().to_string();
        assert_eq!(factory.org_id, expected_hash);
    }

    #[test]
    fn request_trace_parquet_schema_matches_data() {
        // Test that serialize_to_parquet produces valid parquet with matching schema
        // Use valid 64-character hex strings (representing 32-byte blake3 hashes)
        let entries = vec![RequestLogEntry {
            timestamp_start_unix_ms: 1234567890,
            duration_e2e_ms: 100,
            duration_ttft_ms: 50,
            input_tokens: 10,
            output_tokens: 20,
            total_hashes: vec![
                [
                    0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3,
                    0xb4, 0xc5, 0xd6, 0xe7, 0xf8, 0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6,
                    0xc7, 0xd8, 0xe9, 0xf0, 0xa1, 0xb2,
                ],
                [
                    0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3, 0xb4, 0xc5, 0xd6, 0xe7,
                    0xf8, 0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8, 0xe9, 0xf0,
                    0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6,
                ],
                [
                    0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8, 0xe9, 0xf0, 0xa1,
                    0xb2, 0xc3, 0xd4, 0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3, 0xb4,
                    0xc5, 0xd6, 0xe7, 0xf8, 0xa9, 0xb0,
                ],
            ],
            provided_session_id: Some("session-123".to_string()),
            request_canceled: false,
            cached_tokens_reference: Some(5),
            model_name: "test-model".to_string(),
            org_id: "test-org-123".to_string(),
            block_size: DEFAULT_BLOCK_SIZE,
            features: vec!["feature1".to_string()],
            speculation_ratio: Some(3.5),
            __version__: DEFAULT_REQUEST_LOG_VERSION.to_string(),
        }];

        let parquet_bytes = serialize_to_parquet(&entries).expect("Should serialize to parquet");
        assert!(
            !parquet_bytes.is_empty(),
            "Parquet bytes should not be empty"
        );

        // Verify it's valid parquet by reading it back
        use parquet::file::reader::{FileReader, SerializedFileReader};
        let reader = SerializedFileReader::new(bytes::Bytes::from(parquet_bytes))
            .expect("Should be valid parquet");
        let schema = reader.metadata().file_metadata().schema();

        // Verify all expected fields are present
        let expected_fields = vec![
            "timestamp_start_unix_ms",
            "duration_e2e_ms",
            "duration_ttft_ms",
            "input_tokens",
            "output_tokens",
            "total_hashes",
            "provided_session_id",
            "request_canceled",
            "cached_tokens_reference",
            "model_name",
            "org_id",
            "block_size",
            "features",
            "speculation_ratio",
            "__version__",
        ];

        assert_eq!(
            schema.get_fields().len(),
            expected_fields.len(),
            "Schema field count should match expected"
        );

        for (i, expected_name) in expected_fields.iter().enumerate() {
            let actual_name = schema.get_fields()[i].name();
            assert_eq!(
                actual_name, *expected_name,
                "Field at position {} should match",
                i
            );
        }
    }

    #[test]
    fn request_trace_parquet_hashes_roundtrip() {
        use arrow::array::ListArray;
        use parquet::file::reader::{FileReader, SerializedFileReader};
        use std::io::Write;

        // Create test entries with multiple hashes (using valid 64-char hex strings)
        // These represent blake3-256 hashes (32 bytes = 64 hex chars)
        let entries = vec![
            RequestLogEntry {
                timestamp_start_unix_ms: 1000,
                duration_e2e_ms: 100,
                duration_ttft_ms: 50,
                input_tokens: 10,
                output_tokens: 20,
                total_hashes: vec![
                    [
                        0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2,
                        0xa3, 0xb4, 0xc5, 0xd6, 0xe7, 0xf8, 0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4,
                        0xa5, 0xb6, 0xc7, 0xd8, 0xe9, 0xf0, 0xa1, 0xb2,
                    ],
                    [
                        0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3, 0xb4, 0xc5, 0xd6,
                        0xe7, 0xf8, 0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8,
                        0xe9, 0xf0, 0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6,
                    ],
                ],
                provided_session_id: Some("session-123".to_string()),
                request_canceled: false,
                cached_tokens_reference: Some(5),
                model_name: "test-model".to_string(),
                org_id: "test-org".to_string(),
                block_size: DEFAULT_BLOCK_SIZE,
                features: vec!["feature1".to_string()],
                speculation_ratio: None,
                __version__: DEFAULT_REQUEST_LOG_VERSION.to_string(),
            },
            RequestLogEntry {
                timestamp_start_unix_ms: 2000,
                duration_e2e_ms: 200,
                duration_ttft_ms: 100,
                input_tokens: 20,
                output_tokens: 40,
                total_hashes: vec![
                    [
                        0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2,
                        0xa3, 0xb4, 0xc5, 0xd6, 0xe7, 0xf8, 0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4,
                        0xa5, 0xb6, 0xc7, 0xd8, 0xe9, 0xf0, 0xa1, 0xb2,
                    ],
                    [
                        0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3, 0xb4, 0xc5, 0xd6,
                        0xe7, 0xf8, 0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8,
                        0xe9, 0xf0, 0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6,
                    ],
                    [
                        0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8, 0xe9, 0xf0,
                        0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2,
                        0xa3, 0xb4, 0xc5, 0xd6, 0xe7, 0xf8, 0xa9, 0xb0,
                    ],
                ],
                provided_session_id: None,
                request_canceled: true,
                cached_tokens_reference: None,
                model_name: "test-model-2".to_string(),
                org_id: "test-org-2".to_string(),
                block_size: DEFAULT_BLOCK_SIZE,
                features: vec![],
                speculation_ratio: Some(2.5),
                __version__: DEFAULT_REQUEST_LOG_VERSION.to_string(),
            },
        ];

        // Serialize to parquet
        let parquet_bytes = serialize_to_parquet(&entries).expect("Should serialize to parquet");

        // Write to temp file
        let temp_dir =
            std::env::temp_dir().join(format!("pm-parquet-test-{}", unix_ms(SystemTime::now())));
        std::fs::create_dir_all(&temp_dir).unwrap();
        let temp_file = temp_dir.join("test.parquet");
        let mut file = std::fs::File::create(&temp_file).unwrap();
        file.write_all(&parquet_bytes).unwrap();
        file.flush().unwrap();
        drop(file);

        // Read back using parquet reader
        let reader = SerializedFileReader::new(std::fs::File::open(&temp_file).unwrap())
            .expect("Should be valid parquet");

        // Get the row count
        let metadata = reader.metadata().file_metadata();
        assert_eq!(metadata.num_rows(), 2, "Should have 2 rows");

        // Read and validate the data using Arrow
        use arrow::array::{Array, FixedSizeBinaryArray};
        use parquet::arrow::arrow_reader::ParquetRecordBatchReaderBuilder;
        let parquet_file = std::fs::File::open(&temp_file).unwrap();
        let builder = ParquetRecordBatchReaderBuilder::try_new(parquet_file).unwrap();
        let record_reader = builder.build().unwrap();

        // Collect all batches
        let mut batches = Vec::new();
        for batch in record_reader {
            batches.push(batch.unwrap());
        }
        assert!(!batches.is_empty(), "Should have at least one batch");

        // Get the first batch and verify the total_hashes column
        let batch = &batches[0];
        let total_hashes_column = batch.column_by_name("total_hashes").unwrap();
        let total_hashes_list = total_hashes_column
            .as_any()
            .downcast_ref::<ListArray>()
            .unwrap();

        // Verify first row has 2 hashes
        let first_row_hashes = total_hashes_list.value(0);
        let first_row_binary = first_row_hashes
            .as_any()
            .downcast_ref::<FixedSizeBinaryArray>()
            .unwrap();
        assert_eq!(first_row_binary.len(), 2, "First row should have 2 hashes");

        // Expected hash values as raw bytes
        let expected_hash1: HashBytes = [
            0xa1, 0xb2, 0xc3, 0xd4, 0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3, 0xb4,
            0xc5, 0xd6, 0xe7, 0xf8, 0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8,
            0xe9, 0xf0, 0xa1, 0xb2,
        ];
        let expected_hash2: HashBytes = [
            0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3, 0xb4, 0xc5, 0xd6, 0xe7, 0xf8,
            0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8, 0xe9, 0xf0, 0xa1, 0xb2,
            0xc3, 0xd4, 0xe5, 0xf6,
        ];

        // Verify the binary values
        assert_eq!(first_row_binary.value(0), expected_hash1.as_slice());
        assert_eq!(first_row_binary.value(1), expected_hash2.as_slice());

        // Verify second row has 3 hashes
        let second_row_hashes = total_hashes_list.value(1);
        let second_row_binary = second_row_hashes
            .as_any()
            .downcast_ref::<FixedSizeBinaryArray>()
            .unwrap();
        assert_eq!(
            second_row_binary.len(),
            3,
            "Second row should have 3 hashes"
        );

        // Verify the first hash of second row matches
        assert_eq!(second_row_binary.value(0), expected_hash1.as_slice());

        // Verify the third hash
        let expected_hash3: HashBytes = [
            0xa9, 0xb0, 0xc1, 0xd2, 0xe3, 0xf4, 0xa5, 0xb6, 0xc7, 0xd8, 0xe9, 0xf0, 0xa1, 0xb2,
            0xc3, 0xd4, 0xe5, 0xf6, 0xa7, 0xb8, 0xc9, 0xd0, 0xe1, 0xf2, 0xa3, 0xb4, 0xc5, 0xd6,
            0xe7, 0xf8, 0xa9, 0xb0,
        ];
        assert_eq!(second_row_binary.value(2), expected_hash3.as_slice());

        // Cleanup
        std::fs::remove_dir_all(&temp_dir).unwrap();
    }

    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
    }
}
