1use crate::error::{DatalabError, Result};
2use chrono::{DateTime, Utc};
3use serde::{Deserialize, Serialize};
4use sha2::{Digest, Sha256};
5use std::fs;
6use std::path::PathBuf;
7
8#[derive(Serialize, Deserialize)]
9pub struct CacheMetadata {
10 pub created_at: DateTime<Utc>,
11 pub endpoint: String,
12 pub params_hash: String,
13 pub file_hash: Option<String>,
14 pub file_path: Option<String>,
15}
16
17pub struct Cache {
18 base_dir: PathBuf,
19}
20
21impl Cache {
22 pub fn new() -> Result<Self> {
23 let base_dir = dirs::cache_dir()
24 .ok_or_else(|| {
25 DatalabError::CacheError(std::io::Error::new(
26 std::io::ErrorKind::NotFound,
27 "Could not find cache directory",
28 ))
29 })?
30 .join("datalab");
31
32 fs::create_dir_all(base_dir.join("responses"))?;
33 fs::create_dir_all(base_dir.join("files"))?;
34
35 Ok(Self { base_dir })
36 }
37
38 pub fn generate_key(
39 file_hash: Option<&str>,
40 file_url: Option<&str>,
41 endpoint: &str,
42 params: &serde_json::Value,
43 ) -> String {
44 let mut hasher = Sha256::new();
45
46 if let Some(hash) = file_hash {
47 hasher.update(hash.as_bytes());
48 }
49 if let Some(url) = file_url {
50 hasher.update(url.as_bytes());
51 }
52 hasher.update(endpoint.as_bytes());
53
54 let params_str = serde_json::to_string(params).unwrap_or_default();
55 hasher.update(params_str.as_bytes());
56
57 hex::encode(hasher.finalize())
58 }
59
60 pub fn hash_file(path: &PathBuf) -> Result<String> {
61 let content = fs::read(path)?;
62 let mut hasher = Sha256::new();
63 hasher.update(&content);
64 Ok(hex::encode(hasher.finalize()))
65 }
66
67 pub fn get(&self, cache_key: &str) -> Option<serde_json::Value> {
68 let response_path = self
69 .base_dir
70 .join("responses")
71 .join(format!("{}.json", cache_key));
72
73 if response_path.exists() {
74 if let Ok(content) = fs::read_to_string(&response_path) {
75 if let Ok(value) = serde_json::from_str(&content) {
76 return Some(value);
77 }
78 }
79 }
80 None
81 }
82
83 pub fn set(
84 &self,
85 cache_key: &str,
86 response: &serde_json::Value,
87 endpoint: &str,
88 file_hash: Option<&str>,
89 file_path: Option<&str>,
90 ) -> Result<()> {
91 let response_path = self
92 .base_dir
93 .join("responses")
94 .join(format!("{}.json", cache_key));
95 let meta_path = self
96 .base_dir
97 .join("responses")
98 .join(format!("{}.meta.json", cache_key));
99
100 fs::write(&response_path, serde_json::to_string_pretty(response)?)?;
101
102 let metadata = CacheMetadata {
103 created_at: Utc::now(),
104 endpoint: endpoint.to_string(),
105 params_hash: cache_key.to_string(),
106 file_hash: file_hash.map(String::from),
107 file_path: file_path.map(String::from),
108 };
109
110 fs::write(&meta_path, serde_json::to_string_pretty(&metadata)?)?;
111
112 Ok(())
113 }
114
115 #[allow(dead_code)]
116 pub fn save_binary(&self, file_hash: &str, data: &[u8]) -> Result<PathBuf> {
117 let path = self
118 .base_dir
119 .join("files")
120 .join(format!("{}.bin", file_hash));
121 fs::write(&path, data)?;
122 Ok(path)
123 }
124
125 pub fn clear(&self, older_than_days: Option<u64>) -> Result<ClearStats> {
126 let mut stats = ClearStats::default();
127
128 let responses_dir = self.base_dir.join("responses");
129 let files_dir = self.base_dir.join("files");
130
131 let cutoff = older_than_days.map(|days| Utc::now() - chrono::Duration::days(days as i64));
132
133 if responses_dir.exists() {
134 for entry in fs::read_dir(&responses_dir)? {
135 let entry = entry?;
136 let path = entry.path();
137
138 if path.extension().map(|e| e == "json").unwrap_or(false) {
139 let should_delete = if let Some(cutoff) = cutoff {
140 if path.to_string_lossy().ends_with(".meta.json") {
141 if let Ok(content) = fs::read_to_string(&path) {
142 if let Ok(meta) = serde_json::from_str::<CacheMetadata>(&content) {
143 meta.created_at < cutoff
144 } else {
145 false
146 }
147 } else {
148 false
149 }
150 } else {
151 continue;
152 }
153 } else {
154 true
155 };
156
157 if should_delete {
158 let base_name = path.file_stem().unwrap().to_string_lossy();
159 let base_name = base_name.trim_end_matches(".meta");
160 let response_file = responses_dir.join(format!("{}.json", base_name));
161 let meta_file = responses_dir.join(format!("{}.meta.json", base_name));
162
163 if response_file.exists() {
164 fs::remove_file(&response_file)?;
165 stats.responses_cleared += 1;
166 }
167 if meta_file.exists() {
168 fs::remove_file(&meta_file)?;
169 }
170 }
171 }
172 }
173 }
174
175 if cutoff.is_none() && files_dir.exists() {
176 for entry in fs::read_dir(&files_dir)? {
177 let entry = entry?;
178 fs::remove_file(entry.path())?;
179 stats.files_cleared += 1;
180 }
181 }
182
183 Ok(stats)
184 }
185
186 pub fn stats(&self) -> Result<CacheStats> {
187 let mut stats = CacheStats::default();
188
189 let responses_dir = self.base_dir.join("responses");
190 let files_dir = self.base_dir.join("files");
191
192 if responses_dir.exists() {
193 for entry in fs::read_dir(&responses_dir)? {
194 let entry = entry?;
195 let path = entry.path();
196 if path.extension().map(|e| e == "json").unwrap_or(false)
197 && !path.to_string_lossy().ends_with(".meta.json")
198 {
199 stats.response_count += 1;
200 stats.response_size += entry.metadata()?.len();
201 }
202 }
203 }
204
205 if files_dir.exists() {
206 for entry in fs::read_dir(&files_dir)? {
207 let entry = entry?;
208 stats.file_count += 1;
209 stats.file_size += entry.metadata()?.len();
210 }
211 }
212
213 stats.cache_dir = self.base_dir.to_string_lossy().to_string();
214
215 Ok(stats)
216 }
217}
218
219#[derive(Default, Serialize)]
220pub struct ClearStats {
221 pub responses_cleared: usize,
222 pub files_cleared: usize,
223}
224
225#[derive(Default, Serialize)]
226pub struct CacheStats {
227 pub cache_dir: String,
228 pub response_count: usize,
229 pub response_size: u64,
230 pub file_count: usize,
231 pub file_size: u64,
232}