1use std::collections::BTreeMap;
2
3pub struct UsageData {
4 pub primary: Window,
5 pub secondary: Window,
6}
7
8#[derive(Clone)]
9pub struct Window {
10 pub limit_window_seconds: i64,
11 pub used_percent: f64,
12 pub reset_at: i64,
13}
14
15pub struct RenderValues {
16 pub primary_label: String,
17 pub secondary_label: String,
18 pub primary_remaining: i64,
19 pub secondary_remaining: i64,
20 pub primary_reset_epoch: i64,
21 pub secondary_reset_epoch: i64,
22}
23
24pub struct WeeklyValues {
25 pub weekly_remaining: i64,
26 pub weekly_reset_epoch: i64,
27 pub non_weekly_label: String,
28 pub non_weekly_remaining: i64,
29 pub non_weekly_reset_epoch: Option<i64>,
30}
31
32pub fn parse_usage(body: &str) -> Option<UsageData> {
33 parse_usage_body(body)
34}
35
36pub fn parse_usage_body(body: &str) -> Option<UsageData> {
37 let value: serde_json::Value = serde_json::from_str(body).ok()?;
38 parse_wham_usage(&value).or_else(|| parse_code_assist_usage(&value))
39}
40
41fn parse_wham_usage(value: &serde_json::Value) -> Option<UsageData> {
42 let rate_limit = value.get("rate_limit")?;
43 let primary_raw = rate_limit.get("primary_window")?;
44 let secondary_raw = rate_limit.get("secondary_window")?;
45
46 let primary = Window {
47 limit_window_seconds: primary_raw.get("limit_window_seconds")?.as_i64()?,
48 used_percent: primary_raw
49 .get("used_percent")
50 .and_then(|value| value.as_f64())
51 .unwrap_or(0.0),
52 reset_at: primary_raw.get("reset_at")?.as_i64()?,
53 };
54 let secondary = Window {
55 limit_window_seconds: secondary_raw.get("limit_window_seconds")?.as_i64()?,
56 used_percent: secondary_raw
57 .get("used_percent")
58 .and_then(|value| value.as_f64())
59 .unwrap_or(0.0),
60 reset_at: secondary_raw.get("reset_at")?.as_i64()?,
61 };
62
63 Some(UsageData { primary, secondary })
64}
65
66fn parse_code_assist_usage(value: &serde_json::Value) -> Option<UsageData> {
67 let buckets = value.get("buckets")?.as_array()?;
68 let now_epoch = now_epoch_seconds();
69 let mut grouped: BTreeMap<i64, f64> = BTreeMap::new();
70
71 for bucket in buckets {
72 if let Some(token_type) = bucket.get("tokenType").and_then(|value| value.as_str())
73 && !token_type.eq_ignore_ascii_case("REQUESTS")
74 {
75 continue;
76 }
77
78 let remaining_fraction = match bucket
79 .get("remainingFraction")
80 .and_then(|value| value.as_f64())
81 {
82 Some(value) => value.clamp(0.0, 1.0),
83 None => continue,
84 };
85 let used_percent = (100.0 - (remaining_fraction * 100.0)).clamp(0.0, 100.0);
86
87 let reset_at = match bucket
88 .get("resetTime")
89 .and_then(|value| value.as_str())
90 .and_then(parse_rfc3339_epoch)
91 {
92 Some(epoch) if epoch > 0 => epoch,
93 _ => continue,
94 };
95
96 grouped
97 .entry(reset_at)
98 .and_modify(|existing_used| {
100 if used_percent > *existing_used {
101 *existing_used = used_percent;
102 }
103 })
104 .or_insert(used_percent);
105 }
106
107 let mut windows: Vec<Window> = grouped
108 .iter()
109 .map(|(reset_at, used_percent)| {
110 let limit_window_seconds = if now_epoch > 0 {
111 normalize_window_seconds(reset_at.saturating_sub(now_epoch))
112 } else {
113 1
114 };
115 Window {
116 limit_window_seconds,
117 used_percent: *used_percent,
118 reset_at: *reset_at,
119 }
120 })
121 .collect();
122 if windows.is_empty() {
123 return None;
124 }
125 windows.sort_by_key(|window| window.reset_at);
126 let primary = windows.first()?.clone();
127 let secondary = windows.last().cloned().unwrap_or_else(|| primary.clone());
128
129 Some(UsageData { primary, secondary })
130}
131
132fn normalize_window_seconds(seconds: i64) -> i64 {
133 let clamped = seconds.max(1);
134 if (16 * 3_600..=32 * 3_600).contains(&clamped) {
137 return 86_400;
138 }
139 if (5 * 86_400..=9 * 86_400).contains(&clamped) {
140 return 604_800;
141 }
142 if clamped >= 3_600 {
143 return (clamped / 3_600).max(1) * 3_600;
144 }
145 if clamped >= 60 {
146 return (clamped / 60).max(1) * 60;
147 }
148 clamped
149}
150
151fn now_epoch_seconds() -> i64 {
152 std::time::SystemTime::now()
153 .duration_since(std::time::UNIX_EPOCH)
154 .ok()
155 .and_then(|duration| i64::try_from(duration.as_secs()).ok())
156 .unwrap_or(0)
157}
158
159fn parse_rfc3339_epoch(raw: &str) -> Option<i64> {
160 let normalized = normalize_iso(raw);
161 let (datetime, offset_seconds) = if normalized.ends_with('Z') {
162 (&normalized[..normalized.len().saturating_sub(1)], 0i64)
163 } else {
164 if normalized.len() < 6 {
165 return None;
166 }
167 let tail_index = normalized.len() - 6;
168 let sign = normalized.as_bytes().get(tail_index).copied()? as char;
169 if sign != '+' && sign != '-' {
170 return None;
171 }
172 if normalized.as_bytes().get(tail_index + 3).copied()? as char != ':' {
173 return None;
174 }
175 let hours = parse_u32(&normalized[tail_index + 1..tail_index + 3])? as i64;
176 let minutes = parse_u32(&normalized[tail_index + 4..])? as i64;
177 let mut offset = hours * 3_600 + minutes * 60;
178 if sign == '-' {
179 offset = -offset;
180 }
181 (&normalized[..tail_index], offset)
182 };
183
184 if datetime.len() != 19 {
185 return None;
186 }
187 if datetime.as_bytes().get(4).copied()? as char != '-'
188 || datetime.as_bytes().get(7).copied()? as char != '-'
189 || datetime.as_bytes().get(10).copied()? as char != 'T'
190 || datetime.as_bytes().get(13).copied()? as char != ':'
191 || datetime.as_bytes().get(16).copied()? as char != ':'
192 {
193 return None;
194 }
195
196 let year = parse_i64(&datetime[0..4])?;
197 let month = parse_u32(&datetime[5..7])? as i64;
198 let day = parse_u32(&datetime[8..10])? as i64;
199 let hour = parse_u32(&datetime[11..13])? as i64;
200 let minute = parse_u32(&datetime[14..16])? as i64;
201 let second = parse_u32(&datetime[17..19])? as i64;
202
203 if !(1..=12).contains(&month)
204 || !(1..=31).contains(&day)
205 || hour > 23
206 || minute > 59
207 || second > 60
208 {
209 return None;
210 }
211
212 let days = days_from_civil(year, month, day);
213 let local_epoch = days * 86_400 + hour * 3_600 + minute * 60 + second;
214 Some(local_epoch - offset_seconds)
215}
216
217fn normalize_iso(raw: &str) -> String {
218 let mut trimmed = raw
219 .split(&['\n', '\r'][..])
220 .next()
221 .unwrap_or("")
222 .to_string();
223 if let Some(dot) = trimmed.find('.')
224 && trimmed.ends_with('Z')
225 {
226 trimmed.truncate(dot);
227 trimmed.push('Z');
228 }
229 trimmed
230}
231
232pub fn render_values(data: &UsageData) -> RenderValues {
233 let primary_label = format_window_seconds(data.primary.limit_window_seconds)
234 .unwrap_or_else(|| "Primary".to_string());
235 let secondary_label = format_window_seconds(data.secondary.limit_window_seconds)
236 .unwrap_or_else(|| "Secondary".to_string());
237
238 let primary_remaining = remaining_percent(data.primary.used_percent);
239 let secondary_remaining = remaining_percent(data.secondary.used_percent);
240
241 RenderValues {
242 primary_label,
243 secondary_label,
244 primary_remaining,
245 secondary_remaining,
246 primary_reset_epoch: data.primary.reset_at,
247 secondary_reset_epoch: data.secondary.reset_at,
248 }
249}
250
251pub fn weekly_values(values: &RenderValues) -> WeeklyValues {
252 let (
253 weekly_remaining,
254 weekly_reset_epoch,
255 non_weekly_label,
256 non_weekly_remaining,
257 non_weekly_reset_epoch,
258 ) = if values.primary_label == "Weekly" {
259 (
260 values.primary_remaining,
261 values.primary_reset_epoch,
262 values.secondary_label.clone(),
263 values.secondary_remaining,
264 Some(values.secondary_reset_epoch),
265 )
266 } else {
267 (
268 values.secondary_remaining,
269 values.secondary_reset_epoch,
270 values.primary_label.clone(),
271 values.primary_remaining,
272 Some(values.primary_reset_epoch),
273 )
274 };
275
276 WeeklyValues {
277 weekly_remaining,
278 weekly_reset_epoch,
279 non_weekly_label,
280 non_weekly_remaining,
281 non_weekly_reset_epoch,
282 }
283}
284
285pub fn format_window_seconds(raw: i64) -> Option<String> {
286 if raw <= 0 {
287 return None;
288 }
289 if raw % 604_800 == 0 {
290 let weeks = raw / 604_800;
291 if weeks == 1 {
292 return Some("Weekly".to_string());
293 }
294 return Some(format!("{weeks}w"));
295 }
296 if raw % 86_400 == 0 {
297 return Some(format!("{}d", raw / 86_400));
298 }
299 if raw % 3_600 == 0 {
300 return Some(format!("{}h", raw / 3_600));
301 }
302 if raw % 60 == 0 {
303 return Some(format!("{}m", raw / 60));
304 }
305 Some(format!("{raw}s"))
306}
307
308pub fn format_epoch_local_datetime(epoch: i64) -> Option<String> {
309 let components = epoch_components(epoch)?;
310 Some(format!(
311 "{:02}-{:02} {:02}:{:02}",
312 components.1, components.2, components.3, components.4
313 ))
314}
315
316pub fn format_epoch_local_datetime_with_offset(epoch: i64) -> Option<String> {
317 let base = format_epoch_local_datetime(epoch)?;
318 Some(format!("{base} +00:00"))
319}
320
321pub fn format_epoch_local(epoch: i64, fmt: &str) -> Option<String> {
322 let components = epoch_components(epoch)?;
323 Some(format_with_components(fmt, components))
324}
325
326pub fn format_until_epoch_compact(target_epoch: i64, now_epoch: i64) -> Option<String> {
327 if target_epoch <= 0 || now_epoch <= 0 {
328 return None;
329 }
330 let remaining = target_epoch - now_epoch;
331 if remaining <= 0 {
332 return Some(format!("{:>2}h {:>2}m", 0, 0));
333 }
334
335 if remaining >= 86_400 {
336 let days = remaining / 86_400;
337 let hours = (remaining % 86_400) / 3_600;
338 return Some(format!("{:>2}d {:>2}h", days, hours));
339 }
340
341 let hours = remaining / 3_600;
342 let minutes = (remaining % 3_600) / 60;
343 Some(format!("{:>2}h {:>2}m", hours, minutes))
344}
345
346fn remaining_percent(used_percent: f64) -> i64 {
347 let remaining = 100.0 - used_percent;
348 remaining.round() as i64
349}
350
351fn epoch_components(epoch: i64) -> Option<(i64, i64, i64, i64, i64, i64)> {
352 if epoch <= 0 {
353 return None;
354 }
355 let days = epoch.div_euclid(86_400);
356 let seconds_of_day = epoch.rem_euclid(86_400);
357 let (year, month, day) = civil_from_days(days);
358 let hour = seconds_of_day / 3_600;
359 let minute = (seconds_of_day % 3_600) / 60;
360 let second = seconds_of_day % 60;
361 Some((year, month, day, hour, minute, second))
362}
363
364fn format_with_components(fmt: &str, parts: (i64, i64, i64, i64, i64, i64)) -> String {
365 let (year, month, day, hour, minute, second) = parts;
366 let mut out = String::with_capacity(fmt.len() + 16);
367 let chars: Vec<char> = fmt.chars().collect();
368 let mut index = 0usize;
369
370 while index < chars.len() {
371 let ch = chars[index];
372 if ch != '%' {
373 out.push(ch);
374 index += 1;
375 continue;
376 }
377
378 if index + 1 >= chars.len() {
379 out.push('%');
380 index += 1;
381 continue;
382 }
383
384 let next = chars[index + 1];
385 match next {
386 'Y' => out.push_str(&format!("{year:04}")),
387 'm' => out.push_str(&format!("{month:02}")),
388 'd' => out.push_str(&format!("{day:02}")),
389 'H' => out.push_str(&format!("{hour:02}")),
390 'M' => out.push_str(&format!("{minute:02}")),
391 'S' => out.push_str(&format!("{second:02}")),
392 '%' => out.push('%'),
393 ':' => {
394 if index + 2 < chars.len() && chars[index + 2] == 'z' {
395 out.push_str("+00:00");
396 index += 1;
397 } else {
398 out.push('%');
399 out.push(':');
400 }
401 }
402 other => {
403 out.push('%');
404 out.push(other);
405 }
406 }
407 index += 2;
408 }
409
410 out
411}
412
413fn civil_from_days(days_since_epoch: i64) -> (i64, i64, i64) {
414 let z = days_since_epoch + 719_468;
415 let era = if z >= 0 {
416 z / 146_097
417 } else {
418 (z - 146_096) / 146_097
419 };
420 let day_of_era = z - era * 146_097;
421 let year_of_era =
422 (day_of_era - day_of_era / 1460 + day_of_era / 36_524 - day_of_era / 146_096) / 365;
423 let mut year = year_of_era + era * 400;
424 let day_of_year = day_of_era - (365 * year_of_era + year_of_era / 4 - year_of_era / 100);
425 let month_prime = (5 * day_of_year + 2) / 153;
426 let day = day_of_year - (153 * month_prime + 2) / 5 + 1;
427 let month = month_prime + if month_prime < 10 { 3 } else { -9 };
428 if month <= 2 {
429 year += 1;
430 }
431 (year, month, day)
432}
433
434fn days_from_civil(year: i64, month: i64, day: i64) -> i64 {
435 let adjusted_year = year - i64::from(month <= 2);
436 let era = if adjusted_year >= 0 {
437 adjusted_year / 400
438 } else {
439 (adjusted_year - 399) / 400
440 };
441 let year_of_era = adjusted_year - era * 400;
442 let month_prime = month + if month > 2 { -3 } else { 9 };
443 let day_of_year = (153 * month_prime + 2) / 5 + day - 1;
444 let day_of_era = year_of_era * 365 + year_of_era / 4 - year_of_era / 100 + day_of_year;
445 era * 146_097 + day_of_era - 719_468
446}
447
448fn parse_u32(raw: &str) -> Option<u32> {
449 raw.parse::<u32>().ok()
450}
451
452fn parse_i64(raw: &str) -> Option<i64> {
453 raw.parse::<i64>().ok()
454}
455
456#[cfg(test)]
457mod tests {
458 use super::{format_window_seconds, normalize_window_seconds};
459
460 #[test]
461 fn normalize_window_seconds_stabilizes_daily_horizon_to_one_day() {
462 assert_eq!(
463 format_window_seconds(normalize_window_seconds(21 * 3_600)).as_deref(),
464 Some("1d")
465 );
466 assert_eq!(
467 format_window_seconds(normalize_window_seconds(22 * 3_600 + 45 * 60)).as_deref(),
468 Some("1d")
469 );
470 }
471
472 #[test]
473 fn normalize_window_seconds_stabilizes_weekly_horizon_to_weekly() {
474 assert_eq!(
475 format_window_seconds(normalize_window_seconds(6 * 86_400 + 3 * 3_600)).as_deref(),
476 Some("Weekly")
477 );
478 }
479
480 #[test]
481 fn normalize_window_seconds_keeps_short_windows_hourly() {
482 assert_eq!(
483 format_window_seconds(normalize_window_seconds(5 * 3_600 + 20 * 60)).as_deref(),
484 Some("5h")
485 );
486 }
487}