use numpy::PyReadonlyArray1;
use performance_metrics_core::{
    __private::{
        DEFAULT_BLOCK_SIZE as RS_DEFAULT_BLOCK_SIZE,
        DEFAULT_LOCAL_OUTPUT_DIR as RS_DEFAULT_LOCAL_OUTPUT_DIR,
        DEFAULT_REQUEST_LOG_BATCH_SIZE as RS_DEFAULT_REQUEST_LOG_BATCH_SIZE,
        DEFAULT_REQUEST_LOG_QUEUE_CAPACITY as RS_DEFAULT_REQUEST_LOG_QUEUE_CAPACITY,
        DistributionMetricHandle as RsDistributionMetricHandle, FileFormat as RsFileFormat,
        PerformanceMetricsRegistry as RsPerformanceMetricsRegistry,
        RateMetricHandle as RsRateMetricHandle, RatioMetricHandle as RsRatioMetricHandle,
        RequestMetricsOptions as RsRequestMetricsOptions,
        RequestTraceOptions as RsRequestTraceOptions, S3SinkConfig as RsS3SinkConfig,
    },
    AnonymizeConfig as RsAnonymizeConfig, MetricSeriesSnapshot as RsMetricSeriesSnapshot,
    ParquetCompactResult as RsParquetCompactResult, RequestFeatures as RsRequestFeatures,
    RequestMetricsFactory as RsRequestMetricsFactory,
    RequestMetricsRequest as RsRequestMetricsRequest,
    S3ParquetCompactConfig as RsS3ParquetCompactConfig,
    compact_s3_parquet_to_local as RsCompactS3ParquetToLocal,
};
use prometheus::proto::{Counter, Gauge, LabelPair, Metric, MetricFamily, MetricType};
use prometheus::{Encoder, TextEncoder};
use pyo3::types::{PyBytes, PyDict};
use pyo3::{exceptions::PyValueError, prelude::*};
use std::collections::{BTreeMap, HashMap};

#[pyclass(name = "RateMetric")]
pub struct PyRateMetric {
    inner: RsRateMetricHandle,
}

#[pymethods]
impl PyRateMetric {
    #[pyo3(signature = (count = 1))]
    fn record_count(&self, count: u64) -> PyResult<()> {
        self.inner
            .record_count(count)
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }
}

#[pyclass(name = "DistributionMetric")]
pub struct PyDistributionMetric {
    inner: RsDistributionMetricHandle,
}

#[pymethods]
impl PyDistributionMetric {
    fn record_value(&self, value: f64) -> PyResult<()> {
        self.inner
            .record_value(value)
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }
}

#[pyclass(name = "RatioMetric")]
pub struct PyRatioMetric {
    inner: RsRatioMetricHandle,
}

#[pymethods]
impl PyRatioMetric {
    fn record_ratio(&self, numerator: f64, denominator: f64) -> PyResult<()> {
        self.inner
            .record_ratio(numerator, denominator)
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }
}

#[pyclass(name = "PerformanceMetricsRegistry")]
pub struct PyPerformanceMetricsRegistry {
    inner: RsPerformanceMetricsRegistry,
}

#[pymethods]
impl PyPerformanceMetricsRegistry {
    #[new]
    fn new() -> Self {
        Self {
            inner: RsPerformanceMetricsRegistry::default(),
        }
    }

    #[pyo3(signature = (name, quantiles = None, sample_period_seconds = 1.0, window_seconds = None))]
    fn new_rate_metric(
        &self,
        name: String,
        quantiles: Option<Vec<f64>>,
        sample_period_seconds: f64,
        window_seconds: Option<f64>,
    ) -> PyResult<PyRateMetric> {
        self.inner
            .new_rate_metric(
                name,
                quantiles.unwrap_or_default(),
                Some(sample_period_seconds),
                window_seconds,
            )
            .map(|inner| PyRateMetric { inner })
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }

    #[pyo3(signature = (name, quantiles = None, window_seconds = None))]
    fn new_distribution_metric(
        &self,
        name: String,
        quantiles: Option<Vec<f64>>,
        window_seconds: Option<f64>,
    ) -> PyResult<PyDistributionMetric> {
        self.inner
            .new_distribution_metric(name, quantiles.unwrap_or_default(), window_seconds)
            .map(|inner| PyDistributionMetric { inner })
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }

    #[pyo3(signature = (name, quantiles = None, window_seconds = None))]
    fn new_ratio_metric(
        &self,
        name: String,
        quantiles: Option<Vec<f64>>,
        window_seconds: Option<f64>,
    ) -> PyResult<PyRatioMetric> {
        self.inner
            .new_ratio_metric(name, quantiles.unwrap_or_default(), window_seconds)
            .map(|inner| PyRatioMetric { inner })
            .map_err(|e| PyValueError::new_err(e.to_string()))
    }

    fn snapshot_all(&self, py: Python<'_>) -> PyResult<PyObject> {
        let out = PyDict::new(py);
        for (name, metric) in self.inner.snapshot_all() {
            out.set_item(name, Py::new(py, metric_snapshot_to_typed(metric))?)?;
        }
        Ok(out.into_any().unbind())
    }

    fn snapshot_all_series(&self, py: Python<'_>) -> PyResult<PyObject> {
        let out = self
            .inner
            .snapshot_all_series()
            .into_iter()
            .map(metric_series_snapshot_to_typed)
            .map(|series| Py::new(py, series))
            .collect::<PyResult<Vec<_>>>()?;
        Ok(out.into_pyobject(py)?.unbind())
    }

    fn unregister_metric(&self, name: String) {
        self.inner.unregister_metric(name.as_str());
    }
}

#[pyclass(name = "RequestMetricsFactory")]
pub struct PyRequestMetricsFactory {
    inner: RsRequestMetricsFactory,
    request_log_enabled: bool,
}

#[pyclass(name = "RequestMetrics")]
pub struct PyRequestMetricsRequest {
    inner: RsRequestMetricsRequest,
}

#[pyclass(name = "MetricSnapshot")]
#[derive(Clone)]
pub struct PyMetricSnapshot {
    #[pyo3(get)]
    kind: String,
    #[pyo3(get)]
    rate_per_second: Option<f64>,
    #[pyo3(get)]
    average: Option<f64>,
    #[pyo3(get)]
    quantiles: HashMap<String, f64>,
    #[pyo3(get)]
    numerator_sum: Option<f64>,
    #[pyo3(get)]
    denominator_sum: Option<f64>,
    #[pyo3(get)]
    ratio: Option<f64>,
}

#[pyclass(name = "RequestMetricsTotalsSnapshot")]
#[derive(Clone)]
pub struct PyRequestMetricsTotalsSnapshot {
    #[pyo3(get)]
    total_requests: u64,
    #[pyo3(get)]
    total_successful_requests: u64,
    #[pyo3(get)]
    total_pre_first_token_cancellations: u64,
    #[pyo3(get)]
    total_pre_20_token_cancellations: u64,
    #[pyo3(get)]
    total_mid_stream_cancellations: u64,
    #[pyo3(get)]
    total_input_tokens: u64,
    #[pyo3(get)]
    total_output_tokens: u64,
    #[pyo3(get)]
    total_net_new_input_tokens: u64,
    #[pyo3(get)]
    total_kv_cache_hits_tokens: u64,
    #[pyo3(get)]
    total_kv_cache_lookup_tokens: u64,
    #[pyo3(get)]
    total_itl_samples_recorded: u64,
    #[pyo3(get)]
    kv_cache_hit_rate_lifetime: Option<f64>,
}

#[pyclass(name = "RequestMetricsSnapshot")]
pub struct PyRequestMetricsSnapshot {
    series: Vec<PyMetricSeriesSnapshot>,
    totals: PyRequestMetricsTotalsSnapshot,
}

#[pyclass(name = "CompactResult")]
#[derive(Clone)]
pub struct PyCompactResult {
    #[pyo3(get)]
    source_objects: usize,
    #[pyo3(get)]
    rows_written: usize,
    #[pyo3(get)]
    output_path: String,
    #[pyo3(get)]
    min_timestamp: u64,
}

#[pyclass(name = "MetricSeriesSnapshot")]
#[derive(Clone)]
pub struct PyMetricSeriesSnapshot {
    #[pyo3(get)]
    name: String,
    #[pyo3(get)]
    labels: BTreeMap<String, String>,
    snapshot: PyMetricSnapshot,
}

#[pyclass(name = "PrometheusSample")]
#[derive(Clone)]
pub struct PyPrometheusSample {
    #[pyo3(get)]
    name: String,
    #[pyo3(get)]
    value: f64,
    #[pyo3(get)]
    labels: HashMap<String, String>,
    #[pyo3(get)]
    metric_type: String,
}

#[pymethods]
impl PyRequestMetricsSnapshot {
    fn series(&self, py: Python<'_>) -> PyResult<PyObject> {
        let out = self
            .series
            .iter()
            .cloned()
            .map(|series| Py::new(py, series))
            .collect::<PyResult<Vec<_>>>()?;
        Ok(out.into_pyobject(py)?.unbind())
    }

    fn totals(&self, py: Python<'_>) -> PyResult<Py<PyRequestMetricsTotalsSnapshot>> {
        Py::new(py, self.totals.clone())
    }
}

#[pymethods]
impl PyMetricSeriesSnapshot {
    fn snapshot(&self, py: Python<'_>) -> PyResult<Py<PyMetricSnapshot>> {
        Py::new(py, self.snapshot.clone())
    }
}

#[pymethods]
impl PyRequestMetricsFactory {
    #[new]
    #[allow(clippy::too_many_arguments)]
    #[pyo3(signature = (
        block_size = RS_DEFAULT_BLOCK_SIZE,
        request_log_enabled = false,
        request_log_hash_salt = None,
        request_log_batch_size = RS_DEFAULT_REQUEST_LOG_BATCH_SIZE,
        request_log_queue_capacity = RS_DEFAULT_REQUEST_LOG_QUEUE_CAPACITY,
        local_file = true,
        request_log_output_dir = RS_DEFAULT_LOCAL_OUTPUT_DIR.to_string(),
        s3 = false,
        s3_bucket = None,
        s3_region = "us-east-1".to_string(),
        s3_endpoint = None,
        s3_access_key_id = None,
        s3_secret_access_key = None,
        s3_prefix = "llm-request-metrics".to_string(),
        file_format = "json".to_string(),
        model_name = Some("unset_model_name".to_string()),
        org_id = String::new(),
        metric_prefix = "request_metrics".to_string(),
        metrics_window_seconds = 60.0,
        metrics_quantiles = vec![0.5, 0.9, 0.99],
        itl_sample_rate = 0.05
    ))]
    fn new(
        block_size: usize,
        request_log_enabled: bool,
        request_log_hash_salt: Option<String>,
        request_log_batch_size: usize,
        request_log_queue_capacity: usize,
        local_file: bool,
        request_log_output_dir: String,
        s3: bool,
        s3_bucket: Option<String>,
        s3_region: String,
        s3_endpoint: Option<String>,
        s3_access_key_id: Option<String>,
        s3_secret_access_key: Option<String>,
        s3_prefix: String,
        file_format: String,
        model_name: Option<String>,
        org_id: String,
        metric_prefix: String,
        metrics_window_seconds: f64,
        metrics_quantiles: Vec<f64>,
        itl_sample_rate: f64,
    ) -> PyResult<Self> {
        let local_output_dir = if local_file {
            Some(std::path::PathBuf::from(request_log_output_dir))
        } else {
            None
        };

        let file_format = match file_format.as_str() {
            "json" => RsFileFormat::Json,
            "parquet" => RsFileFormat::Parquet,
            _ => {
                return Err(PyValueError::new_err(format!(
                    "Invalid file_format: {}. Must be 'json' or 'parquet'",
                    file_format
                )));
            }
        };

        let s3_config = if s3 {
            let bucket = s3_bucket
                .ok_or_else(|| PyValueError::new_err("s3=True requires s3_bucket to be set"))?;
            Some(RsS3SinkConfig {
                bucket,
                region: s3_region,
                endpoint: s3_endpoint,
                access_key_id: s3_access_key_id,
                secret_access_key: s3_secret_access_key,
                prefix: s3_prefix,
            })
        } else {
            None
        };

        let kv_cache_hit_rate_quantiles = kv_cache_quantiles(metrics_quantiles.clone());

        let inner = RsRequestMetricsFactory::new(RsRequestMetricsOptions {
            trace: RsRequestTraceOptions {
                block_size,
                request_log_enabled,
                request_log_hash_salt,
                request_log_batch_size,
                request_log_queue_capacity,
                local_output_dir,
                s3_config,
                model_name: model_name.unwrap_or_else(|| "unset_model_name".to_string()),
                org_id,
                file_format,
            },
            metric_prefix,
            request_quantiles: metrics_quantiles.clone(),
            request_sample_period_seconds: Some(1.0),
            request_window_seconds: Some(metrics_window_seconds),
            input_tokens_quantiles: metrics_quantiles.clone(),
            input_tokens_sample_period_seconds: Some(1.0),
            input_tokens_window_seconds: Some(metrics_window_seconds),
            ttft_quantiles: metrics_quantiles.clone(),
            ttft_window_seconds: Some(metrics_window_seconds),
            ttft_per_input_token_quantiles: metrics_quantiles.clone(),
            ttft_per_input_token_window_seconds: Some(metrics_window_seconds),
            itl_quantiles: vec![0.1, 0.5, 0.9, 0.99, 0.999],
            itl_window_seconds: Some(metrics_window_seconds),
            pre_first_token_cancellation_quantiles: Vec::new(),
            pre_first_token_cancellation_sample_period_seconds: Some(1.0),
            pre_first_token_cancellation_window_seconds: Some(metrics_window_seconds),
            pre_20_token_cancellation_quantiles: Vec::new(),
            pre_20_token_cancellation_sample_period_seconds: Some(1.0),
            pre_20_token_cancellation_window_seconds: Some(metrics_window_seconds),
            mid_stream_cancellation_quantiles: Vec::new(),
            mid_stream_cancellation_sample_period_seconds: Some(1.0),
            mid_stream_cancellation_window_seconds: Some(metrics_window_seconds),
            successful_request_quantiles: Vec::new(),
            successful_request_sample_period_seconds: Some(1.0),
            successful_request_window_seconds: Some(metrics_window_seconds),
            kv_cache_hit_rate_quantiles,
            kv_cache_hit_rate_window_seconds: Some(metrics_window_seconds),
            itl_sample_rate,
        })
        .map_err(|e| PyValueError::new_err(e.to_string()))?;

        Ok(Self {
            inner,
            request_log_enabled,
        })
    }

    #[pyo3(signature = (input_token_ids, features = 0, session_id = None))]
    fn new_request(
        &self,
        input_token_ids: &Bound<'_, PyAny>,
        features: u8,
        session_id: Option<String>,
    ) -> PyResult<PyRequestMetricsRequest> {
        // Fast path: when logging is disabled, just get the length without copying
        if !self.request_log_enabled {
            // this branch is around 100-200x faster for 50k tokens
            // since we don't need to create a rust vec to get len() if log is disabled
            let count = extract_length(input_token_ids)?;
            return Ok(PyRequestMetricsRequest {
                inner: self.inner.new_request_with_count(
                    count,
                    RsRequestFeatures::from(features),
                    session_id.as_deref(),
                ),
            });
        }

        // Full path: extract and copy tokens for hashing
        let input_token_ids = extract_list_or_numpy_u32(input_token_ids)?;
        Ok(PyRequestMetricsRequest {
            inner: self.inner.new_request(
                input_token_ids.as_slice(),
                RsRequestFeatures::from(features),
                session_id.as_deref(),
            ),
        })
    }

    fn snapshot(&self) -> PyRequestMetricsSnapshot {
        request_metrics_snapshot_to_typed(self.inner.snapshot())
    }

    #[pyo3(signature = (base_labels = None))]
    fn prometheus_samples(
        &self,
        base_labels: Option<HashMap<String, String>>,
    ) -> Vec<PyPrometheusSample> {
        flatten_snapshot(self.inner.snapshot(), base_labels.as_ref())
    }

    #[pyo3(signature = (base_labels = None))]
    fn prometheus_strfmt(&self, base_labels: Option<HashMap<String, String>>) -> PyResult<String> {
        let samples = self.prometheus_samples(base_labels);
        encode_samples_as_prometheus_text(samples)
    }
}

#[pymethods]
impl PyRequestMetricsRequest {
    #[pyo3(signature = (token_ids, cached_tokens = None, is_diff = true, speculation_ratio = None))]
    fn record_tokens(
        &mut self,
        token_ids: &Bound<'_, PyAny>,
        cached_tokens: Option<u64>,
        is_diff: bool,
        speculation_ratio: Option<f64>,
    ) -> PyResult<()> {
        let token_ids = extract_list_or_numpy_u32(token_ids)?;
        self.inner.record_tokens(
            token_ids.as_slice(),
            cached_tokens,
            is_diff,
            speculation_ratio,
        );
        Ok(())
    }

    fn success(&mut self) {
        self.inner.success();
    }

    fn cancel(&mut self) {
        self.inner.cancel();
    }

    fn total_hashes<'py>(&self, py: Python<'py>) -> Vec<Bound<'py, PyBytes>> {
        self.inner
            .total_hashes()
            .iter()
            .map(|hash| PyBytes::new(py, hash))
            .collect()
    }

    fn features(&self) -> u8 {
        self.inner.features().bits()
    }
}

#[pyfunction]
#[allow(clippy::too_many_arguments)]
#[pyo3(signature = (
    s3_bucket,
    output_path,
    s3_region = "us-east-1".to_string(),
    s3_prefix = String::new(),
    s3_endpoint = None,
    s3_access_key_id = None,
    s3_secret_access_key = None,
    max_concurrent_downloads = 128,
    anonymize = false
))]
fn compact_s3_parquet_to_local(
    s3_bucket: String,
    output_path: String,
    s3_region: String,
    s3_prefix: String,
    s3_endpoint: Option<String>,
    s3_access_key_id: Option<String>,
    s3_secret_access_key: Option<String>,
    max_concurrent_downloads: usize,
    anonymize: bool,
) -> PyResult<PyCompactResult> {
    let anonymize_config = if anonymize {
        Some(RsAnonymizeConfig {
            columns_to_pseudonymize: vec![
                "provided_session_id".to_string(),
                "org_id".to_string(),
                "total_hashes".to_string(),
            ],
            columns_to_drop: vec!["org_id".to_string()],
        })
    } else {
        None
    };

    let result = RsCompactS3ParquetToLocal(RsS3ParquetCompactConfig {
        s3: RsS3SinkConfig {
            bucket: s3_bucket,
            region: s3_region,
            endpoint: s3_endpoint,
            access_key_id: s3_access_key_id,
            secret_access_key: s3_secret_access_key,
            prefix: s3_prefix,
        },
        output_path: std::path::PathBuf::from(output_path),
        max_concurrent_downloads,
        anonymize: anonymize_config,
    })
    .map_err(|e| PyValueError::new_err(e.to_string()))?;

    Ok(compact_result_to_typed(result))
}

#[pyfunction]
#[pyo3(signature = (
    input_path,
    output_path,
    columns_to_pseudonymize = vec!["provided_session_id".to_string(), "org_id".to_string(), "total_hashes".to_string()],
    columns_to_drop = vec!["org_id".to_string()],
    min_timestamp = 0
))]
fn anonymize_parquet(
    input_path: String,
    output_path: String,
    columns_to_pseudonymize: Vec<String>,
    columns_to_drop: Vec<String>,
    min_timestamp: u64,
) -> PyResult<u64> {
    let columns_to_pseudonymize_refs: Vec<&str> =
        columns_to_pseudonymize.iter().map(|s| s.as_str()).collect();
    let columns_to_drop_refs: Vec<&str> = columns_to_drop.iter().map(|s| s.as_str()).collect();

    performance_metrics_core::anonymize_parquet_file(
        std::path::Path::new(&input_path),
        std::path::Path::new(&output_path),
        &columns_to_pseudonymize_refs,
        &columns_to_drop_refs,
        min_timestamp,
    )
    .map_err(|e| PyValueError::new_err(e.to_string()))
}

#[pymodule]
fn _core(_py: Python<'_>, m: &Bound<'_, PyModule>) -> PyResult<()> {
    m.add_class::<PyPerformanceMetricsRegistry>()?;
    m.add_class::<PyRateMetric>()?;
    m.add_class::<PyDistributionMetric>()?;
    m.add_class::<PyRatioMetric>()?;
    m.add_class::<PyRequestMetricsFactory>()?;
    m.add_class::<PyRequestMetricsRequest>()?;
    m.add_class::<PyMetricSnapshot>()?;
    m.add_class::<PyMetricSeriesSnapshot>()?;
    m.add_class::<PyRequestMetricsTotalsSnapshot>()?;
    m.add_class::<PyRequestMetricsSnapshot>()?;
    m.add_class::<PyPrometheusSample>()?;
    m.add_class::<PyCompactResult>()?;
    m.add_function(wrap_pyfunction!(compact_s3_parquet_to_local, m)?)?;
    m.add_function(wrap_pyfunction!(anonymize_parquet, m)?)?;
    m.add("REQUEST_FEATURE_NONE", RsRequestFeatures::NONE.bits())?;
    m.add(
        "REQUEST_FEATURE_XGRAMMAR",
        RsRequestFeatures::XGRAMMAR.bits(),
    )?;
    m.add("REQUEST_FEATURE_TOOLS", RsRequestFeatures::TOOLS.bits())?;
    m.add("REQUEST_FEATURE_IMAGE", RsRequestFeatures::IMAGE.bits())?;
    m.add("__version__", env!("CARGO_PKG_VERSION"))?;
    Ok(())
}

/// Extract just the length from a Python list or NumPy array without copying data.
fn extract_length(token_ids: &Bound<'_, PyAny>) -> PyResult<usize> {
    // Try NumPy array first (fastest path)
    if let Ok(arr) = token_ids.extract::<PyReadonlyArray1<'_, u32>>() {
        return arr.len().map_err(|e| PyValueError::new_err(e.to_string()));
    }
    if let Ok(arr) = token_ids.extract::<PyReadonlyArray1<'_, i64>>() {
        return arr.len().map_err(|e| PyValueError::new_err(e.to_string()));
    }
    // Fall back to Python list
    if let Ok(list) = token_ids.downcast::<pyo3::types::PyList>() {
        return Ok(list.len());
    }
    Err(PyValueError::new_err(
        "token_ids must be list[int] or numpy.ndarray (uint32/int64)",
    ))
}

fn extract_list_or_numpy_u32(token_ids: &Bound<'_, PyAny>) -> PyResult<Vec<u32>> {
    if let Ok(arr) = token_ids.extract::<PyReadonlyArray1<'_, u32>>() {
        return arr
            .as_slice()
            .map(|s| s.to_vec())
            .map_err(|e| PyValueError::new_err(e.to_string()));
    }

    if let Ok(arr) = token_ids.extract::<PyReadonlyArray1<'_, i64>>() {
        let values = arr
            .as_slice()
            .map_err(|e| PyValueError::new_err(e.to_string()))?;
        let mut out = Vec::with_capacity(values.len());
        for &v in values {
            let u = u32::try_from(v)
                .map_err(|_| PyValueError::new_err("numpy int64 values must fit in u32"))?;
            out.push(u);
        }
        return Ok(out);
    }

    if let Ok(ids) = token_ids.extract::<Vec<u32>>() {
        return Ok(ids);
    }
    Err(PyValueError::new_err(
        "token_ids must be list[int] or numpy.ndarray (uint32/int64)",
    ))
}

fn metric_snapshot_to_typed(m: performance_metrics_core::MetricSnapshot) -> PyMetricSnapshot {
    PyMetricSnapshot {
        kind: match m.kind {
            performance_metrics_core::MetricKind::Rate => "rate".to_string(),
            performance_metrics_core::MetricKind::Distribution => "distribution".to_string(),
            performance_metrics_core::MetricKind::Ratio => "ratio".to_string(),
        },
        rate_per_second: m.rate_per_second,
        average: m.average,
        quantiles: m.quantiles,
        numerator_sum: m.numerator_sum,
        denominator_sum: m.denominator_sum,
        ratio: m.ratio,
    }
}

fn compact_result_to_typed(result: RsParquetCompactResult) -> PyCompactResult {
    PyCompactResult {
        source_objects: result.source_objects,
        rows_written: result.rows_written,
        output_path: result.output_path.display().to_string(),
        min_timestamp: result.min_timestamp,
    }
}

fn metric_series_snapshot_to_typed(series: RsMetricSeriesSnapshot) -> PyMetricSeriesSnapshot {
    PyMetricSeriesSnapshot {
        name: series.name,
        labels: series.labels,
        snapshot: metric_snapshot_to_typed(series.snapshot),
    }
}

fn request_metrics_snapshot_to_typed(
    snapshot: performance_metrics_core::RequestMetricsSnapshot,
) -> PyRequestMetricsSnapshot {
    PyRequestMetricsSnapshot {
        series: snapshot
            .series
            .into_iter()
            .map(metric_series_snapshot_to_typed)
            .collect(),
        totals: totals_to_typed(snapshot.totals),
    }
}

fn totals_to_typed(
    totals: performance_metrics_core::RequestMetricsTotalsSnapshot,
) -> PyRequestMetricsTotalsSnapshot {
    PyRequestMetricsTotalsSnapshot {
        total_requests: totals.total_requests,
        total_successful_requests: totals.total_successful_requests,
        total_pre_first_token_cancellations: totals.total_pre_first_token_cancellations,
        total_pre_20_token_cancellations: totals.total_pre_20_token_cancellations,
        total_mid_stream_cancellations: totals.total_mid_stream_cancellations,
        total_input_tokens: totals.total_input_tokens,
        total_output_tokens: totals.total_output_tokens,
        total_net_new_input_tokens: totals.total_net_new_input_tokens,
        total_kv_cache_hits_tokens: totals.total_kv_cache_hits_tokens,
        total_kv_cache_lookup_tokens: totals.total_kv_cache_lookup_tokens,
        total_itl_samples_recorded: totals.total_itl_samples_recorded,
        kv_cache_hit_rate_lifetime: totals.kv_cache_hit_rate_lifetime,
    }
}

fn flatten_snapshot(
    snapshot: performance_metrics_core::RequestMetricsSnapshot,
    base_labels: Option<&HashMap<String, String>>,
) -> Vec<PyPrometheusSample> {
    let mut out = flatten_series(snapshot.series, base_labels);
    out.extend(flatten_totals(snapshot.totals, base_labels));
    out
}

fn flatten_series(
    rolling: Vec<RsMetricSeriesSnapshot>,
    base_labels: Option<&HashMap<String, String>>,
) -> Vec<PyPrometheusSample> {
    let mut out = Vec::new();
    for series in rolling {
        let metric_name = series.name;
        let extra_labels = Some(
            series
                .labels
                .into_iter()
                .collect::<HashMap<String, String>>(),
        );
        let snapshot = series.snapshot;
        if let Some(v) = snapshot.rate_per_second {
            out.push(make_sample(
                format!(
                    "{}_rate_per_second",
                    sanitize_metric_name(metric_name.as_str())
                ),
                v,
                base_labels,
                extra_labels.clone(),
                "gauge",
            ));
        }
        if let Some(v) = snapshot.average {
            out.push(make_sample(
                format!("{}_average", sanitize_metric_name(metric_name.as_str())),
                v,
                base_labels,
                extra_labels.clone(),
                "gauge",
            ));
        }
        if let Some(v) = snapshot.numerator_sum {
            out.push(make_sample(
                format!(
                    "{}_numerator_sum",
                    sanitize_metric_name(metric_name.as_str())
                ),
                v,
                base_labels,
                extra_labels.clone(),
                "gauge",
            ));
        }
        if let Some(v) = snapshot.denominator_sum {
            out.push(make_sample(
                format!(
                    "{}_denominator_sum",
                    sanitize_metric_name(metric_name.as_str())
                ),
                v,
                base_labels,
                extra_labels.clone(),
                "gauge",
            ));
        }
        if let Some(v) = snapshot.ratio {
            out.push(make_sample(
                format!("{}_ratio", sanitize_metric_name(metric_name.as_str())),
                v,
                base_labels,
                extra_labels.clone(),
                "gauge",
            ));
        }
        for (q, v) in snapshot.quantiles {
            let mut extra = extra_labels.clone().unwrap_or_default();
            extra.insert("quantile".to_string(), q);
            out.push(make_sample(
                format!("{}_quantile", sanitize_metric_name(metric_name.as_str())),
                v,
                base_labels,
                Some(extra),
                "gauge",
            ));
        }
    }
    out
}

fn flatten_totals(
    totals: performance_metrics_core::RequestMetricsTotalsSnapshot,
    base_labels: Option<&HashMap<String, String>>,
) -> Vec<PyPrometheusSample> {
    let mut out = Vec::new();
    let counters = [
        ("total_requests", totals.total_requests as f64),
        (
            "total_successful_requests",
            totals.total_successful_requests as f64,
        ),
        (
            "total_pre_first_token_cancellations",
            totals.total_pre_first_token_cancellations as f64,
        ),
        (
            "total_pre_20_token_cancellations",
            totals.total_pre_20_token_cancellations as f64,
        ),
        (
            "total_mid_stream_cancellations",
            totals.total_mid_stream_cancellations as f64,
        ),
        ("total_input_tokens", totals.total_input_tokens as f64),
        ("total_output_tokens", totals.total_output_tokens as f64),
        (
            "total_net_new_input_tokens",
            totals.total_net_new_input_tokens as f64,
        ),
        (
            "total_kv_cache_hits_tokens",
            totals.total_kv_cache_hits_tokens as f64,
        ),
        (
            "total_kv_cache_lookup_tokens",
            totals.total_kv_cache_lookup_tokens as f64,
        ),
        (
            "total_itl_samples_recorded",
            totals.total_itl_samples_recorded as f64,
        ),
    ];
    for (name, value) in counters {
        out.push(make_sample(
            sanitize_metric_name(name),
            value,
            base_labels,
            None,
            "counter",
        ));
    }
    if let Some(v) = totals.kv_cache_hit_rate_lifetime {
        out.push(make_sample(
            "kv_cache_hit_rate_lifetime".to_string(),
            v,
            base_labels,
            None,
            "gauge",
        ));
    }
    out
}

fn make_sample(
    name: String,
    value: f64,
    base_labels: Option<&HashMap<String, String>>,
    extra_labels: Option<HashMap<String, String>>,
    metric_type: &str,
) -> PyPrometheusSample {
    let mut labels = base_labels.cloned().unwrap_or_default();
    if let Some(extra) = extra_labels {
        labels.extend(extra);
    }
    let labels = labels
        .into_iter()
        .map(|(k, v)| (sanitize_label_name(k.as_str()), v))
        .collect();
    PyPrometheusSample {
        name,
        value,
        labels,
        metric_type: metric_type.to_string(),
    }
}

fn sanitize_metric_name(value: &str) -> String {
    let mut out = String::new();
    for ch in value.chars() {
        if ch.is_ascii_alphanumeric() || ch == '_' {
            out.push(ch);
        } else {
            out.push('_');
        }
    }
    if out.is_empty() {
        return "metric".to_string();
    }
    if out.chars().next().is_some_and(|c| c.is_ascii_digit()) {
        out.insert(0, '_');
    }
    out
}

fn sanitize_label_name(value: &str) -> String {
    sanitize_metric_name(value)
}

fn kv_cache_quantiles(mut base: Vec<f64>) -> Vec<f64> {
    base.extend([0.01, 0.05, 0.1, 0.25, 0.5, 0.75, 0.9, 0.95, 0.99, 0.999]);
    base.retain(|q| q.is_finite());
    for q in &mut base {
        *q = q.clamp(0.0, 1.0);
    }
    base.sort_unstable_by(f64::total_cmp);
    base.dedup_by(|a, b| (*a - *b).abs() < 1e-9);
    base
}

fn metric_type_from_str(metric_type: &str) -> PyResult<MetricType> {
    match metric_type {
        "counter" => Ok(MetricType::COUNTER),
        "gauge" => Ok(MetricType::GAUGE),
        _ => Err(PyValueError::new_err(format!(
            "unsupported prometheus metric type: {metric_type}"
        ))),
    }
}

fn encode_samples_as_prometheus_text(samples: Vec<PyPrometheusSample>) -> PyResult<String> {
    let mut grouped: BTreeMap<(String, String), Vec<PyPrometheusSample>> = BTreeMap::new();
    for sample in samples {
        grouped
            .entry((sample.name.clone(), sample.metric_type.clone()))
            .or_default()
            .push(sample);
    }

    let mut families = Vec::with_capacity(grouped.len());
    for ((name, metric_type), mut samples) in grouped {
        let parsed_type = metric_type_from_str(metric_type.as_str())?;
        samples.sort_by(|a, b| {
            let mut a_labels = a.labels.iter().collect::<Vec<_>>();
            let mut b_labels = b.labels.iter().collect::<Vec<_>>();
            a_labels.sort_by(|x, y| x.0.cmp(y.0).then_with(|| x.1.cmp(y.1)));
            b_labels.sort_by(|x, y| x.0.cmp(y.0).then_with(|| x.1.cmp(y.1)));
            a_labels.cmp(&b_labels)
        });

        let mut family = MetricFamily::new();
        family.set_name(name.clone());
        family.set_help(name);
        family.set_field_type(parsed_type);

        for sample in samples {
            let mut metric = Metric::new();
            let mut labels = sample.labels.into_iter().collect::<Vec<_>>();
            labels.sort_by(|a, b| a.0.cmp(&b.0));
            for (k, v) in labels {
                let mut pair = LabelPair::new();
                pair.set_name(sanitize_label_name(k.as_str()));
                pair.set_value(v);
                metric.mut_label().push(pair);
            }

            match parsed_type {
                MetricType::COUNTER => {
                    let mut counter = Counter::new();
                    counter.set_value(sample.value);
                    metric.set_counter(counter);
                }
                MetricType::GAUGE => {
                    let mut gauge = Gauge::new();
                    gauge.set_value(sample.value);
                    metric.set_gauge(gauge);
                }
                _ => {
                    return Err(PyValueError::new_err(format!(
                        "unsupported prometheus type in encoder: {metric_type}"
                    )));
                }
            }

            family.mut_metric().push(metric);
        }
        families.push(family);
    }

    let mut out = Vec::new();
    TextEncoder::new()
        .encode(&families, &mut out)
        .map_err(|e| PyValueError::new_err(e.to_string()))?;
    String::from_utf8(out).map_err(|e| PyValueError::new_err(e.to_string()))
}

#[cfg(test)]
mod tests {
    use super::*;
    use performance_metrics_core::{
        __private::RequestMetricsOptions as RsRequestMetricsOptions,
        RequestMetricsFactory as RsRequestMetricsFactory,
    };
    use std::collections::{HashMap, HashSet};

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

        let mut base_labels = HashMap::new();
        base_labels.insert("engine".to_string(), "vllm".to_string());
        base_labels.insert("model".to_string(), "demo".to_string());

        let samples = flatten_snapshot(factory.snapshot(), Some(&base_labels));

        let mut series = HashSet::new();
        for sample in samples {
            let mut labels = sample.labels.into_iter().collect::<Vec<_>>();
            labels.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1)));
            series.insert((sample.name, labels));
        }

        assert!(series.len() < 10_000, "series count was {}", series.len());
    }

    #[test]
    fn request_metrics_feature_exports_use_feature_label() {
        let factory = RsRequestMetricsFactory::new(RsRequestMetricsOptions::default()).unwrap();
        let mut req = factory.new_request(
            &[1, 2, 3, 4],
            RsRequestFeatures::TOOLS | RsRequestFeatures::IMAGE,
            None,
        );
        req.record_tokens(&[10, 11, 12], Some(2), true, None);
        req.record_tokens(&[13], None, true, None);
        req.success();

        let mut base_labels = HashMap::new();
        base_labels.insert("engine".to_string(), "vllm".to_string());

        let samples = flatten_snapshot(factory.snapshot(), Some(&base_labels));

        let mut series = HashSet::new();
        for sample in samples {
            let mut labels = sample.labels.into_iter().collect::<Vec<_>>();
            labels.sort_by(|a, b| a.0.cmp(&b.0).then_with(|| a.1.cmp(&b.1)));
            series.insert((sample.name, labels));
        }

        assert!(series.contains(&(
            "request_metrics_ttft_ms_average".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "all".to_string()),
            ],
        )));
        assert!(series.contains(&(
            "request_metrics_ttft_ms_average".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "tools".to_string()),
            ],
        )));
        assert!(series.contains(&(
            "request_metrics_ttft_ms_average".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "image".to_string()),
            ],
        )));
        assert!(series.contains(&(
            "request_metrics_ttft_ms_average".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "none".to_string()),
            ],
        )));
        assert!(series.contains(&(
            "request_metrics_request_rate_per_second".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "all".to_string()),
            ],
        )));
        assert!(series.contains(&(
            "request_metrics_request_rate_per_second".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "tools".to_string()),
            ],
        )));
        assert!(series.contains(&(
            "request_metrics_request_rate_per_second".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "image".to_string()),
            ],
        )));
        assert!(series.contains(&(
            "request_metrics_request_rate_per_second".to_string(),
            vec![
                ("engine".to_string(), "vllm".to_string()),
                ("feature".to_string(), "none".to_string()),
            ],
        )));
        assert!(!series.iter().any(|(name, _)| name.contains("_with_tools")));
        assert!(!series.iter().any(|(name, _)| name.contains("_with_image")));
    }
}
