seiso.moe
1// this entire file could probably be replaced with an actual orm or
2// something but im too lazy to find one thats good enough
3
4use crate::config::*;
5use crate::utils::calculate_image_hash;
6
7use std::collections::HashMap;
8use anyhow::{Context, Result};
9use rand::prelude::*;
10use rusqlite::{Connection, OptionalExtension, Result as SqlResult, params};
11use std::collections::HashSet;
12use std::fs::{self};
13use std::path::{Path, PathBuf};
14use std::process::{Command, Stdio};
15use std::time::{SystemTime, UNIX_EPOCH};
16use walkdir::WalkDir;
17
18pub fn setup_database(path: &Path) -> Result<Connection> {
19 let conn = Connection::open(path)
20 .with_context(|| format!("Failed to open or create database: {}", path.display()))?;
21
22 conn.execute("PRAGMA foreign_keys = ON;", [])?;
23
24 conn.execute(
25 "CREATE TABLE IF NOT EXISTS images (
26 id TEXT PRIMARY KEY NOT NULL,
27 path TEXT NOT NULL UNIQUE,
28 last_used INTEGER NOT NULL DEFAULT 0,
29 used_count INTEGER NOT NULL DEFAULT 0,
30 favorite BOOLEAN NOT NULL DEFAULT 0 CHECK (favorite IN (0, 1))
31 )",
32 [],
33 )?;
34
35 conn.execute(
36 "CREATE TABLE IF NOT EXISTS tags (
37 id INTEGER PRIMARY KEY AUTOINCREMENT,
38 name TEXT NOT NULL UNIQUE COLLATE NOCASE
39 )",
40 [],
41 )?;
42
43 conn.execute(
44 "CREATE TABLE IF NOT EXISTS image_tags (
45 image_id TEXT NOT NULL,
46 tag_id INTEGER NOT NULL,
47 PRIMARY KEY (image_id, tag_id),
48 FOREIGN KEY (image_id) REFERENCES images (id) ON DELETE CASCADE,
49 FOREIGN KEY (tag_id) REFERENCES tags (id) ON DELETE CASCADE
50 )",
51 [],
52 )?;
53
54 Ok(conn)
55}
56
57pub fn create_indices(conn: &Connection) -> Result<()> {
58 conn.execute(
59 "CREATE INDEX IF NOT EXISTS idx_images_path ON images(path);",
60 [],
61 )?;
62 conn.execute(
63 "CREATE INDEX IF NOT EXISTS idx_images_favorite ON images(favorite);",
64 [],
65 )?;
66 conn.execute(
67 "CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name);",
68 [],
69 )?;
70 conn.execute(
71 "CREATE INDEX IF NOT EXISTS idx_image_tags_tag_id ON image_tags(tag_id);",
72 [],
73 )?;
74 Ok(())
75}
76
77pub fn get_image_id(conn: &Connection, path: &Path) -> SqlResult<Option<String>> {
78 conn.query_row(
79 "SELECT id FROM images WHERE path = ?",
80 params![path.to_string_lossy()],
81 |row| row.get(0),
82 )
83 .optional()
84}
85
86pub fn get_or_create_tag_id(conn: &Connection, tag: &str) -> SqlResult<i64> {
87 conn.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", params![tag])?;
88
89 conn.query_row(
90 "SELECT id FROM tags WHERE name = ? COLLATE NOCASE",
91 params![tag],
92 |row| row.get(0),
93 )
94}
95
96pub fn get_or_insert_image_id(conn: &Connection, path: &Path) -> Result<String> {
97 let path_str = path.to_string_lossy();
98
99 if let Some(id) = get_image_id(conn, path)? {
100 Ok(id)
101 } else {
102 println!(
103 "Image not found in DB, calculating hash and adding: {}",
104 path.display()
105 );
106 let hash_id = calculate_image_hash(path).with_context(|| {
107 format!("Failed to calculate hash for new image: {}", path.display())
108 })?;
109
110 let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
111
112 conn.execute(
113 "INSERT INTO images (id, path, last_used, used_count, favorite) VALUES (?, ?, ?, 0, 0)",
114 params![hash_id, path_str, now],
115 )
116 .with_context(|| {
117 format!(
118 "Failed to insert new image into database: {}",
119 path.display()
120 )
121 })?;
122
123 Ok(hash_id)
124 }
125}
126
127pub fn add_tag(conn: &Connection, path: &Path, tag: &str) -> Result<()> {
128 let image_id = get_or_insert_image_id(conn, path)?;
129 let tag_id = get_or_create_tag_id(conn, tag)?;
130
131 let changed = conn.execute(
132 "INSERT OR IGNORE INTO image_tags (image_id, tag_id) VALUES (?, ?)",
133 params![image_id, tag_id],
134 )?;
135
136 if changed > 0 {
137 println!("Added tag '{}' to image: {}", tag, path.display());
138 } else {
139 println!(
140 "Tag '{}' already present for image: {}",
141 tag,
142 path.display()
143 );
144 }
145
146 Ok(())
147}
148
149pub fn remove_tag(conn: &Connection, path: &Path, tag: &str) -> Result<()> {
150 let image_id_opt = get_image_id(conn, path)?;
151
152 if let Some(image_id) = image_id_opt {
153 let tag_id_opt: SqlResult<Option<i64>> = conn
154 .query_row(
155 "SELECT id FROM tags WHERE name = ? COLLATE NOCASE",
156 params![tag],
157 |row| row.get(0),
158 )
159 .optional();
160
161 match tag_id_opt {
162 Ok(Some(tag_id)) => {
163 let changed = conn.execute(
164 "DELETE FROM image_tags WHERE image_id = ? AND tag_id = ?",
165 params![image_id, tag_id],
166 )?;
167
168 if changed > 0 {
169 println!("Removed tag '{}' from image: {}", tag, path.display());
170 } else {
171 println!("Tag '{}' was not found on image: {}", tag, path.display());
172 }
173 }
174 Ok(None) => {
175 println!("Tag '{}' not found in database.", tag);
176 }
177 Err(e) => {
178 return Err(e).context("Database error checking for tag ID");
179 }
180 }
181 } else {
182 println!("Image not found in database: {}", path.display());
183 }
184
185 Ok(())
186}
187
188pub fn list_tags(conn: &Connection, path: &Path) -> Result<()> {
189 let image_id_opt = get_image_id(conn, path)?;
190
191 if let Some(image_id) = image_id_opt {
192 let mut stmt = conn.prepare(
193 "SELECT t.name FROM tags t
194 JOIN image_tags it ON t.id = it.tag_id
195 WHERE it.image_id = ?
196 ORDER BY t.name COLLATE NOCASE",
197 )?;
198 let tags: Vec<String> = stmt
199 .query_map(params![image_id], |row| row.get(0))?
200 .collect::<SqlResult<_>>()?;
201
202 let favorite: bool = conn.query_row(
203 "SELECT favorite FROM images WHERE id = ?",
204 params![image_id],
205 |row| row.get(0),
206 )?;
207
208 println!("Details for image: {}", path.display());
209 println!(
210 " Status: {}",
211 if favorite {
212 "Favorite ★"
213 } else {
214 "Not Favorite"
215 }
216 );
217
218 if tags.is_empty() {
219 println!(" Tags: None");
220 } else {
221 println!(" Tags:");
222 for tag in tags {
223 println!(" - {}", tag);
224 }
225 }
226 } else {
227 println!("Image not found in database: {}", path.display());
228 }
229
230 Ok(())
231}
232
233pub fn add_favorite(conn: &Connection, path: &Path) -> Result<()> {
234 let image_id = get_or_insert_image_id(conn, path)?;
235
236 let changed = conn.execute(
237 "UPDATE images SET favorite = 1 WHERE id = ?",
238 params![image_id],
239 )?;
240
241 if changed > 0 {
242 println!("Marked image as favorite: {}", path.display());
243 } else {
244 println!(
245 "Image already marked as favorite or not found for update: {}",
246 path.display()
247 );
248 }
249
250 Ok(())
251}
252
253pub fn remove_favorite(conn: &Connection, path: &Path) -> Result<()> {
254 let image_id_opt = get_image_id(conn, path)?;
255
256 if let Some(image_id) = image_id_opt {
257 let changed = conn.execute(
258 "UPDATE images SET favorite = 0 WHERE id = ?",
259 params![image_id],
260 )?;
261
262 if changed > 0 {
263 println!("Removed image from favorites: {}", path.display());
264 } else {
265 println!(
266 "Image was not marked as favorite or not found for update: {}",
267 path.display()
268 );
269 }
270 } else {
271 println!("Image not found in database: {}", path.display());
272 }
273
274 Ok(())
275}
276
277pub fn list_images(conn: &Connection, tag: Option<&str>, favorite_only: bool) -> Result<()> {
278 let mut query_parts = vec!["SELECT i.path, i.used_count, i.favorite FROM images i"];
279 let mut conditions = vec![];
280
281 if let Some(_t) = tag {
282 query_parts.push("JOIN image_tags it ON i.id = it.image_id");
283 query_parts.push("JOIN tags t ON it.tag_id = t.id");
284 conditions.push("t.name = ? COLLATE NOCASE");
285 }
286
287 if favorite_only {
288 conditions.push("i.favorite = 1");
289 }
290
291 let conditions_str;
292 if !conditions.is_empty() {
293 query_parts.push("WHERE");
294 conditions_str = conditions.join(" AND ");
295 query_parts.push(&conditions_str);
296 }
297
298 query_parts.push("ORDER BY i.path COLLATE NOCASE");
299
300 let final_query = query_parts.join(" ");
301 let mut stmt = conn.prepare(&final_query)?;
302
303 let row_mapper = |row: &rusqlite::Row| {
304 Ok((
305 row.get::<_, String>(0)?,
306 row.get::<_, u32>(1)?,
307 row.get::<_, bool>(2)?,
308 ))
309 };
310
311 let images_result = if let Some(t) = tag {
312 stmt.query_map(params![t], row_mapper)
313 } else {
314 stmt.query_map([], row_mapper)
315 };
316
317 let images: Vec<(String, u32, bool)> = images_result?
318 .collect::<SqlResult<Vec<_>>>()
319 .context("Failed to retrieve image list from database")?;
320
321 if images.is_empty() {
322 let filter_desc = match (tag, favorite_only) {
323 (Some(t), true) => format!("favorite images with tag '{}'", t),
324 (Some(t), false) => format!("images with tag '{}'", t),
325 (None, true) => "favorite images".to_string(),
326 (None, false) => "all images".to_string(),
327 };
328 println!("No {} found in the database.", filter_desc);
329 } else {
330 let filter_desc = match (tag, favorite_only) {
331 (Some(t), true) => format!("Favorite images with tag '{}'", t),
332 (Some(t), false) => format!("Images with tag '{}'", t),
333 (None, true) => "Favorite images".to_string(),
334 (None, false) => "All images".to_string(),
335 };
336
337 println!("Listing {}:", filter_desc);
338 println!("{:<2} {:<6} {}", "★", "Uses", "Path");
339 println!("{:-<2} {:-<6} {:-<10}", "", "", "");
340 for (path, uses, favorite) in &images {
341 let star = if *favorite { "★" } else { " " };
342 println!("{:<2} {:<6} {}", star, uses, path);
343 }
344 println!("\nTotal: {} images listed.", images.len());
345 }
346
347 Ok(())
348}
349
350pub fn get_current_wallpaper(conn: &Connection) -> Result<String> {
351 conn.execute(
352 "CREATE TABLE IF NOT EXISTS settings (
353 key TEXT PRIMARY KEY NOT NULL,
354 value TEXT NOT NULL
355 )",
356 [],
357 )?;
358
359 match conn.query_row(
360 "SELECT value FROM settings WHERE key = 'current_wallpaper'",
361 [],
362 |row| row.get::<_, String>(0),
363 ) {
364 Ok(path) => {
365 let path_buf = PathBuf::from(&path);
366 if path_buf.exists() {
367 Ok(path)
368 } else {
369 anyhow::bail!("Current wallpaper file no longer exists: {}", path)
370 }
371 }
372 Err(rusqlite::Error::QueryReturnedNoRows) => {
373 anyhow::bail!(
374 "Current wallpaper not set. Please use 'erm random' first or specify a path."
375 )
376 }
377 Err(e) => {
378 anyhow::bail!("Error retrieving current wallpaper: {}", e)
379 }
380 }
381}
382
383
384pub fn refresh_database(conn: &mut Connection, config: &Config) -> Result<()> {
385 println!("Starting database refresh...");
386 let start_time = std::time::Instant::now();
387
388 let tx = conn.transaction()?;
389
390 let mut db_images_stmt = tx.prepare("SELECT id, path FROM images")?;
391 let db_images: HashMap<String, String> = db_images_stmt
392 .query_map([], |row| Ok((row.get(1)?, row.get(0)?)))?
393 .collect::<SqlResult<HashMap<String, String>>>()?;
394 drop(db_images_stmt);
395
396 let mut current_fs_paths: HashSet<PathBuf> = HashSet::new();
397 let mut discovered_count = 0;
398
399 let extensions = config
400 .file_extensions
401 .as_ref()
402 .map(|ext| {
403 ext.iter()
404 .map(|s| format!(".{}", s.to_lowercase()))
405 .collect::<HashSet<_>>()
406 })
407 .unwrap_or_else(|| {
408 [".jpg", ".jpeg", ".png", ".webp"]
409 .iter()
410 .map(|&s| s.to_string())
411 .collect()
412 });
413
414 let exclude_patterns = config
415 .exclude_patterns
416 .as_ref()
417 .map(|patterns| {
418 patterns
419 .iter()
420 .filter_map(|p| glob::Pattern::new(p).ok())
421 .collect::<Vec<_>>()
422 })
423 .unwrap_or_default();
424
425 for dir_path_str in &config.directories {
426 let expanded_path_str = shellexpand::tilde(dir_path_str).into_owned();
427 let start_path = Path::new(&expanded_path_str);
428
429 if !start_path.is_dir() {
430 eprintln!(
431 "Warning: Configured directory does not exist or is not a directory: {}",
432 start_path.display()
433 );
434 continue;
435 }
436
437 println!("Scanning directory: {}", start_path.display());
438
439 let include_subdirs = config.include_subdirectories.unwrap_or(true);
440 let mut walker = WalkDir::new(start_path).follow_links(true);
441 if !include_subdirs {
442 walker = walker.max_depth(1);
443 }
444
445 for entry_result in walker.into_iter() {
446 let entry = match entry_result {
447 Ok(e) => e,
448 Err(e) => {
449 eprintln!("Warning: Error traversing directory: {}", e);
450 continue;
451 }
452 };
453
454 if !entry.file_type().is_file() {
455 continue;
456 }
457
458 let file_path = entry.path();
459
460 if let Some(ext) = file_path.extension().and_then(|e| e.to_str()) {
461 if !extensions.contains(&format!(".{}", ext.to_lowercase())) {
462 continue;
463 }
464 } else {
465 continue;
466 }
467
468 let canonical_path = match fs::canonicalize(file_path) {
469 Ok(p) => p,
470 Err(e) => {
471 eprintln!(
472 "Warning: Could not canonicalize path {}: {}",
473 file_path.display(),
474 e
475 );
476 continue;
477 }
478 };
479 let canonical_path_str = canonical_path.to_string_lossy();
480
481 let should_exclude = exclude_patterns
482 .iter()
483 .any(|pattern| pattern.matches(&canonical_path_str));
484
485 if should_exclude {
486 continue;
487 }
488
489 discovered_count += 1;
490 current_fs_paths.insert(canonical_path);
491 }
492 }
493
494 println!(
495 "Discovered {} potential image files on disk.",
496 discovered_count
497 );
498
499 let db_paths: HashSet<String> = db_images.keys().cloned().collect();
500 let fs_paths_str: HashSet<String> = current_fs_paths
501 .iter()
502 .map(|p| p.to_string_lossy().into_owned())
503 .collect();
504
505 let paths_to_remove = db_paths
506 .difference(&fs_paths_str)
507 .cloned()
508 .collect::<Vec<_>>();
509 let mut removed_count = 0;
510 if !paths_to_remove.is_empty() {
511 println!(
512 "Removing {} entries from database that no longer exist on disk...",
513 paths_to_remove.len()
514 );
515 let mut delete_stmt = tx.prepare("DELETE FROM images WHERE path = ?")?;
516 for path_str in paths_to_remove {
517 match delete_stmt.execute(params![path_str]) {
518 Ok(count) => removed_count += count,
519 Err(e) => eprintln!("Error deleting image with path {}: {}", path_str, e),
520 }
521 }
522 }
523
524 let paths_to_add = fs_paths_str.difference(&db_paths).collect::<Vec<_>>();
525 let mut added_count = 0;
526 let mut hash_errors = 0;
527 if !paths_to_add.is_empty() {
528 println!(
529 "Adding {} new images to the database...",
530 paths_to_add.len()
531 );
532 let mut insert_stmt = tx.prepare(
533 "INSERT OR IGNORE INTO images (id, path, last_used, used_count, favorite) VALUES (?, ?, 0, 0, 0)",
534 )?;
535 for path_str in paths_to_add {
536 print!("added {} to database\r", added_count);
537 let path = PathBuf::from(path_str);
538 match calculate_image_hash(&path) {
539 Ok(hash_id) => {
540 match insert_stmt.execute(params![hash_id, path_str]) {
541 Ok(1) => added_count += 1,
542 Ok(0) => { /* id or path already existed*/ }
543 Ok(_) => { unreachable!() }
544 Err(e) => eprintln!("error inserting image with path {}: {}", path_str, e),
545 }
546 }
547 Err(e) => {
548 eprintln!(
549 "error calculating hash for {}: {}. Skipping file.",
550 path_str, e
551 );
552 hash_errors += 1;
553 }
554 }
555 }
556 println!();
557 }
558
559 tx.commit()?;
560
561 let duration = start_time.elapsed();
562 println!(
563 "database refresh completed in {:.2?}. Added: {}, Removed: {}{}",
564 duration,
565 added_count,
566 removed_count,
567 if hash_errors > 0 {
568 format!(", hash errors: {}", hash_errors)
569 } else {
570 "".to_string()
571 }
572 );
573
574 Ok(())
575}
576
577pub fn select_wallpaper(
578 conn: &Connection,
579 tag: Option<&str>,
580 favorite_only: bool,
581) -> Result<Option<String>> {
582 let mut query_parts = vec!["SELECT i.path FROM images i"];
583 let mut conditions = vec![];
584
585 if let Some(_t) = tag {
586 query_parts.push("JOIN image_tags it ON i.id = it.image_id");
587 query_parts.push("JOIN tags t ON it.tag_id = t.id");
588 conditions.push("t.name = ? COLLATE NOCASE");
589 }
590
591 if favorite_only {
592 conditions.push("where i.favorite = 1");
593 }
594
595 let conditions_str = conditions.join(" AND ");
596 query_parts.push(&conditions_str);
597
598 query_parts.push("ORDER BY i.used_count ASC, RANDOM()");
599 query_parts.push("LIMIT 50");
600
601 let final_query = query_parts.join(" ");
602 let mut stmt = conn.prepare(&final_query)?;
603
604 let candidate_paths: Vec<String> = if let Some(t) = tag {
605 stmt.query_map(params![t], |row| row.get(0))?
606 .collect::<SqlResult<Vec<String>>>()?
607 } else {
608 stmt.query_map([], |row| row.get(0))?
609 .collect::<SqlResult<Vec<String>>>()?
610 };
611
612 if candidate_paths.is_empty() {
613 Ok(None)
614 } else {
615 let mut rng = rand::rng();
616 Ok(candidate_paths.choose(&mut rng).cloned())
617 }
618}
619
620pub fn update_usage(conn: &Connection, path: &str) -> Result<()> {
621 let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
622
623 let updated_rows = conn.execute(
624 "UPDATE images SET last_used = ?, used_count = used_count + 1 WHERE path = ?",
625 params![now, path],
626 )?;
627
628 if updated_rows == 0 {
629 eprintln!(
630 "Warning: Tried to update usage for path not found in DB: {}",
631 path
632 );
633 }
634
635 Ok(())
636}
637
638pub fn set_wallpaper(conn: &Connection, path_str: &str, config: &Config) -> Result<()> {
639 let path = Path::new(path_str);
640 if !path.exists() {
641 anyhow::bail!("Wallpaper file does not exist: {}", path_str);
642 }
643
644 let pgrep_output = Command::new("pgrep").arg("-x").arg("swaybg").output();
645
646 if let Ok(output) = pgrep_output {
647 if output.status.success() {
648 println!("Found existing swaybg process, terminating it...");
649 Command::new("pkill")
650 .arg("-x")
651 .arg("swaybg")
652 .status()
653 .context("Failed to terminate existing swaybg process")?;
654 std::thread::sleep(std::time::Duration::from_millis(100));
655 }
656 }
657
658 let mut cmd = Command::new("swaybg");
659 cmd.arg("-i").arg(path);
660
661 if let Some(args) = &config.swaybg_args {
662 cmd.args(args);
663 }
664
665 println!("Setting wallpaper: {}", path.display());
666 let child = cmd
667 .stdin(Stdio::null())
668 .stdout(Stdio::null())
669 .stderr(Stdio::null())
670 .spawn()
671 .context("Failed to spawn swaybg in daemon mode")?;
672 println!("Started swaybg daemon (PID: {})", child.id());
673 std::mem::forget(child);
674
675 conn.execute(
676 "CREATE TABLE IF NOT EXISTS settings (
677 key TEXT PRIMARY KEY NOT NULL,
678 value TEXT NOT NULL
679 )",
680 [],
681 )?;
682
683 conn.execute(
684 "INSERT OR REPLACE INTO settings (key, value) VALUES ('current_wallpaper', ?)",
685 params![path_str],
686 )?;
687
688 Ok(())
689}