···
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
4
+
use crate::config::*;
5
+
use crate::utils::calculate_image_hash;
7
+
use std::collections::HashMap;
8
+
use anyhow::{Context, Result};
9
+
use rand::prelude::*;
10
+
use rusqlite::{Connection, OptionalExtension, Result as SqlResult, params};
11
+
use std::collections::HashSet;
12
+
use std::fs::{self};
13
+
use std::path::{Path, PathBuf};
14
+
use std::process::{Command, Stdio};
15
+
use std::time::{SystemTime, UNIX_EPOCH};
16
+
use walkdir::WalkDir;
18
+
pub 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()))?;
22
+
conn.execute("PRAGMA foreign_keys = ON;", [])?;
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))
36
+
"CREATE TABLE IF NOT EXISTS tags (
37
+
id INTEGER PRIMARY KEY AUTOINCREMENT,
38
+
name TEXT NOT NULL UNIQUE COLLATE NOCASE
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
57
+
pub fn create_indices(conn: &Connection) -> Result<()> {
59
+
"CREATE INDEX IF NOT EXISTS idx_images_path ON images(path);",
63
+
"CREATE INDEX IF NOT EXISTS idx_images_favorite ON images(favorite);",
67
+
"CREATE INDEX IF NOT EXISTS idx_tags_name ON tags(name);",
71
+
"CREATE INDEX IF NOT EXISTS idx_image_tags_tag_id ON image_tags(tag_id);",
77
+
pub fn get_image_id(conn: &Connection, path: &Path) -> SqlResult<Option<String>> {
79
+
"SELECT id FROM images WHERE path = ?",
80
+
params![path.to_string_lossy()],
86
+
pub fn get_or_create_tag_id(conn: &Connection, tag: &str) -> SqlResult<i64> {
87
+
conn.execute("INSERT OR IGNORE INTO tags (name) VALUES (?)", params![tag])?;
90
+
"SELECT id FROM tags WHERE name = ? COLLATE NOCASE",
96
+
pub fn get_or_insert_image_id(conn: &Connection, path: &Path) -> Result<String> {
97
+
let path_str = path.to_string_lossy();
99
+
if let Some(id) = get_image_id(conn, path)? {
103
+
"Image not found in DB, calculating hash and adding: {}",
106
+
let hash_id = calculate_image_hash(path).with_context(|| {
107
+
format!("Failed to calculate hash for new image: {}", path.display())
110
+
let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
113
+
"INSERT INTO images (id, path, last_used, used_count, favorite) VALUES (?, ?, ?, 0, 0)",
114
+
params![hash_id, path_str, now],
118
+
"Failed to insert new image into database: {}",
127
+
pub 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)?;
131
+
let changed = conn.execute(
132
+
"INSERT OR IGNORE INTO image_tags (image_id, tag_id) VALUES (?, ?)",
133
+
params![image_id, tag_id],
137
+
println!("Added tag '{}' to image: {}", tag, path.display());
140
+
"Tag '{}' already present for image: {}",
149
+
pub fn remove_tag(conn: &Connection, path: &Path, tag: &str) -> Result<()> {
150
+
let image_id_opt = get_image_id(conn, path)?;
152
+
if let Some(image_id) = image_id_opt {
153
+
let tag_id_opt: SqlResult<Option<i64>> = conn
155
+
"SELECT id FROM tags WHERE name = ? COLLATE NOCASE",
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],
169
+
println!("Removed tag '{}' from image: {}", tag, path.display());
171
+
println!("Tag '{}' was not found on image: {}", tag, path.display());
175
+
println!("Tag '{}' not found in database.", tag);
178
+
return Err(e).context("Database error checking for tag ID");
182
+
println!("Image not found in database: {}", path.display());
188
+
pub fn list_tags(conn: &Connection, path: &Path) -> Result<()> {
189
+
let image_id_opt = get_image_id(conn, path)?;
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",
198
+
let tags: Vec<String> = stmt
199
+
.query_map(params![image_id], |row| row.get(0))?
200
+
.collect::<SqlResult<_>>()?;
202
+
let favorite: bool = conn.query_row(
203
+
"SELECT favorite FROM images WHERE id = ?",
208
+
println!("Details for image: {}", path.display());
218
+
if tags.is_empty() {
219
+
println!(" Tags: None");
221
+
println!(" Tags:");
223
+
println!(" - {}", tag);
227
+
println!("Image not found in database: {}", path.display());
233
+
pub fn add_favorite(conn: &Connection, path: &Path) -> Result<()> {
234
+
let image_id = get_or_insert_image_id(conn, path)?;
236
+
let changed = conn.execute(
237
+
"UPDATE images SET favorite = 1 WHERE id = ?",
242
+
println!("Marked image as favorite: {}", path.display());
245
+
"Image already marked as favorite or not found for update: {}",
253
+
pub fn remove_favorite(conn: &Connection, path: &Path) -> Result<()> {
254
+
let image_id_opt = get_image_id(conn, path)?;
256
+
if let Some(image_id) = image_id_opt {
257
+
let changed = conn.execute(
258
+
"UPDATE images SET favorite = 0 WHERE id = ?",
263
+
println!("Removed image from favorites: {}", path.display());
266
+
"Image was not marked as favorite or not found for update: {}",
271
+
println!("Image not found in database: {}", path.display());
277
+
pub 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![];
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");
288
+
conditions.push("i.favorite = 1");
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);
298
+
query_parts.push("ORDER BY i.path COLLATE NOCASE");
300
+
let final_query = query_parts.join(" ");
301
+
let mut stmt = conn.prepare(&final_query)?;
303
+
let row_mapper = |row: &rusqlite::Row| {
305
+
row.get::<_, String>(0)?,
306
+
row.get::<_, u32>(1)?,
307
+
row.get::<_, bool>(2)?,
311
+
let images_result = if let Some(t) = tag {
312
+
stmt.query_map(params![t], row_mapper)
314
+
stmt.query_map([], row_mapper)
317
+
let images: Vec<(String, u32, bool)> = images_result?
318
+
.collect::<SqlResult<Vec<_>>>()
319
+
.context("Failed to retrieve image list from database")?;
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(),
328
+
println!("No {} found in the database.", filter_desc);
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(),
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);
344
+
println!("\nTotal: {} images listed.", images.len());
350
+
pub fn get_current_wallpaper(conn: &Connection) -> Result<String> {
352
+
"CREATE TABLE IF NOT EXISTS settings (
353
+
key TEXT PRIMARY KEY NOT NULL,
354
+
value TEXT NOT NULL
359
+
match conn.query_row(
360
+
"SELECT value FROM settings WHERE key = 'current_wallpaper'",
362
+
|row| row.get::<_, String>(0),
365
+
let path_buf = PathBuf::from(&path);
366
+
if path_buf.exists() {
369
+
anyhow::bail!("Current wallpaper file no longer exists: {}", path)
372
+
Err(rusqlite::Error::QueryReturnedNoRows) => {
374
+
"Current wallpaper not set. Please use 'erm random' first or specify a path."
378
+
anyhow::bail!("Error retrieving current wallpaper: {}", e)
384
+
pub fn refresh_database(conn: &mut Connection, config: &Config) -> Result<()> {
385
+
println!("Starting database refresh...");
386
+
let start_time = std::time::Instant::now();
388
+
let tx = conn.transaction()?;
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);
396
+
let mut current_fs_paths: HashSet<PathBuf> = HashSet::new();
397
+
let mut discovered_count = 0;
399
+
let extensions = config
404
+
.map(|s| format!(".{}", s.to_lowercase()))
405
+
.collect::<HashSet<_>>()
407
+
.unwrap_or_else(|| {
408
+
[".jpg", ".jpeg", ".png", ".webp"]
410
+
.map(|&s| s.to_string())
414
+
let exclude_patterns = config
420
+
.filter_map(|p| glob::Pattern::new(p).ok())
421
+
.collect::<Vec<_>>()
423
+
.unwrap_or_default();
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);
429
+
if !start_path.is_dir() {
431
+
"Warning: Configured directory does not exist or is not a directory: {}",
432
+
start_path.display()
437
+
println!("Scanning directory: {}", start_path.display());
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);
445
+
for entry_result in walker.into_iter() {
446
+
let entry = match entry_result {
449
+
eprintln!("Warning: Error traversing directory: {}", e);
454
+
if !entry.file_type().is_file() {
458
+
let file_path = entry.path();
460
+
if let Some(ext) = file_path.extension().and_then(|e| e.to_str()) {
461
+
if !extensions.contains(&format!(".{}", ext.to_lowercase())) {
468
+
let canonical_path = match fs::canonicalize(file_path) {
472
+
"Warning: Could not canonicalize path {}: {}",
473
+
file_path.display(),
479
+
let canonical_path_str = canonical_path.to_string_lossy();
481
+
let should_exclude = exclude_patterns
483
+
.any(|pattern| pattern.matches(&canonical_path_str));
485
+
if should_exclude {
489
+
discovered_count += 1;
490
+
current_fs_paths.insert(canonical_path);
495
+
"Discovered {} potential image files on disk.",
499
+
let db_paths: HashSet<String> = db_images.keys().cloned().collect();
500
+
let fs_paths_str: HashSet<String> = current_fs_paths
502
+
.map(|p| p.to_string_lossy().into_owned())
505
+
let paths_to_remove = db_paths
506
+
.difference(&fs_paths_str)
508
+
.collect::<Vec<_>>();
509
+
let mut removed_count = 0;
510
+
if !paths_to_remove.is_empty() {
512
+
"Removing {} entries from database that no longer exist on disk...",
513
+
paths_to_remove.len()
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),
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() {
529
+
"Adding {} new images to the database...",
532
+
let mut insert_stmt = tx.prepare(
533
+
"INSERT OR IGNORE INTO images (id, path, last_used, used_count, favorite) VALUES (?, ?, 0, 0, 0)",
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) {
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),
549
+
"error calculating hash for {}: {}. Skipping file.",
561
+
let duration = start_time.elapsed();
563
+
"database refresh completed in {:.2?}. Added: {}, Removed: {}{}",
567
+
if hash_errors > 0 {
568
+
format!(", hash errors: {}", hash_errors)
577
+
pub fn select_wallpaper(
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![];
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");
592
+
conditions.push("i.favorite = 1");
595
+
let conditions_str = conditions.join(" AND ");
596
+
query_parts.push(&conditions_str);
598
+
query_parts.push("ORDER BY i.used_count ASC, RANDOM()");
599
+
query_parts.push("LIMIT 50");
601
+
let final_query = query_parts.join(" ");
602
+
let mut stmt = conn.prepare(&final_query)?;
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>>>()?
608
+
stmt.query_map([], |row| row.get(0))?
609
+
.collect::<SqlResult<Vec<String>>>()?
612
+
if candidate_paths.is_empty() {
615
+
let mut rng = rand::rng();
616
+
Ok(candidate_paths.choose(&mut rng).cloned())
620
+
pub fn update_usage(conn: &Connection, path: &str) -> Result<()> {
621
+
let now = SystemTime::now().duration_since(UNIX_EPOCH)?.as_secs();
623
+
let updated_rows = conn.execute(
624
+
"UPDATE images SET last_used = ?, used_count = used_count + 1 WHERE path = ?",
625
+
params![now, path],
628
+
if updated_rows == 0 {
630
+
"Warning: Tried to update usage for path not found in DB: {}",
638
+
pub 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);
644
+
let pgrep_output = Command::new("pgrep").arg("-x").arg("swaybg").output();
646
+
if let Ok(output) = pgrep_output {
647
+
if output.status.success() {
648
+
println!("Found existing swaybg process, terminating it...");
649
+
Command::new("pkill")
653
+
.context("Failed to terminate existing swaybg process")?;
654
+
std::thread::sleep(std::time::Duration::from_millis(100));
658
+
let mut cmd = Command::new("swaybg");
659
+
cmd.arg("-i").arg(path);
661
+
if let Some(args) = &config.swaybg_args {
665
+
println!("Setting wallpaper: {}", path.display());
667
+
.stdin(Stdio::null())
668
+
.stdout(Stdio::null())
669
+
.stderr(Stdio::null())
671
+
.context("Failed to spawn swaybg in daemon mode")?;
672
+
println!("Started swaybg daemon (PID: {})", child.id());
673
+
std::mem::forget(child);
676
+
"CREATE TABLE IF NOT EXISTS settings (
677
+
key TEXT PRIMARY KEY NOT NULL,
678
+
value TEXT NOT NULL
684
+
"INSERT OR REPLACE INTO settings (key, value) VALUES ('current_wallpaper', ?)",
seiso.moe