1use std::{
2 cmp::Reverse,
3 collections::{BTreeMap, HashMap},
4};
5
6use chrono::{DateTime, Duration, Local, Utc};
7
8use crate::models::{
9 DimensionDistributions, DistributionPoint, TimelinePoint, Totals, UNKNOWN, UsageEvent,
10 UsageSummary, UsageWarning,
11};
12
13pub fn build_summary(
14 mut events: Vec<UsageEvent>,
15 warnings: Vec<UsageWarning>,
16 session_count: usize,
17 now: DateTime<Utc>,
18) -> UsageSummary {
19 let now_date = now.with_timezone(&Local).date_naive();
20 let last_7_days_start = now_date - Duration::days(6);
21
22 let mut totals = Totals {
23 event_count: events.len(),
24 session_count,
25 ..Totals::default()
26 };
27 let mut timeline = BTreeMap::<String, (u64, usize)>::new();
28
29 for event in &events {
30 let tokens = event.usage.total_tokens;
31 let event_date = event.timestamp.with_timezone(&Local).date_naive();
32
33 totals.all_time_tokens += tokens;
34 totals.input_tokens += event.usage.input_tokens;
35 totals.cached_input_tokens += event.usage.cached_input_tokens;
36 totals.output_tokens += event.usage.output_tokens;
37 totals.reasoning_output_tokens += event.usage.reasoning_output_tokens;
38
39 if event_date == now_date {
40 totals.today_tokens += tokens;
41 }
42 if event_date >= last_7_days_start && event_date <= now_date {
43 totals.last_7_days_tokens += tokens;
44 }
45
46 let bucket = event_date.to_string();
47 let entry = timeline.entry(bucket).or_insert((0, 0));
48 entry.0 += tokens;
49 entry.1 += 1;
50 }
51
52 events.sort_by_key(|event| Reverse(event.timestamp));
53
54 UsageSummary {
55 generated_at: now,
56 totals,
57 timeline: timeline
58 .into_iter()
59 .map(|(bucket, (total_tokens, event_count))| TimelinePoint {
60 bucket,
61 total_tokens,
62 event_count,
63 })
64 .collect(),
65 distributions: DimensionDistributions {
66 projects: distribution_by(&events, |event| event.project.clone()),
67 models: distribution_by(&events, |event| event.model.clone()),
68 threads: distribution_by(&events, |event| event.thread.clone()),
69 efforts: distribution_by(&events, |event| event.effort.clone()),
70 context_windows: distribution_by(&events, |event| {
71 event
72 .context_window
73 .map(|window| window.to_string())
74 .unwrap_or_else(|| UNKNOWN.to_string())
75 }),
76 },
77 events,
78 warnings,
79 }
80}
81
82fn distribution_by(
83 events: &[UsageEvent],
84 label_for: impl Fn(&UsageEvent) -> String,
85) -> Vec<DistributionPoint> {
86 let mut totals = HashMap::<String, (u64, usize)>::new();
87
88 for event in events {
89 let label = normalize_label(label_for(event));
90 let entry = totals.entry(label).or_insert((0, 0));
91 entry.0 += event.usage.total_tokens;
92 entry.1 += 1;
93 }
94
95 let mut points = totals
96 .into_iter()
97 .map(|(label, (total_tokens, event_count))| DistributionPoint {
98 label,
99 total_tokens,
100 event_count,
101 })
102 .collect::<Vec<_>>();
103
104 points.sort_by(|left, right| {
105 right
106 .total_tokens
107 .cmp(&left.total_tokens)
108 .then_with(|| right.event_count.cmp(&left.event_count))
109 .then_with(|| left.label.cmp(&right.label))
110 });
111 points
112}
113
114fn normalize_label(label: String) -> String {
115 if label.trim().is_empty() {
116 UNKNOWN.to_string()
117 } else {
118 label
119 }
120}
121
122#[cfg(test)]
123mod tests {
124 use chrono::{DateTime, Utc};
125
126 use super::build_summary;
127 use crate::models::{TokenUsage, UsageEvent};
128
129 fn at(timestamp: &str) -> DateTime<Utc> {
130 DateTime::parse_from_rfc3339(timestamp)
131 .expect("valid timestamp")
132 .with_timezone(&Utc)
133 }
134
135 fn event(
136 timestamp: &str,
137 project: &str,
138 model: &str,
139 thread: &str,
140 effort: &str,
141 context_window: Option<u64>,
142 total_tokens: u64,
143 ) -> UsageEvent {
144 UsageEvent {
145 timestamp: at(timestamp),
146 session_id: format!("session-{thread}"),
147 thread: thread.to_string(),
148 project: project.to_string(),
149 model: model.to_string(),
150 effort: effort.to_string(),
151 context_window,
152 usage: TokenUsage {
153 input_tokens: total_tokens / 2,
154 cached_input_tokens: 1,
155 output_tokens: total_tokens / 4,
156 reasoning_output_tokens: 2,
157 total_tokens,
158 },
159 cumulative_total_tokens: Some(total_tokens * 10),
160 source_file: "fixture.jsonl".to_string(),
161 }
162 }
163
164 #[test]
165 fn aggregates_totals_time_windows_and_token_breakdown() {
166 let events = vec![
167 event(
168 "2026-07-06T08:00:00Z",
169 "alpha",
170 "gpt-5",
171 "one",
172 "xhigh",
173 Some(200),
174 10,
175 ),
176 event(
177 "2026-07-05T08:00:00Z",
178 "alpha",
179 "gpt-5",
180 "two",
181 "high",
182 Some(200),
183 20,
184 ),
185 event(
186 "2026-06-28T08:00:00Z",
187 "beta",
188 "gpt-4.1",
189 "three",
190 "medium",
191 None,
192 30,
193 ),
194 ];
195
196 let summary = build_summary(events, Vec::new(), 3, at("2026-07-06T12:00:00Z"));
197
198 assert_eq!(summary.totals.all_time_tokens, 60);
199 assert_eq!(summary.totals.today_tokens, 10);
200 assert_eq!(summary.totals.last_7_days_tokens, 30);
201 assert_eq!(summary.totals.event_count, 3);
202 assert_eq!(summary.totals.session_count, 3);
203 assert_eq!(summary.totals.input_tokens, 30);
204 assert_eq!(summary.totals.cached_input_tokens, 3);
205 assert_eq!(summary.totals.output_tokens, 14);
206 assert_eq!(summary.totals.reasoning_output_tokens, 6);
207 }
208
209 #[test]
210 fn builds_sorted_distributions_and_timeline() {
211 let events = vec![
212 event(
213 "2026-07-06T08:00:00Z",
214 "alpha",
215 "gpt-5",
216 "one",
217 "xhigh",
218 Some(200),
219 10,
220 ),
221 event(
222 "2026-07-05T08:00:00Z",
223 "alpha",
224 "gpt-5",
225 "two",
226 "high",
227 Some(200),
228 20,
229 ),
230 event(
231 "2026-07-05T09:00:00Z",
232 "beta",
233 "gpt-4.1",
234 "one",
235 "xhigh",
236 None,
237 30,
238 ),
239 ];
240
241 let summary = build_summary(events, Vec::new(), 2, at("2026-07-06T12:00:00Z"));
242
243 assert_eq!(summary.distributions.projects[0].label, "alpha");
244 assert_eq!(summary.distributions.projects[0].total_tokens, 30);
245 assert_eq!(summary.distributions.projects[1].label, "beta");
246 assert_eq!(summary.distributions.projects[1].total_tokens, 30);
247 assert_eq!(summary.distributions.models[0].label, "gpt-5");
248 assert_eq!(summary.distributions.models[0].total_tokens, 30);
249 assert_eq!(summary.distributions.threads[0].label, "one");
250 assert_eq!(summary.distributions.threads[0].total_tokens, 40);
251 assert_eq!(summary.distributions.efforts[0].label, "xhigh");
252 assert_eq!(summary.distributions.efforts[0].total_tokens, 40);
253 assert_eq!(summary.distributions.context_windows[0].label, "200");
254 assert_eq!(summary.distributions.context_windows[0].total_tokens, 30);
255 assert_eq!(summary.distributions.context_windows[1].label, "Unknown");
256 assert_eq!(summary.distributions.context_windows[1].total_tokens, 30);
257
258 assert_eq!(summary.timeline.len(), 2);
259 assert_eq!(summary.timeline[0].bucket, "2026-07-05");
260 assert_eq!(summary.timeline[0].total_tokens, 50);
261 assert_eq!(summary.timeline[1].bucket, "2026-07-06");
262 assert_eq!(summary.timeline[1].total_tokens, 10);
263
264 assert_eq!(summary.events[0].timestamp, at("2026-07-06T08:00:00Z"));
265 }
266}