//! What a classification on the ACP runtime cost, exported as the counter
//! `hyperhive.review_classifier.cost.usage` with the run's `repo`, `pr`,
//! `model` and `currency`, and the tokens it used, as the counter
//! `hyperhive.review_classifier.token.usage` with the same labels bar
//! `currency` and a `type` of `input`, `output`, `cacheRead` or
//! `cacheCreation`. Those are `claude_code.token.usage`'s types, which
//! `claude` exports for itself on the claude runtime, so the cache hit rate is
//! `cacheRead` over the three input types on either runtime.
//!
//! The amount is computed here from the prices the run carries and the
//! tokens the agent reports, so it does not depend on how, or whether, the
//! agent prices a custom provider.
//!
//! The exporter is built on the first record, from the standard `OTEL_*`
//! environment, as `hive-agent`'s `otel_turn_metrics` does; with no
//! `OTEL_EXPORTER_OTLP_ENDPOINT` a record does nothing.

use std::sync::OnceLock;
use std::time::Duration;

use hive_runtime::TokenUsage;
use opentelemetry::KeyValue;
use opentelemetry::metrics::{Counter, MeterProvider as _};
use opentelemetry_otlp::{MetricExporter, Protocol, WithExportConfig};
use opentelemetry_sdk::metrics::{PeriodicReader, SdkMeterProvider};
use swarm_wire::review_classifier::Prices;

/// The currency every price is in.
pub const CURRENCY: &str = "USD";

/// How often the counter is exported.
const INTERVAL: Duration = Duration::from_mins(1);

/// One classification's cost and tokens, and what they are labelled with.
#[derive(Debug, Clone, PartialEq)]
pub struct Cost {
    pub repo: String,
    pub pr: u64,
    pub model: String,
    /// In [`CURRENCY`].
    pub amount: f64,
    pub usage: TokenUsage,
}

/// Where a classification's cost goes: the metrics store in production, a
/// list in the tests.
pub trait CostSink: Send + Sync {
    fn record(&self, cost: &Cost);
}

/// `usage` at `prices`. Cache reads and writes without a price of their own
/// are priced as input.
#[must_use]
pub fn amount(prices: &Prices, usage: &TokenUsage) -> f64 {
    let per_token = |price: f64| price / 1_000_000.0;
    // One model call's count is far below `u32::MAX`.
    let tokens = |n: u64| f64::from(u32::try_from(n).unwrap_or(u32::MAX));
    tokens(usage.input_tokens) * per_token(prices.input)
        + tokens(usage.output_tokens) * per_token(prices.output)
        + tokens(usage.cache_read_input_tokens)
            * per_token(prices.cache_read.unwrap_or(prices.input))
        + tokens(usage.cache_creation_input_tokens)
            * per_token(prices.cache_write.unwrap_or(prices.input))
}

/// The counters in the metrics store.
pub struct Otel;

impl CostSink for Otel {
    fn record(&self, cost: &Cost) {
        static COUNTERS: OnceLock<Option<(SdkMeterProvider, Counters)>> = OnceLock::new();
        let Some((_, counters)) = COUNTERS.get_or_init(build) else {
            return;
        };
        add(counters, cost);
    }
}

/// The cost counter, and the token counter beside it.
struct Counters {
    cost: Counter<f64>,
    tokens: Counter<u64>,
}

/// The counters on `provider`.
fn counters(provider: &SdkMeterProvider) -> Counters {
    let meter = provider.meter("hyperhive.review_classifier");
    Counters {
        cost: meter
            .f64_counter("hyperhive.review_classifier.cost.usage")
            .build(),
        tokens: meter
            .u64_counter("hyperhive.review_classifier.token.usage")
            .build(),
    }
}

/// Adds `cost` to `counters` with the run's labels.
fn add(counters: &Counters, cost: &Cost) {
    let labels = [
        KeyValue::new("repo", cost.repo.clone()),
        KeyValue::new("pr", i64::try_from(cost.pr).unwrap_or(i64::MAX)),
        KeyValue::new("model", cost.model.clone()),
    ];
    let with = |extra: KeyValue| {
        let mut all = labels.to_vec();
        all.push(extra);
        all
    };
    counters
        .cost
        .add(cost.amount, &with(KeyValue::new("currency", CURRENCY)));
    let usage = &cost.usage;
    for (kind, n) in [
        ("input", usage.input_tokens),
        ("output", usage.output_tokens),
        ("cacheRead", usage.cache_read_input_tokens),
        ("cacheCreation", usage.cache_creation_input_tokens),
    ] {
        counters.tokens.add(n, &with(KeyValue::new("type", kind)));
    }
}

/// The provider, kept alive with the counters since it exports only while it
/// lives, or `None` when no collector is configured or the exporter failed.
fn build() -> Option<(SdkMeterProvider, Counters)> {
    if !std::env::var("OTEL_EXPORTER_OTLP_ENDPOINT").is_ok_and(|s| !s.trim().is_empty()) {
        return None;
    }
    // The blocking client: `PeriodicReader` exports from a thread of its own
    // with no Tokio reactor. The resource comes from `OTEL_RESOURCE_ATTRIBUTES`.
    let exporter = MetricExporter::builder()
        .with_http()
        .with_protocol(Protocol::HttpJson)
        .build();
    let exporter = match exporter {
        Ok(exporter) => exporter,
        Err(e) => {
            tracing::warn!(error = %e, "building the cost metric exporter failed; no cost is exported");
            return None;
        }
    };
    let reader = PeriodicReader::builder(exporter)
        .with_interval(INTERVAL)
        .build();
    let provider = SdkMeterProvider::builder().with_reader(reader).build();
    let counters = counters(&provider);
    Some((provider, counters))
}

#[cfg(test)]
mod tests {
    use super::*;
    use opentelemetry::Value;
    use opentelemetry_sdk::metrics::InMemoryMetricExporter;
    use opentelemetry_sdk::metrics::data::{
        AggregatedMetrics, MetricData, ResourceMetrics, ScopeMetrics,
    };

    fn usage(input: u64, output: u64, cache_read: u64, cache_write: u64) -> TokenUsage {
        let mut usage = TokenUsage::default();
        usage.input_tokens = input;
        usage.output_tokens = output;
        usage.cache_read_input_tokens = cache_read;
        usage.cache_creation_input_tokens = cache_write;
        usage
    }

    #[test]
    fn each_token_class_is_charged_at_its_price_per_million() {
        let prices = Prices {
            input: 3.0,
            output: 15.0,
            cache_read: Some(0.3),
            cache_write: Some(3.75),
        };
        let cost = amount(&prices, &usage(1_000, 200, 10_000, 2_000));
        assert!(
            (cost - (0.003 + 0.003 + 0.003 + 0.0075)).abs() < 1e-12,
            "{cost}"
        );
    }

    #[test]
    fn cache_without_a_price_of_its_own_is_charged_as_input() {
        let prices = Prices {
            input: 2.0,
            output: 0.0,
            cache_read: None,
            cache_write: None,
        };
        let cost = amount(&prices, &usage(0, 0, 500_000, 500_000));
        assert!((cost - 2.0).abs() < 1e-12, "{cost}");
    }

    #[test]
    fn a_record_adds_the_amount_under_the_counter_name_and_the_run_labels() {
        let exporter = InMemoryMetricExporter::default();
        let reader = PeriodicReader::builder(exporter.clone()).build();
        let provider = SdkMeterProvider::builder().with_reader(reader).build();
        let counters = counters(&provider);
        let prices = Prices {
            input: 3.0,
            output: 15.0,
            cache_read: None,
            cache_write: None,
        };
        let cost = Cost {
            repo: "hyperhive/hyperhive".to_owned(),
            pr: 5214,
            model: "some-model".to_owned(),
            amount: amount(&prices, &usage(1_000_000, 100_000, 0, 0)),
            usage: usage(1_000_000, 100_000, 0, 0),
        };
        add(&counters, &cost);
        provider.force_flush().expect("flush");

        let metrics = exporter.get_finished_metrics().expect("metrics");
        let metric = metrics
            .iter()
            .flat_map(ResourceMetrics::scope_metrics)
            .flat_map(ScopeMetrics::metrics)
            .find(|m| m.name() == "hyperhive.review_classifier.cost.usage")
            .expect("the cost counter was exported");
        let AggregatedMetrics::F64(MetricData::Sum(sum)) = metric.data() else {
            panic!("not an f64 sum: {:?}", metric.data());
        };
        let points: Vec<_> = sum.data_points().collect();
        assert_eq!(points.len(), 1);
        let point = points[0];
        assert!((point.value() - 4.5).abs() < 1e-12, "{}", point.value());
        let labels: Vec<(String, Value)> = point
            .attributes()
            .map(|kv| (kv.key.to_string(), kv.value.clone()))
            .collect();
        let label = |key: &str| {
            labels
                .iter()
                .find(|(k, _)| k == key)
                .map(|(_, v)| v.clone())
        };
        assert_eq!(label("repo"), Some(Value::from("hyperhive/hyperhive")));
        assert_eq!(label("pr"), Some(Value::from(5214_i64)));
        assert_eq!(label("model"), Some(Value::from("some-model")));
        assert_eq!(label("currency"), Some(Value::from("USD")));
        assert_eq!(labels.len(), 4);
    }

    #[test]
    fn a_record_adds_each_token_class_under_its_type() {
        let exporter = InMemoryMetricExporter::default();
        let reader = PeriodicReader::builder(exporter.clone()).build();
        let provider = SdkMeterProvider::builder().with_reader(reader).build();
        let cost = Cost {
            repo: "hyperhive/hyperhive".to_owned(),
            pr: 5245,
            model: "some-model".to_owned(),
            amount: 0.0,
            usage: usage(100, 20, 9_000, 400),
        };
        add(&counters(&provider), &cost);
        provider.force_flush().expect("flush");

        let metrics = exporter.get_finished_metrics().expect("metrics");
        let metric = metrics
            .iter()
            .flat_map(ResourceMetrics::scope_metrics)
            .flat_map(ScopeMetrics::metrics)
            .find(|m| m.name() == "hyperhive.review_classifier.token.usage")
            .expect("the token counter was exported");
        let AggregatedMetrics::U64(MetricData::Sum(sum)) = metric.data() else {
            panic!("not a u64 sum: {:?}", metric.data());
        };
        let mut by_type: Vec<(String, u64, usize)> = sum
            .data_points()
            .map(|point| {
                let kind = point
                    .attributes()
                    .find(|kv| kv.key.as_str() == "type")
                    .map(|kv| kv.value.to_string())
                    .unwrap_or_default();
                (kind, point.value(), point.attributes().count())
            })
            .collect();
        by_type.sort();
        assert_eq!(
            by_type,
            [
                ("cacheCreation".to_owned(), 400, 4),
                ("cacheRead".to_owned(), 9_000, 4),
                ("input".to_owned(), 100, 4),
                ("output".to_owned(), 20, 4),
            ]
        );
    }
}