1use std::collections::{HashMap, HashSet};
4use std::path::Path;
5
6use image::DynamicImage;
7use rusqlite::{params, Connection, Result as SqliteResult};
8
9use crate::db::{PhotoStackRepo, StackCandidate};
10use crate::services::path_util::safe_join_relative;
11
12#[derive(Debug, Clone, Default)]
13pub struct StackRefreshResult {
14 pub stacks_found: usize,
15}
16
17#[derive(Debug, Clone)]
18struct CandidatePhoto {
19 id: i64,
20 file_path: String,
21 thumbnail_path: Option<String>,
22 file_size: i64,
23 width: Option<i32>,
24 height: Option<i32>,
25 suggested: bool,
26 face_count: i64,
27}
28
29#[derive(Debug, Clone)]
30struct Score {
31 value: f32,
32 reasons: String,
33}
34
35pub struct PhotoStackService;
36
37impl PhotoStackService {
38 pub fn refresh(conn: &Connection, drive_root: &Path) -> SqliteResult<StackRefreshResult> {
39 let mut candidates = Vec::new();
40 let mut claimed = HashSet::new();
41
42 candidates.extend(Self::duplicate_candidates(
43 conn,
44 drive_root,
45 "exact",
46 "exact_duplicate",
47 1.0,
48 &mut claimed,
49 )?);
50 candidates.extend(Self::duplicate_candidates(
51 conn,
52 drive_root,
53 "perceptual",
54 "perceptual_duplicate",
55 0.94,
56 &mut claimed,
57 )?);
58 candidates.extend(Self::burst_candidates(conn, drive_root, &mut claimed)?);
59
60 let count = candidates.len();
61 PhotoStackRepo::new(conn).sync_stacks(&candidates)?;
62 Ok(StackRefreshResult {
63 stacks_found: count,
64 })
65 }
66
67 fn duplicate_candidates(
68 conn: &Connection,
69 drive_root: &Path,
70 duplicate_type: &str,
71 stack_kind: &str,
72 confidence: f32,
73 claimed: &mut HashSet<i64>,
74 ) -> SqliteResult<Vec<StackCandidate>> {
75 let mut stmt =
76 conn.prepare("SELECT id, group_hash FROM duplicate_groups WHERE duplicate_type = ?1")?;
77 let groups: Vec<(i64, String)> = stmt
78 .query_map(params![duplicate_type], |r| Ok((r.get(0)?, r.get(1)?)))?
79 .collect::<SqliteResult<Vec<_>>>()?;
80
81 let mut out = Vec::new();
82 for (group_id, group_hash) in groups {
83 let mut photos = Self::duplicate_members(conn, group_id)?;
84 photos.retain(|p| !claimed.contains(&p.id));
85 if duplicate_type == "exact" {
86 photos = Self::verified_exact_members(drive_root, photos);
87 }
88 if photos.len() < 2 {
89 continue;
90 }
91 let c = Self::build_candidate(
92 stack_kind,
93 group_id,
94 Some(group_hash),
95 confidence,
96 photos,
97 drive_root,
98 );
99 for id in &c.photo_ids {
100 claimed.insert(*id);
101 }
102 out.push(c);
103 }
104 Ok(out)
105 }
106
107 fn duplicate_members(conn: &Connection, group_id: i64) -> SqliteResult<Vec<CandidatePhoto>> {
108 let mut stmt = conn.prepare(
109 r#"
110 SELECT p.id, p.file_path, p.thumbnail_path, p.file_size, p.width, p.height,
111 m.is_suggested_keep,
112 (SELECT COUNT(*) FROM faces f WHERE f.photo_id = p.id) AS face_count
113 FROM duplicate_group_members m
114 JOIN photos p ON p.id = m.photo_id
115 WHERE m.group_id = ?1 AND p.is_trashed = FALSE
116 ORDER BY m.is_suggested_keep DESC, p.date_taken ASC, p.id ASC
117 "#,
118 )?;
119 let rows = stmt.query_map(params![group_id], row_to_candidate)?;
120 rows.collect()
121 }
122
123 fn verified_exact_members(
124 drive_root: &Path,
125 photos: Vec<CandidatePhoto>,
126 ) -> Vec<CandidatePhoto> {
127 let mut by_hash: HashMap<String, Vec<CandidatePhoto>> = HashMap::new();
128 for p in photos {
129 let abs = match safe_join_relative(drive_root, &p.file_path) {
130 Ok(path) => path,
131 Err(e) => {
132 tracing::debug!(
133 "stack exact verification skipped invalid path {}: {}",
134 p.file_path,
135 e
136 );
137 continue;
138 }
139 };
140 match crate::services::scanner::calculate_hash(&abs) {
141 Ok(hash) => by_hash.entry(hash).or_default().push(p),
142 Err(e) => {
143 tracing::debug!("stack exact verification skipped {}: {}", abs.display(), e)
144 }
145 }
146 }
147 by_hash
148 .into_values()
149 .filter(|v| v.len() >= 2)
150 .max_by_key(|v| v.len())
151 .unwrap_or_default()
152 }
153
154 fn burst_candidates(
155 conn: &Connection,
156 drive_root: &Path,
157 claimed: &mut HashSet<i64>,
158 ) -> SqliteResult<Vec<StackCandidate>> {
159 let mut stmt = conn.prepare("SELECT id FROM burst_groups")?;
160 let groups: Vec<i64> = stmt
161 .query_map([], |r| r.get(0))?
162 .collect::<SqliteResult<Vec<_>>>()?;
163
164 let mut out = Vec::new();
165 for group_id in groups {
166 let mut photos = Self::burst_members(conn, group_id)?;
167 photos.retain(|p| !claimed.contains(&p.id));
168 if photos.len() < 2 || !Self::has_visual_evidence(&photos) {
169 continue;
170 }
171 let c = Self::build_candidate("burst", group_id, None, 0.80, photos, drive_root);
172 for id in &c.photo_ids {
173 claimed.insert(*id);
174 }
175 out.push(c);
176 }
177 Ok(out)
178 }
179
180 fn burst_members(conn: &Connection, group_id: i64) -> SqliteResult<Vec<CandidatePhoto>> {
181 let mut stmt = conn.prepare(
182 r#"
183 SELECT p.id, p.file_path, p.thumbnail_path, p.file_size, p.width, p.height,
184 m.is_suggested_best,
185 (SELECT COUNT(*) FROM faces f WHERE f.photo_id = p.id) AS face_count
186 FROM burst_group_members m
187 JOIN photos p ON p.id = m.photo_id
188 WHERE m.group_id = ?1 AND p.is_trashed = FALSE
189 ORDER BY m.is_suggested_best DESC, p.date_taken ASC, p.id ASC
190 "#,
191 )?;
192 let rows = stmt.query_map(params![group_id], row_to_candidate)?;
193 rows.collect()
194 }
195
196 fn has_visual_evidence(photos: &[CandidatePhoto]) -> bool {
197 photos.iter().all(|p| p.thumbnail_path.is_some())
198 }
199
200 fn build_candidate(
201 kind: &str,
202 source_group_id: i64,
203 source_group_hash: Option<String>,
204 confidence: f32,
205 photos: Vec<CandidatePhoto>,
206 drive_root: &Path,
207 ) -> StackCandidate {
208 let mut scored: Vec<(CandidatePhoto, Score)> = photos
209 .into_iter()
210 .map(|p| {
211 let s = Self::score_photo(&p, drive_root);
212 (p, s)
213 })
214 .collect();
215 scored.sort_by(|a, b| {
216 b.1.value
217 .partial_cmp(&a.1.value)
218 .unwrap_or(std::cmp::Ordering::Equal)
219 .then_with(|| b.0.suggested.cmp(&a.0.suggested))
220 .then_with(|| b.0.file_size.cmp(&a.0.file_size))
221 .then_with(|| a.0.id.cmp(&b.0.id))
222 });
223 let cover_photo_id = scored.first().map(|(p, _)| p.id).unwrap_or_default();
224 StackCandidate {
225 kind: kind.to_string(),
226 source_group_id,
227 source_group_hash,
228 photo_ids: scored.iter().map(|(p, _)| p.id).collect(),
229 cover_photo_id,
230 confidence,
231 member_scores: scored
232 .into_iter()
233 .map(|(p, s)| (p.id, s.value, s.reasons))
234 .collect(),
235 }
236 }
237
238 fn score_photo(photo: &CandidatePhoto, drive_root: &Path) -> Score {
239 let mut score = 0.0f32;
240 let mut reasons = Vec::new();
241
242 if photo.suggested {
243 score += 20.0;
244 reasons.push("existing pick");
245 }
246 let megapixels = match (photo.width, photo.height) {
247 (Some(w), Some(h)) if w > 0 && h > 0 => (w as f32 * h as f32) / 1_000_000.0,
248 _ => 0.0,
249 };
250 if megapixels > 0.0 {
251 score += megapixels.min(60.0) * 1.2;
252 reasons.push("resolution");
253 }
254 if photo.file_size > 0 {
255 score += ((photo.file_size as f32).ln() / 20.0).min(4.0);
256 reasons.push("file size");
257 }
258 if photo.face_count > 0 {
259 score += (photo.face_count as f32).min(5.0) * 2.0;
260 reasons.push("faces");
261 }
262 if let Some(img) = Self::open_scoring_image(photo, drive_root) {
263 let (sharpness, brightness_score) = Self::image_quality(&img);
264 score += sharpness.min(40.0);
265 score += brightness_score * 8.0;
266 reasons.push("sharpness");
267 reasons.push("exposure");
268 }
269
270 Score {
271 value: score,
272 reasons: reasons.join(", "),
273 }
274 }
275
276 fn open_scoring_image(photo: &CandidatePhoto, drive_root: &Path) -> Option<DynamicImage> {
277 let path = photo
278 .thumbnail_path
279 .as_ref()
280 .and_then(|p| safe_join_relative(drive_root, p).ok())
281 .or_else(|| safe_join_relative(drive_root, &photo.file_path).ok())?;
282 crate::services::image_io::open_image(&path).ok()
283 }
284
285 fn image_quality(img: &DynamicImage) -> (f32, f32) {
286 let gray = img
287 .resize(128, 128, image::imageops::FilterType::Triangle)
288 .to_luma8();
289 let (w, h) = gray.dimensions();
290 if w < 3 || h < 3 {
291 return (0.0, 0.0);
292 }
293
294 let mut edge_sum = 0.0f32;
295 let mut count = 0.0f32;
296 let mut luminance = 0.0f32;
297 for y in 1..(h - 1) {
298 for x in 1..(w - 1) {
299 let c = gray.get_pixel(x, y)[0] as f32;
300 let l = gray.get_pixel(x - 1, y)[0] as f32;
301 let r = gray.get_pixel(x + 1, y)[0] as f32;
302 let u = gray.get_pixel(x, y - 1)[0] as f32;
303 let d = gray.get_pixel(x, y + 1)[0] as f32;
304 edge_sum += ((r - l).abs() + (d - u).abs()) / 255.0;
305 luminance += c / 255.0;
306 count += 1.0;
307 }
308 }
309 let sharpness = if count > 0.0 {
310 (edge_sum / count) * 80.0
311 } else {
312 0.0
313 };
314 let avg = if count > 0.0 { luminance / count } else { 0.5 };
315 let brightness_score = (1.0 - ((avg - 0.5).abs() * 2.0)).clamp(0.0, 1.0);
316 (sharpness, brightness_score)
317 }
318}
319
320fn row_to_candidate(row: &rusqlite::Row) -> SqliteResult<CandidatePhoto> {
321 Ok(CandidatePhoto {
322 id: row.get(0)?,
323 file_path: row.get(1)?,
324 thumbnail_path: row.get(2)?,
325 file_size: row.get(3)?,
326 width: row.get(4)?,
327 height: row.get(5)?,
328 suggested: row.get(6)?,
329 face_count: row.get(7)?,
330 })
331}
332
333#[cfg(test)]
334mod tests {
335 use super::*;
336
337 fn candidate(id: i64, width: i32, height: i32, file_size: i64) -> CandidatePhoto {
338 CandidatePhoto {
339 id,
340 file_path: format!("{id}.jpg"),
341 thumbnail_path: None,
342 file_size,
343 width: Some(width),
344 height: Some(height),
345 suggested: false,
346 face_count: 0,
347 }
348 }
349
350 #[test]
351 fn stack_cover_uses_quality_score_not_first_member() {
352 let photos = vec![
353 candidate(1, 800, 600, 250_000),
354 candidate(2, 6000, 4000, 5_000_000),
355 ];
356
357 let stack = PhotoStackService::build_candidate(
358 "burst",
359 10,
360 None,
361 0.8,
362 photos,
363 Path::new("/missing"),
364 );
365
366 assert_eq!(stack.cover_photo_id, 2);
367 }
368}