1use std::fs;
9use std::path::{Path, PathBuf};
10
11use crate::cache::atomic_write;
12use crate::error::{AppError, Result};
13use crate::vendor::VendorId;
14
15fn state_dir() -> Result<PathBuf> {
16 let base = directories::BaseDirs::new()
17 .ok_or_else(|| AppError::Other("could not resolve XDG cache dir".into()))?;
18 Ok(base.cache_dir().join("ai-usagebar"))
19}
20
21fn state_path() -> Result<PathBuf> {
22 Ok(state_dir()?.join("active_vendor"))
23}
24
25pub fn read() -> Option<VendorId> {
28 read_from(&state_path().ok()?)
29}
30
31pub fn read_from(path: &Path) -> Option<VendorId> {
35 let raw = fs::read_to_string(path).ok()?;
36 parse_slug(raw.trim())
37}
38
39pub fn write(vendor: VendorId) -> Result<()> {
41 write_to(&state_path()?, vendor)
42}
43
44pub fn write_to(path: &Path, vendor: VendorId) -> Result<()> {
47 atomic_write(path, vendor.slug().as_bytes())
48}
49
50pub fn cycle(enabled: &[VendorId], start: VendorId, delta: i32) -> Result<VendorId> {
54 cycle_at(&state_path()?, enabled, start, delta)
55}
56
57pub fn cycle_at(
61 path: &Path,
62 enabled: &[VendorId],
63 start: VendorId,
64 delta: i32,
65) -> Result<VendorId> {
66 if enabled.is_empty() {
67 return Err(AppError::Other("no enabled vendors to cycle".into()));
68 }
69 let current = read_from(path)
70 .filter(|v| enabled.contains(v))
71 .unwrap_or(start);
72 let cur_idx = enabled.iter().position(|v| *v == current).unwrap_or(0);
73 let n = enabled.len() as i32;
74 let next_idx = ((cur_idx as i32 + delta).rem_euclid(n)) as usize;
75 let next = enabled[next_idx];
76 write_to(path, next)?;
77 Ok(next)
78}
79
80fn parse_slug(s: &str) -> Option<VendorId> {
81 match s {
82 "anthropic" => Some(VendorId::Anthropic),
83 "openai" => Some(VendorId::Openai),
84 "zai" => Some(VendorId::Zai),
85 "openrouter" => Some(VendorId::Openrouter),
86 "deepseek" => Some(VendorId::Deepseek),
87 "kimi" => Some(VendorId::Kimi),
88 _ => None,
89 }
90}
91
92#[cfg(test)]
93mod tests {
94 use super::*;
95 use tempfile::TempDir;
96
97 const CYCLE_SET: [VendorId; 4] = [
100 VendorId::Anthropic,
101 VendorId::Openai,
102 VendorId::Zai,
103 VendorId::Openrouter,
104 ];
105
106 #[test]
107 fn parse_slug_round_trip() {
108 for id in VendorId::all() {
109 assert_eq!(parse_slug(id.slug()), Some(*id));
110 }
111 }
112
113 #[test]
114 fn parse_slug_unknown_returns_none() {
115 assert!(parse_slug("not-a-vendor").is_none());
116 assert!(parse_slug("").is_none());
117 }
118
119 #[test]
120 fn read_from_missing_or_garbage_returns_none() {
121 let td = TempDir::new().unwrap();
122 assert!(read_from(&td.path().join("active_vendor")).is_none());
124 let path = td.path().join("active_vendor");
126 write_to(&path, VendorId::Zai).unwrap();
127 assert_eq!(read_from(&path), Some(VendorId::Zai));
128 fs::write(&path, "not-a-vendor").unwrap();
130 assert!(read_from(&path).is_none());
131 }
132
133 #[test]
134 fn cycle_at_persists_state_across_calls() {
135 let td = TempDir::new().unwrap();
136 let path = td.path().join("active_vendor");
137
138 let v = cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap();
140 assert_eq!(v, VendorId::Openai);
141 assert_eq!(read_from(&path), Some(VendorId::Openai));
142
143 let v = cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap();
145 assert_eq!(v, VendorId::Zai);
146 assert_eq!(read_from(&path), Some(VendorId::Zai));
147 }
148
149 #[test]
150 fn cycle_at_wraps_forward_and_backward() {
151 let td = TempDir::new().unwrap();
152 let path = td.path().join("active_vendor");
153 write_to(&path, VendorId::Anthropic).unwrap();
154
155 assert_eq!(
157 cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, -1).unwrap(),
158 VendorId::Openrouter
159 );
160 assert_eq!(
162 cycle_at(&path, &CYCLE_SET, VendorId::Anthropic, 1).unwrap(),
163 VendorId::Anthropic
164 );
165 }
166
167 #[test]
168 fn cycle_at_ignores_persisted_vendor_not_in_enabled_set() {
169 let td = TempDir::new().unwrap();
170 let path = td.path().join("active_vendor");
171 write_to(&path, VendorId::Deepseek).unwrap();
173 let enabled = [VendorId::Anthropic, VendorId::Openai];
174 let v = cycle_at(&path, &enabled, VendorId::Openai, 1).unwrap();
176 assert_eq!(v, VendorId::Anthropic);
177 }
178
179 #[test]
180 fn cycle_at_errors_on_empty_enabled() {
181 let td = TempDir::new().unwrap();
182 let path = td.path().join("active_vendor");
183 let res = cycle_at(&path, &[], VendorId::Anthropic, 1);
184 assert!(matches!(res, Err(AppError::Other(_))));
185 }
186}