1use std::collections::BTreeMap;
4use std::fs;
5
6use chrono::Utc;
7use serde::{Deserialize, Serialize};
8use serde_json::Value;
9use sha2::{Digest, Sha256};
10
11use crate::cache::atomic_write_json;
12use crate::error::Result;
13use crate::model::CommandDef;
14use crate::paths::usage_file;
15
16#[derive(Debug, Clone, Serialize, Deserialize, Default)]
17pub struct UsageEntry {
18 pub count: u64,
19 pub last_used: String,
20}
21
22pub type UsageBucket = BTreeMap<String, UsageEntry>;
23pub type UsageData = BTreeMap<String, UsageBucket>;
24
25pub fn load_usage() -> UsageData {
26 let path = usage_file();
27 if !path.exists() {
28 return UsageData::new();
29 }
30 match fs::read_to_string(&path) {
31 Ok(text) => serde_json::from_str(&text).unwrap_or_default(),
32 Err(_) => UsageData::new(),
33 }
34}
35
36pub fn save_usage(data: &UsageData) -> Result<()> {
37 atomic_write_json(&usage_file(), data)
38}
39
40pub fn record_usage(source_hash: &str, tool_name: &str) -> Result<()> {
41 let mut usage = load_usage();
42 let bucket = usage.entry(source_hash.to_string()).or_default();
43 let entry = bucket.entry(tool_name.to_string()).or_default();
44 entry.count += 1;
45 entry.last_used = Utc::now().to_rfc3339();
46 save_usage(&usage)
47}
48
49pub fn source_hash_for(source: &str) -> String {
50 let digest = Sha256::digest(source.as_bytes());
51 hex::encode(&digest[..8])
52}
53
54fn usage_key(c: &CommandDef) -> &str {
55 c.tool_name
56 .as_deref()
57 .or(c.graphql_field_name.as_deref())
58 .unwrap_or(c.name.as_str())
59}
60
61pub fn sort_commands(
62 mut commands: Vec<CommandDef>,
63 sort_mode: &str,
64 source_hash: &str,
65) -> Vec<CommandDef> {
66 match sort_mode {
67 "default" => commands,
68 "alpha" => {
69 commands.sort_by(|a, b| a.name.cmp(&b.name));
70 commands
71 }
72 "usage" | "recent" => {
73 let usage = load_usage();
74 let bucket = usage.get(source_hash);
75 if bucket.is_none() {
76 return commands;
77 }
78 let bucket = bucket.unwrap();
79 if sort_mode == "usage" {
80 commands.sort_by(|a, b| {
81 let ca = bucket.get(usage_key(a)).map(|e| e.count).unwrap_or(0);
82 let cb = bucket.get(usage_key(b)).map(|e| e.count).unwrap_or(0);
83 cb.cmp(&ca)
84 });
85 } else {
86 commands.sort_by(|a, b| {
87 let la = bucket
88 .get(usage_key(a))
89 .map(|e| e.last_used.as_str())
90 .unwrap_or("");
91 let lb = bucket
92 .get(usage_key(b))
93 .map(|e| e.last_used.as_str())
94 .unwrap_or("");
95 lb.cmp(la)
96 });
97 }
98 commands
99 }
100 _ => commands,
101 }
102}
103
104pub fn resolve_sort_mode(explicit: Option<&str>, source_hash: &str) -> String {
105 if let Some(s) = explicit {
106 return s.to_string();
107 }
108 let usage = load_usage();
109 if usage.get(source_hash).is_some_and(|b| !b.is_empty()) {
110 "usage".into()
111 } else {
112 "default".into()
113 }
114}
115
116pub fn usage_as_value() -> Value {
118 serde_json::to_value(load_usage()).unwrap_or(Value::Object(Default::default()))
119}
120
121#[cfg(test)]
122mod tests {
123 use super::*;
124 use crate::paths::{set_cache_dir_override, set_config_dir_override, TEST_PATHS_LOCK};
125 use tempfile::tempdir;
126
127 fn with_tmp_cache<F: FnOnce()>(f: F) {
128 let _guard = TEST_PATHS_LOCK.lock().unwrap_or_else(|e| e.into_inner());
129 let dir = tempdir().unwrap();
130 set_cache_dir_override(Some(dir.path().to_path_buf()));
131 std::fs::create_dir_all(dir.path()).unwrap();
132 f();
133 set_cache_dir_override(None);
134 set_config_dir_override(None);
135 }
136
137 #[test]
138 fn load_save_roundtrip_and_corrupt() {
139 with_tmp_cache(|| {
140 assert!(load_usage().is_empty());
141 let mut data = UsageData::new();
142 let mut bucket = UsageBucket::new();
143 bucket.insert(
144 "my-tool".into(),
145 UsageEntry {
146 count: 5,
147 last_used: "2026-01-01T00:00:00+00:00".into(),
148 },
149 );
150 data.insert("abc123".into(), bucket);
151 save_usage(&data).unwrap();
152 assert_eq!(load_usage()["abc123"]["my-tool"].count, 5);
153
154 fs::write(usage_file(), "{bad json").unwrap();
155 assert!(load_usage().is_empty());
156 });
157 }
158
159 #[test]
160 fn record_and_hash() {
161 with_tmp_cache(|| {
162 record_usage("src1", "tool-a").unwrap();
163 record_usage("src1", "tool-a").unwrap();
164 record_usage("src1", "tool-b").unwrap();
165 record_usage("src2", "tool-a").unwrap();
166 let u = load_usage();
167 assert_eq!(u["src1"]["tool-a"].count, 2);
168 assert_eq!(u["src1"]["tool-b"].count, 1);
169 assert_eq!(u["src2"]["tool-a"].count, 1);
170
171 let h1 = source_hash_for("http://example.com");
172 let h2 = source_hash_for("http://example.com");
173 assert_eq!(h1, h2);
174 assert_eq!(h1.len(), 16);
175 assert_ne!(
176 source_hash_for("http://example.com"),
177 source_hash_for("http://other.com")
178 );
179 });
180 }
181
182 fn make_commands(names: &[&str]) -> Vec<CommandDef> {
183 names
184 .iter()
185 .map(|n| {
186 let mut c = CommandDef::new(*n);
187 c.tool_name = Some((*n).into());
188 c.description = format!("desc-{n}");
189 c
190 })
191 .collect()
192 }
193
194 #[test]
195 fn sort_modes() {
196 with_tmp_cache(|| {
197 let cmds = make_commands(&["c", "a", "b"]);
198 let names: Vec<_> = sort_commands(cmds.clone(), "default", "src1")
199 .into_iter()
200 .map(|c| c.name)
201 .collect();
202 assert_eq!(names, ["c", "a", "b"]);
203
204 let names: Vec<_> = sort_commands(cmds.clone(), "alpha", "src1")
205 .into_iter()
206 .map(|c| c.name)
207 .collect();
208 assert_eq!(names, ["a", "b", "c"]);
209
210 record_usage("src1", "b").unwrap();
211 record_usage("src1", "b").unwrap();
212 record_usage("src1", "a").unwrap();
213 let names: Vec<_> = sort_commands(make_commands(&["c", "a", "b"]), "usage", "src1")
214 .into_iter()
215 .map(|c| c.name)
216 .collect();
217 assert_eq!(names[0], "b");
218 assert_eq!(names[1], "a");
219
220 assert_eq!(resolve_sort_mode(Some("alpha"), "src1"), "alpha");
221 assert_eq!(resolve_sort_mode(None, "src1"), "usage");
222 assert_eq!(resolve_sort_mode(None, "empty"), "default");
223 });
224 }
225}