use std::collections::HashMap;
use chrono::{DateTime, Utc};
use crate::{api::MetricValue, Error, Result};
use super::{utils, MetricSummary, StatsRecord, StatsSummary};
#[derive(Debug, Clone)]
pub struct StatsAggregator {
records: Vec<StatsRecord>,
}
impl StatsAggregator {
pub fn new() -> Self {
Self {
records: Vec::new(),
}
}
pub fn add_record(&mut self, record: StatsRecord) {
self.records.push(record);
}
pub fn add_records(&mut self, records: Vec<StatsRecord>) {
self.records.extend(records);
}
pub fn clear(&mut self) {
self.records.clear();
}
pub fn record_count(&self) -> usize {
self.records.len()
}
pub fn summarize(&self) -> Result<StatsSummary> {
if self.records.is_empty() {
return Err(Error::statistics("No records to summarize"));
}
let time_range = self.get_time_range();
let sources = self.get_unique_sources();
let metrics = self.aggregate_metrics()?;
Ok(StatsSummary {
count: self.records.len() as u64,
time_range,
metrics,
sources,
})
}
pub fn summarize_time_range(
&self,
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Result<StatsSummary> {
let filtered_records: Vec<_> = self
.records
.iter()
.filter(|record| record.timestamp >= start && record.timestamp <= end)
.collect();
if filtered_records.is_empty() {
return Err(Error::statistics("No records in specified time range"));
}
let mut temp_aggregator = StatsAggregator::new();
temp_aggregator.records = filtered_records.into_iter().cloned().collect();
temp_aggregator.summarize()
}
pub fn summarize_by_source(&self) -> Result<HashMap<String, StatsSummary>> {
let mut summaries = HashMap::new();
let mut records_by_source: HashMap<String, Vec<StatsRecord>> = HashMap::new();
for record in &self.records {
records_by_source
.entry(record.source.clone())
.or_default()
.push(record.clone());
}
for (source, records) in records_by_source {
let mut aggregator = StatsAggregator::new();
aggregator.add_records(records);
summaries.insert(source, aggregator.summarize()?);
}
Ok(summaries)
}
pub fn filter_by_source(&self, source: &str) -> Vec<&StatsRecord> {
self.records
.iter()
.filter(|record| record.source == source)
.collect()
}
pub fn filter_by_time_range(
&self,
start: DateTime<Utc>,
end: DateTime<Utc>,
) -> Vec<&StatsRecord> {
self.records
.iter()
.filter(|record| record.timestamp >= start && record.timestamp <= end)
.collect()
}
fn get_time_range(&self) -> (DateTime<Utc>, DateTime<Utc>) {
let mut min_time = self.records[0].timestamp;
let mut max_time = self.records[0].timestamp;
for record in &self.records {
if record.timestamp < min_time {
min_time = record.timestamp;
}
if record.timestamp > max_time {
max_time = record.timestamp;
}
}
(min_time, max_time)
}
fn get_unique_sources(&self) -> Vec<String> {
let mut sources: Vec<String> = self
.records
.iter()
.map(|record| record.source.clone())
.collect();
sources.sort();
sources.dedup();
sources
}
fn aggregate_metrics(&self) -> Result<HashMap<String, MetricSummary>> {
let mut metric_values: HashMap<String, Vec<f64>> = HashMap::new();
for record in &self.records {
for (name, value) in &record.metrics {
let numeric_value = match value {
MetricValue::Integer(i) => *i as f64,
MetricValue::Float(f) => *f,
MetricValue::Duration(d) => *d as f64,
MetricValue::Boolean(b) => {
if *b {
1.0
} else {
0.0
}
}
MetricValue::String(_) => continue, MetricValue::Histogram { sum, .. } => *sum,
};
metric_values
.entry(name.clone())
.or_default()
.push(numeric_value);
}
}
let mut summaries = HashMap::new();
for (name, mut values) in metric_values {
if values.is_empty() {
continue;
}
values.sort_by(|a, b| a.partial_cmp(b).unwrap());
let min = values[0];
let max = values[values.len() - 1];
let sum = values.iter().sum();
let count = values.len() as u64;
let avg = sum / count as f64;
let std_dev = if values.len() > 1 {
Some(utils::std_deviation(&values))
} else {
None
};
let percentiles = if values.len() >= 4 {
Some([
utils::percentile(&values, 50.0), utils::percentile(&values, 90.0), utils::percentile(&values, 95.0), utils::percentile(&values, 99.0), ])
} else {
None
};
summaries.insert(
name.clone(),
MetricSummary {
name,
min,
max,
avg,
sum,
count,
std_dev,
percentiles,
},
);
}
Ok(summaries)
}
}
impl Default for StatsAggregator {
fn default() -> Self {
Self::new()
}
}