// SPDX-License-Identifier: Apache-2.0
// SPDX-FileCopyrightText: Copyright The Lance Authors

//! Bridge from the [`metrics`] crate facade to Python OpenTelemetry.
//!
//! Lance core publishes metrics through the global [`metrics`] facade without
//! choosing a backend. This module installs a process-global [`Recorder`] that
//! aggregates those metrics into lock-free cumulative storage, and exposes that
//! state to Python so the bindings can feed it into the user's OpenTelemetry
//! `MeterProvider`.
//!
//! While it targets OpenTelemetry on the Python side, the recorder is agnostic
//! to the metric *source*: it records any metric emitted through the facade,
//! keyed by name and labels. Object store metrics are the first producer, but
//! nothing here is specific to them. New metrics flow through automatically;
//! they only need to be described (see [`describe_all`]) so the Python layer can
//! discover their name, kind, and unit up front.
//!
//! ## Why pull, not push
//!
//! OpenTelemetry collects on its own schedule and invokes observable-instrument
//! callbacks at collection time. Cumulative counters map directly onto OTel's
//! `ObservableCounter` semantics. So the bridge aggregates in Rust and lets the
//! Python collection thread pull a [`snapshot`](snapshot_lance_metrics). The
//! snapshot is lock-free, but it still walks every registered series and
//! allocates owned copies of their names and labels, so it runs with the GIL
//! released to avoid stalling other Python threads during collection.
//!
//! ## Histograms
//!
//! OpenTelemetry has no asynchronous histogram instrument, so histograms cannot
//! be pulled as-is. Instead each histogram is aggregated into fixed buckets
//! (Prometheus style) and exposed as cumulative `le` bucket counts plus a count
//! and sum, which the Python layer surfaces as observable counters.

use std::collections::HashMap;
use std::sync::atomic::{AtomicU64, Ordering};
use std::sync::{Arc, LazyLock, Mutex, OnceLock, RwLock};

use metrics::{Counter, Gauge, Histogram, Key, KeyName, Metadata, Recorder, SharedString, Unit};
use metrics_util::registry::{Registry, Storage};
use pyo3::prelude::*;

/// Bucket boundaries used when a histogram has no registered bounds. Covers a
/// broad latency range so unknown histograms still produce useful buckets.
const DEFAULT_BOUNDS: &[f64] = &[
    0.005, 0.01, 0.025, 0.05, 0.1, 0.25, 0.5, 1.0, 2.5, 5.0, 10.0, 30.0, 60.0, 120.0, 300.0,
];

/// The kind of a metric, mirroring the three `metrics` instrument types.
#[derive(Clone, Copy)]
enum MetricKind {
    Counter,
    Gauge,
    Histogram,
}

impl MetricKind {
    fn as_str(self) -> &'static str {
        match self {
            Self::Counter => "counter",
            Self::Gauge => "gauge",
            Self::Histogram => "histogram",
        }
    }
}

/// Description of a metric, populated by the recorder's `describe_*` methods.
struct MetricDescription {
    kind: MetricKind,
    unit: Option<String>,
    description: String,
}

/// Catalog of described metrics, keyed by metric name. The Python layer reads
/// this to create one OpenTelemetry instrument per metric up front.
static CATALOG: LazyLock<Mutex<HashMap<String, MetricDescription>>> =
    LazyLock::new(|| Mutex::new(HashMap::new()));

/// Per-metric histogram bucket boundaries, keyed by metric name. Producers
/// register their recommended bounds before any metric is recorded.
static HISTOGRAM_BOUNDS: LazyLock<RwLock<HashMap<String, Arc<[f64]>>>> =
    LazyLock::new(|| RwLock::new(HashMap::new()));

/// The installed recorder's registry, available once installation succeeds.
static REGISTRY: OnceLock<Arc<Registry<Key, LanceStorage>>> = OnceLock::new();

fn bounds_for(name: &str) -> Arc<[f64]> {
    HISTOGRAM_BOUNDS
        .read()
        .unwrap()
        .get(name)
        .cloned()
        .unwrap_or_else(|| Arc::from(DEFAULT_BOUNDS))
}

/// A histogram that buckets samples at record time into fixed boundaries,
/// keeping a cumulative count and sum. Bucketing eagerly keeps memory bounded
/// (unlike retaining raw samples) and produces Prometheus-style `le` buckets.
struct BucketedHistogram {
    /// Sorted, finite upper bounds. A sample `v` falls in the first bucket whose
    /// bound is `>= v`; samples above all bounds fall in the implicit `+Inf`
    /// bucket stored as the final entry of `counts`.
    bounds: Arc<[f64]>,
    /// Per-bucket (non-cumulative) counts; length is `bounds.len() + 1`.
    counts: Box<[AtomicU64]>,
    count: AtomicU64,
    /// Running sum of recorded values, stored as `f64` bits (there is no atomic
    /// f64, so the bit pattern is held in a `u64`; see [`Self::add_to_sum`]).
    sum_bits: AtomicU64,
}

// All atomics here use `Ordering::Relaxed`: each metric counter is independent,
// so no happens-before relationship is needed between them, and a snapshot
// reader tolerates slightly stale values. This matches `metrics_util`'s
// `AtomicStorage`.

impl BucketedHistogram {
    fn new(bounds: Arc<[f64]>) -> Self {
        let counts = (0..bounds.len() + 1)
            .map(|_| AtomicU64::new(0))
            .collect::<Vec<_>>()
            .into_boxed_slice();
        Self {
            bounds,
            counts,
            count: AtomicU64::new(0),
            sum_bits: AtomicU64::new(0),
        }
    }

    fn add_to_sum(&self, value: f64) {
        // No atomic offers an f64 add, so read the current bit pattern, add in
        // float space, and CAS it back, retrying if another thread won the race.
        let mut current = self.sum_bits.load(Ordering::Relaxed);
        loop {
            let updated = (f64::from_bits(current) + value).to_bits();
            match self.sum_bits.compare_exchange_weak(
                current,
                updated,
                Ordering::Relaxed,
                Ordering::Relaxed,
            ) {
                Ok(_) => break,
                Err(actual) => current = actual,
            }
        }
    }

    /// Cumulative `le` buckets, total count, and sum at this instant.
    fn snapshot(&self) -> MetricValue {
        let mut cumulative = 0u64;
        let mut buckets = Vec::with_capacity(self.bounds.len() + 1);
        for (i, bound) in self.bounds.iter().enumerate() {
            cumulative += self.counts[i].load(Ordering::Relaxed);
            buckets.push((format!("{}", bound), cumulative));
        }
        cumulative += self.counts[self.bounds.len()].load(Ordering::Relaxed);
        buckets.push(("+Inf".to_string(), cumulative));
        MetricValue::Histogram {
            buckets,
            count: self.count.load(Ordering::Relaxed),
            sum: f64::from_bits(self.sum_bits.load(Ordering::Relaxed)),
        }
    }
}

impl metrics::HistogramFn for BucketedHistogram {
    fn record(&self, value: f64) {
        let idx = self.bounds.partition_point(|&bound| bound < value);
        self.counts[idx].fetch_add(1, Ordering::Relaxed);
        self.count.fetch_add(1, Ordering::Relaxed);
        self.add_to_sum(value);
    }
}

/// Storage backing the registry. Counters and gauges are plain atomics (as in
/// `metrics_util`'s `AtomicStorage`); histograms use [`BucketedHistogram`].
struct LanceStorage;

impl Storage<Key> for LanceStorage {
    type Counter = Arc<AtomicU64>;
    type Gauge = Arc<AtomicU64>;
    type Histogram = Arc<BucketedHistogram>;

    fn counter(&self, _key: &Key) -> Self::Counter {
        Arc::new(AtomicU64::new(0))
    }

    fn gauge(&self, _key: &Key) -> Self::Gauge {
        // The `metrics` facade writes the f64 bit pattern into this `u64` (the
        // snapshot decodes it with `f64::from_bits`), matching `AtomicStorage`.
        // `0` decodes to `0.0`, the correct initial value.
        Arc::new(AtomicU64::new(0))
    }

    fn histogram(&self, key: &Key) -> Self::Histogram {
        Arc::new(BucketedHistogram::new(bounds_for(key.name())))
    }
}

struct LanceRecorder {
    registry: Arc<Registry<Key, LanceStorage>>,
}

impl LanceRecorder {
    fn describe(
        &self,
        key: KeyName,
        kind: MetricKind,
        unit: Option<Unit>,
        description: SharedString,
    ) {
        CATALOG.lock().unwrap().insert(
            key.as_str().to_string(),
            MetricDescription {
                kind,
                unit: unit.map(|u| u.as_canonical_label().to_string()),
                description: description.into_owned(),
            },
        );
    }
}

impl Recorder for LanceRecorder {
    fn describe_counter(&self, key: KeyName, unit: Option<Unit>, description: SharedString) {
        self.describe(key, MetricKind::Counter, unit, description);
    }

    fn describe_gauge(&self, key: KeyName, unit: Option<Unit>, description: SharedString) {
        self.describe(key, MetricKind::Gauge, unit, description);
    }

    fn describe_histogram(&self, key: KeyName, unit: Option<Unit>, description: SharedString) {
        self.describe(key, MetricKind::Histogram, unit, description);
    }

    fn register_counter(&self, key: &Key, _metadata: &Metadata<'_>) -> Counter {
        self.registry
            .get_or_create_counter(key, |c| Counter::from_arc(c.clone()))
    }

    fn register_gauge(&self, key: &Key, _metadata: &Metadata<'_>) -> Gauge {
        self.registry
            .get_or_create_gauge(key, |g| Gauge::from_arc(g.clone()))
    }

    fn register_histogram(&self, key: &Key, _metadata: &Metadata<'_>) -> Histogram {
        self.registry
            .get_or_create_histogram(key, |h| Histogram::from_arc(h.clone()))
    }
}

/// Register the recommended histogram bounds for every metric-emitting
/// subsystem. New subsystems add their `histogram_bounds()` here.
fn register_bounds() {
    let mut bounds = HISTOGRAM_BOUNDS.write().unwrap();
    for (name, values) in lance_io::object_store::metrics::histogram_bounds() {
        bounds.insert((*name).to_string(), Arc::from(*values));
    }
}

/// Describe every metric-emitting subsystem so the catalog is populated. Must
/// run after the recorder is installed. New subsystems add their
/// `describe_metrics()` here.
fn describe_all() {
    lance_io::object_store::metrics::describe_metrics();
}

enum MetricValue {
    Scalar(f64),
    Histogram {
        buckets: Vec<(String, u64)>,
        count: u64,
        sum: f64,
    },
}

struct MetricPoint {
    name: String,
    kind: &'static str,
    attributes: HashMap<String, String>,
    value: MetricValue,
}

fn labels(key: &Key) -> HashMap<String, String> {
    key.labels()
        .map(|label| (label.key().to_string(), label.value().to_string()))
        .collect()
}

fn collect_points(registry: &Registry<Key, LanceStorage>) -> Vec<MetricPoint> {
    let mut points = Vec::new();
    for (key, handle) in registry.get_counter_handles() {
        points.push(MetricPoint {
            name: key.name().to_string(),
            kind: "counter",
            attributes: labels(&key),
            // OpenTelemetry observations are float; counts stay well within the
            // f64-exact integer range (2^53), so this cast is lossless in practice.
            value: MetricValue::Scalar(handle.load(Ordering::Relaxed) as f64),
        });
    }
    for (key, handle) in registry.get_gauge_handles() {
        points.push(MetricPoint {
            name: key.name().to_string(),
            kind: "gauge",
            attributes: labels(&key),
            value: MetricValue::Scalar(f64::from_bits(handle.load(Ordering::Relaxed))),
        });
    }
    for (key, handle) in registry.get_histogram_handles() {
        points.push(MetricPoint {
            name: key.name().to_string(),
            kind: "histogram",
            attributes: labels(&key),
            value: handle.snapshot(),
        });
    }
    points
}

/// One metric data point exposed to Python. For counters and gauges only
/// `value` is set; for histograms `buckets` (cumulative `le` counts), `count`,
/// and `sum` are set.
#[pyclass(name = "MetricPoint", get_all)]
pub struct PyMetricPoint {
    name: String,
    kind: String,
    attributes: HashMap<String, String>,
    value: Option<f64>,
    buckets: Option<Vec<(String, u64)>>,
    count: Option<u64>,
    sum: Option<f64>,
}

impl From<MetricPoint> for PyMetricPoint {
    fn from(point: MetricPoint) -> Self {
        let (value, buckets, count, sum) = match point.value {
            MetricValue::Scalar(v) => (Some(v), None, None, None),
            MetricValue::Histogram {
                buckets,
                count,
                sum,
            } => (None, Some(buckets), Some(count), Some(sum)),
        };
        Self {
            name: point.name,
            kind: point.kind.to_string(),
            attributes: point.attributes,
            value,
            buckets,
            count,
            sum,
        }
    }
}

/// A described metric, used by the Python layer to create instruments up front.
#[pyclass(name = "MetricDescription", get_all)]
pub struct PyMetricDescription {
    name: String,
    kind: String,
    unit: Option<String>,
    description: String,
}

/// Install the Lance metrics recorder as the process-global `metrics` recorder.
///
/// Returns `True` if the recorder is installed (now or previously). Returns
/// `False` if a *different* recorder is already installed — `metrics` allows
/// only one global recorder per process, so Lance cannot coexist with another.
#[pyfunction]
pub fn register_lance_metrics_recorder() -> bool {
    if REGISTRY.get().is_some() {
        return true;
    }
    let registry = Arc::new(Registry::new(LanceStorage));
    let recorder = LanceRecorder {
        registry: registry.clone(),
    };
    // Register histogram bounds *before* installing the recorder. Once the
    // recorder is global, another thread can emit a metric and create the
    // histogram handle concurrently; if the bounds aren't registered yet that
    // handle would be built with the fallback bounds and keep them for the
    // life of the process. `register_bounds()` doesn't need the recorder.
    register_bounds();
    match metrics::set_global_recorder(recorder) {
        Ok(()) => {
            let _ = REGISTRY.set(registry);
            describe_all();
            true
        }
        Err(_) => false,
    }
}

/// The catalog of described Lance metrics. Empty until the recorder is installed.
#[pyfunction]
pub fn lance_metrics_catalog() -> Vec<PyMetricDescription> {
    CATALOG
        .lock()
        .unwrap()
        .iter()
        .map(|(name, desc)| PyMetricDescription {
            name: name.clone(),
            kind: desc.kind.as_str().to_string(),
            unit: desc.unit.clone(),
            description: desc.description.clone(),
        })
        .collect()
}

/// A point-in-time snapshot of every recorded metric. Empty until the recorder
/// is installed. The read is lock-free but walks every series and allocates, so
/// it runs with the GIL released.
#[pyfunction]
pub fn snapshot_lance_metrics(py: Python<'_>) -> Vec<PyMetricPoint> {
    let Some(registry) = REGISTRY.get() else {
        return Vec::new();
    };
    let points = py.detach(|| collect_points(registry));
    points.into_iter().map(PyMetricPoint::from).collect()
}

#[cfg(test)]
mod tests {
    use super::*;
    use metrics::HistogramFn;

    fn bucket_count(buckets: &[(String, u64)], le: &str) -> u64 {
        buckets
            .iter()
            .find(|(b, _)| b == le)
            .map(|(_, c)| *c)
            .unwrap_or_else(|| panic!("no bucket with le={le}"))
    }

    #[test]
    fn bucketed_histogram_records_cumulative_buckets() {
        let hist = BucketedHistogram::new(Arc::from([0.1f64, 1.0, 10.0].as_slice()));
        hist.record(0.05); // le=0.1
        hist.record(0.5); // le=1
        hist.record(0.5); // le=1
        hist.record(50.0); // +Inf

        let MetricValue::Histogram {
            buckets,
            count,
            sum,
        } = hist.snapshot()
        else {
            panic!("expected histogram");
        };

        // Buckets are cumulative (Prometheus `le` semantics).
        assert_eq!(bucket_count(&buckets, "0.1"), 1);
        assert_eq!(bucket_count(&buckets, "1"), 3);
        assert_eq!(bucket_count(&buckets, "10"), 3);
        assert_eq!(bucket_count(&buckets, "+Inf"), 4);
        assert_eq!(count, 4);
        assert!((sum - 51.05).abs() < 1e-9);
    }

    #[test]
    fn bucketed_histogram_boundary_is_inclusive() {
        let hist = BucketedHistogram::new(Arc::from([1.0f64].as_slice()));
        hist.record(1.0); // exactly the bound -> le=1, not +Inf
        let MetricValue::Histogram { buckets, .. } = hist.snapshot() else {
            panic!("expected histogram");
        };
        assert_eq!(bucket_count(&buckets, "1"), 1);
        assert_eq!(bucket_count(&buckets, "+Inf"), 1);
    }

    #[test]
    fn bucketed_histogram_boundary_is_inclusive_mid_range() {
        // A value equal to a middle bound lands in that bucket, not the next.
        let hist = BucketedHistogram::new(Arc::from([0.1f64, 1.0, 10.0].as_slice()));
        hist.record(1.0);
        let MetricValue::Histogram { buckets, .. } = hist.snapshot() else {
            panic!("expected histogram");
        };
        assert_eq!(bucket_count(&buckets, "0.1"), 0);
        assert_eq!(bucket_count(&buckets, "1"), 1);
        assert_eq!(bucket_count(&buckets, "10"), 1); // cumulative, so still 1
        assert_eq!(bucket_count(&buckets, "+Inf"), 1);
    }

    #[test]
    fn recorder_aggregates_counters_with_labels() {
        let registry = Arc::new(Registry::new(LanceStorage));
        let recorder = LanceRecorder {
            registry: registry.clone(),
        };
        metrics::with_local_recorder(&recorder, || {
            metrics::counter!("test_requests_total", "operation" => "get", "scheme" => "s3")
                .increment(2);
            metrics::counter!("test_requests_total", "operation" => "get", "scheme" => "s3")
                .increment(3);
            // A distinct label set must produce a separate point, not merge.
            metrics::counter!("test_requests_total", "operation" => "put", "scheme" => "gs")
                .increment(7);
        });

        let scalar = |attrs: &[(&str, &str)]| {
            let points = collect_points(&registry);
            let point = points
                .into_iter()
                .find(|p| {
                    p.name == "test_requests_total"
                        && attrs
                            .iter()
                            .all(|(k, v)| p.attributes.get(*k).map(String::as_str) == Some(*v))
                })
                .expect("counter recorded for label set");
            assert_eq!(point.kind, "counter");
            match point.value {
                MetricValue::Scalar(v) => v,
                _ => panic!("expected scalar"),
            }
        };

        // Same labels aggregate; distinct labels stay separate.
        assert!((scalar(&[("operation", "get"), ("scheme", "s3")]) - 5.0).abs() < 1e-9);
        assert!((scalar(&[("operation", "put"), ("scheme", "gs")]) - 7.0).abs() < 1e-9);
    }

    #[test]
    fn recorder_records_gauges() {
        let registry = Arc::new(Registry::new(LanceStorage));
        let recorder = LanceRecorder {
            registry: registry.clone(),
        };
        // Gauges store the f64 bit pattern in a u64; the snapshot must decode it.
        metrics::with_local_recorder(&recorder, || {
            metrics::gauge!("test_gauge", "scheme" => "s3").set(3.5);
        });

        let points = collect_points(&registry);
        let point = points
            .iter()
            .find(|p| p.name == "test_gauge")
            .expect("gauge recorded");
        assert_eq!(point.kind, "gauge");
        assert!(matches!(point.value, MetricValue::Scalar(v) if (v - 3.5).abs() < 1e-9));
    }

    #[test]
    fn recorder_falls_back_to_default_bounds() {
        // A histogram with no registered bounds uses DEFAULT_BOUNDS.
        let name = "test_unregistered_histogram";
        assert!(!HISTOGRAM_BOUNDS.read().unwrap().contains_key(name));

        let registry = Arc::new(Registry::new(LanceStorage));
        let recorder = LanceRecorder {
            registry: registry.clone(),
        };
        metrics::with_local_recorder(&recorder, || {
            metrics::histogram!(name).record(0.02);
        });

        let points = collect_points(&registry);
        let point = points.iter().find(|p| p.name == name).expect("recorded");
        let MetricValue::Histogram { buckets, count, .. } = &point.value else {
            panic!("expected histogram");
        };
        assert_eq!(*count, 1);
        // DEFAULT_BOUNDS yields one bucket per bound plus the implicit `+Inf`.
        assert_eq!(buckets.len(), DEFAULT_BOUNDS.len() + 1);
        // 0.02 falls in the le=0.025 bucket (the third DEFAULT_BOUNDS entry).
        assert_eq!(bucket_count(buckets, "0.025"), 1);
        assert_eq!(bucket_count(buckets, "0.01"), 0);
        assert_eq!(bucket_count(buckets, "+Inf"), 1);
    }

    #[test]
    fn recorder_uses_registered_histogram_bounds() {
        let name = "test_recorder_bounds_seconds";
        HISTOGRAM_BOUNDS
            .write()
            .unwrap()
            .insert(name.to_string(), Arc::from([0.1f64, 1.0].as_slice()));

        let registry = Arc::new(Registry::new(LanceStorage));
        let recorder = LanceRecorder {
            registry: registry.clone(),
        };
        metrics::with_local_recorder(&recorder, || {
            metrics::histogram!(name).record(0.05);
            metrics::histogram!(name).record(5.0);
        });

        let points = collect_points(&registry);
        let point = points.iter().find(|p| p.name == name).expect("recorded");
        let MetricValue::Histogram { buckets, count, .. } = &point.value else {
            panic!("expected histogram");
        };
        assert_eq!(*count, 2);
        assert_eq!(bucket_count(buckets, "0.1"), 1);
        assert_eq!(bucket_count(buckets, "+Inf"), 2);
    }

    #[test]
    fn describe_populates_catalog() {
        let name = "test_describe_catalog_total";
        let registry = Arc::new(Registry::new(LanceStorage));
        let recorder = LanceRecorder { registry };
        metrics::with_local_recorder(&recorder, || {
            metrics::describe_counter!(name, Unit::Count, "a test counter");
        });

        let catalog = CATALOG.lock().unwrap();
        let desc = catalog.get(name).expect("described");
        assert!(matches!(desc.kind, MetricKind::Counter));
        assert_eq!(desc.description, "a test counter");
    }

    #[test]
    fn describe_all_covers_object_store_metrics() {
        use lance_io::object_store::metrics as os;

        let registry = Arc::new(Registry::new(LanceStorage));
        let recorder = LanceRecorder { registry };
        metrics::with_local_recorder(&recorder, describe_all);

        let catalog = CATALOG.lock().unwrap();
        let kind = |name: &str| catalog.get(name).expect("described").kind;
        // Every emitted object store metric must be described so the OTel
        // bridge can create an instrument for it, including the gauge and the
        // retryable counter that a plain request path might never emit.
        assert!(matches!(kind(os::METRIC_REQUESTS), MetricKind::Counter));
        assert!(matches!(kind(os::METRIC_BYTES), MetricKind::Counter));
        assert!(matches!(kind(os::METRIC_ERRORS), MetricKind::Counter));
        assert!(matches!(kind(os::METRIC_THROTTLE), MetricKind::Counter));
        assert!(matches!(kind(os::METRIC_RETRYABLE), MetricKind::Counter));
        assert!(matches!(kind(os::METRIC_IN_FLIGHT), MetricKind::Gauge));
        assert!(matches!(kind(os::METRIC_DURATION), MetricKind::Histogram));
    }
}
