1use std::path::{Path, PathBuf};
8use std::sync::Arc;
9
10use localcache::{CacheEntry, CacheStatus, ConnectionPool, ReadPool};
11use rayon::prelude::*;
12
13use crate::core::engine::{
14 CacheConfig, DbLocation, IMAGE_PAYLOAD_VERSION, NAMESPACE_IMAGE, Result, cache_options,
15 ensure_db_dir, ensure_schema, is_fresh, read_pool_size,
16};
17use crate::core::payload::ImagePayload;
18use crate::core::thumbnail::{generate_image_thumbnail, thumbnail_dest};
19use crate::types::{
20 CacheRead, DirCacheSummary, ImageCacheEntry, ImageFeatures, LookupResult, UpsertImageRequest,
21};
22
23#[derive(Debug, Clone, Default)]
28pub struct ImageCacheConfig {
29 pub cache_config: CacheConfig,
30}
31
32#[derive(Clone)]
41pub struct ImageCacheWriter {
42 write: ConnectionPool<ImagePayload>,
43 read: ReadPool<ImagePayload>,
44 config: Arc<ImageCacheConfig>,
45}
46
47impl ImageCacheWriter {
48 pub fn as_session(config: ImageCacheConfig) -> Result<Self> {
49 let options = cache_options(&config.cache_config, NAMESPACE_IMAGE, IMAGE_PAYLOAD_VERSION);
50 ensure_db_dir(&options)?;
52 let write = ConnectionPool::open(options.clone())?;
55 write.with(|e| e.purge_stale_versions())?;
57 let read = ReadPool::open(options, read_pool_size(&config.cache_config))?;
58 Ok(Self {
59 write,
60 read,
61 config: Arc::new(config),
62 })
63 }
64
65 pub fn onetime(location: DbLocation) -> Result<Self> {
66 Self::as_session(ImageCacheConfig {
67 cache_config: CacheConfig {
68 db_location: location,
69 ..CacheConfig::default()
70 },
71 })
72 }
73
74 pub fn upsert(&self, req: UpsertImageRequest) -> Result<()> {
79 let status = self.write.check_status(&req.path)?;
80 self.write_payload(&req, status)
81 }
82
83 pub fn upsert_all(&self, reqs: Vec<UpsertImageRequest>) -> Vec<(PathBuf, Result<()>)> {
97 let prepared: Vec<(UpsertImageRequest, Result<Prepared>)> = reqs
99 .into_par_iter()
100 .map(|req| {
101 let prep = self.prepare(&req);
102 (req, prep)
103 })
104 .collect();
105
106 prepared
108 .into_iter()
109 .map(|(req, prep)| {
110 let path = req.path.clone();
111 let result = prep.and_then(|p| self.commit(&req, p));
112 (path, result)
113 })
114 .collect()
115 }
116
117 pub fn delete(&self, path: &Path) -> Result<bool> {
118 Ok(self.write.remove(path)?)
119 }
120
121 pub fn delete_in_dir(&self, dir: &Path) -> Result<usize> {
126 let entries = self.read.query_run(|q| q.path_in_dir(dir, false))?;
127 let mut removed = 0;
128 for entry in entries {
129 if let Some(thumb) = &entry.payload.thumbnail_path {
130 let _ = std::fs::remove_file(thumb);
132 }
133 if self.write.remove(&entry.path)? {
134 removed += 1;
135 }
136 }
137 Ok(removed)
138 }
139
140 pub fn list_paths(&self) -> Result<Vec<String>> {
141 let keys = self.write.with(|e| e.keys(None))?;
142 Ok(keys
143 .into_iter()
144 .map(|p| p.to_string_lossy().into_owned())
145 .collect())
146 }
147
148 pub fn as_reader(&self) -> ImageCacheReader {
149 ImageCacheReader {
150 read: self.read.clone(),
151 }
152 }
153
154 pub fn lookup(&self, path: &Path) -> Result<LookupResult<ImageCacheEntry>> {
155 self.as_reader().lookup(path)
156 }
157
158 fn prepare(&self, req: &UpsertImageRequest) -> Result<Prepared> {
165 let status = self.read.check_status(&req.path)?;
166 let thumbnail = self.ensure_thumbnail(&req.path, &status)?;
167 Ok(Prepared { status, thumbnail })
168 }
169
170 fn commit(&self, req: &UpsertImageRequest, prep: Prepared) -> Result<()> {
173 let existing = if is_fresh(&prep.status) {
174 self.write.get(&req.path)?.map(|e| e.payload)
175 } else {
176 None
177 };
178
179 if is_fresh(&prep.status)
182 && req.clip_vector.is_none()
183 && let Some(p) = &existing
184 && (prep.thumbnail.is_none() || p.thumbnail_path == prep.thumbnail)
185 {
186 return Ok(());
187 }
188
189 let mut payload = existing.unwrap_or_default();
190 if let Some(t) = prep.thumbnail {
191 payload.thumbnail_path = Some(t);
192 }
193 if let Some(v) = &req.clip_vector {
194 payload.clip_vector = Some(v.clone());
195 }
196 self.write.set(&req.path, &payload)?;
197 Ok(())
198 }
199
200 fn write_payload(&self, req: &UpsertImageRequest, status: CacheStatus) -> Result<()> {
202 let thumbnail = self.ensure_thumbnail(&req.path, &status)?;
203 self.commit(req, Prepared { status, thumbnail })
204 }
205
206 fn ensure_thumbnail(&self, path: &Path, status: &CacheStatus) -> Result<Option<String>> {
211 let Some(thumb_dir) = &self.config.cache_config.thumbnail_dir else {
212 return Ok(None);
213 };
214 let dest = thumbnail_dest(thumb_dir, path)?;
215 if !dest.exists() || !is_fresh(status) {
216 generate_image_thumbnail(path, &dest)?;
217 }
218 Ok(Some(dest.to_string_lossy().into_owned()))
219 }
220}
221
222struct Prepared {
224 status: CacheStatus,
225 thumbnail: Option<String>,
226}
227
228#[derive(Clone)]
235pub struct ImageCacheReader {
236 read: ReadPool<ImagePayload>,
237}
238
239impl ImageCacheReader {
240 pub fn as_session(config: ImageCacheConfig) -> Result<Self> {
241 let options = cache_options(&config.cache_config, NAMESPACE_IMAGE, IMAGE_PAYLOAD_VERSION);
242 ensure_db_dir(&options)?;
244 ensure_schema::<ImagePayload>(&options)?;
247 let read = ReadPool::open(options, read_pool_size(&config.cache_config))?;
248 Ok(Self { read })
249 }
250
251 pub fn onetime(location: DbLocation) -> Result<Self> {
252 Self::as_session(ImageCacheConfig {
253 cache_config: CacheConfig {
254 db_location: location,
255 ..CacheConfig::default()
256 },
257 })
258 }
259
260 pub fn lookup(&self, path: &Path) -> Result<LookupResult<ImageCacheEntry>> {
261 let canonical = match path.canonicalize() {
262 Ok(p) => p,
263 Err(_) => return Ok(LookupResult::Miss),
264 };
265
266 match self.read.check_status(&canonical)? {
267 CacheStatus::Missing => Ok(LookupResult::Miss),
268 CacheStatus::Stale => Ok(LookupResult::Invalidated),
269 CacheStatus::Fresh => match self.read.get(&canonical)? {
270 None => Ok(LookupResult::Miss),
271 Some(entry) => Ok(LookupResult::Hit(to_image_entry(entry))),
272 },
273 }
274 }
275
276 pub fn lookup_all(
279 &self,
280 paths: &[&Path],
281 ) -> Vec<(PathBuf, Result<LookupResult<ImageCacheEntry>>)> {
282 paths
283 .par_iter()
284 .map(|p| (p.to_path_buf(), self.lookup(p)))
285 .collect()
286 }
287
288 pub fn check(&self, path: &Path) -> Result<bool> {
289 Ok(is_fresh(&self.read.check_status(path)?))
290 }
291
292 pub fn list_paths(&self) -> Result<Vec<String>> {
293 Ok(self
294 .read
295 .keys(None)?
296 .into_iter()
297 .map(|p| p.to_string_lossy().into_owned())
298 .collect())
299 }
300
301 pub fn all(&self) -> Result<Vec<Result<ImageCacheEntry>>> {
302 let entries = self.read.query_run(|q| q)?;
303 Ok(entries.into_iter().map(|e| Ok(to_image_entry(e))).collect())
304 }
305
306 pub fn summarize_by_dir(&self) -> Result<Vec<DirCacheSummary>> {
310 summarize_entries(self.read.list_entries()?)
311 }
312
313 pub fn all_in_dir(&self, path: &Path) -> Result<Vec<Result<ImageCacheEntry>>> {
319 let dir = dir_of(path);
320 let entries = self.read.query_run(|q| q.path_in_dir(dir, false))?;
321 Ok(entries.into_iter().map(|e| Ok(to_image_entry(e))).collect())
322 }
323
324 pub fn all_in_dir_and_sub_dirs(&self, path: &Path) -> Result<Vec<Result<ImageCacheEntry>>> {
329 let dir = dir_of(path);
330 let entries = self.read.query_run(|q| q.path_in_dir(dir, true))?;
331 Ok(entries.into_iter().map(|e| Ok(to_image_entry(e))).collect())
332 }
333}
334
335impl CacheRead for ImageCacheReader {
336 fn check(&self, path: &Path) -> Result<bool> {
337 ImageCacheReader::check(self, path)
338 }
339
340 fn check_all(&self, paths: &[&Path]) -> Vec<(PathBuf, Result<bool>)> {
341 paths
342 .par_iter()
343 .map(|p| (p.to_path_buf(), self.check(p)))
344 .collect()
345 }
346
347 fn list_paths(&self) -> Result<Vec<String>> {
348 ImageCacheReader::list_paths(self)
349 }
350}
351
352pub(crate) fn summarize_entries(
358 entries: Vec<localcache::EntryInfo>,
359) -> Result<Vec<DirCacheSummary>> {
360 use std::collections::BTreeMap;
361
362 let mut map: BTreeMap<PathBuf, (usize, u64, i64)> = BTreeMap::new();
363 for e in entries {
364 let dir = e.path.parent().map(|p| p.to_path_buf()).unwrap_or_default();
365 let agg = map.entry(dir).or_insert((0, 0, 0));
366 agg.0 += 1;
367 agg.1 += e.metadata.file_size;
368 agg.2 = agg.2.max(e.updated_at);
369 }
370 Ok(map
371 .into_iter()
372 .map(
373 |(dir, (file_count, total_size, latest_cached_at))| DirCacheSummary {
374 dir_path: dir.to_string_lossy().into_owned(),
375 file_count,
376 total_size,
377 latest_cached_at,
378 },
379 )
380 .collect())
381}
382
383fn dir_of(path: &Path) -> &Path {
389 if path.is_dir() {
390 path
391 } else {
392 path.parent().unwrap_or(path)
393 }
394}
395
396fn to_image_entry(entry: CacheEntry<ImagePayload>) -> ImageCacheEntry {
401 ImageCacheEntry {
402 path: entry.path.to_string_lossy().into_owned(),
403 thumbnail_path: entry.payload.thumbnail_path,
404 features: entry
405 .payload
406 .clip_vector
407 .map(|v| ImageFeatures { clip_vector: v }),
408 }
409}