ImageApi/src/versions/mod.rs

//! Photo versions (#191): AI restorations / upscales / edits that sit near
//! their original and are shown as one photo.
//!
//! ImageApi only *detects* versions; an external process (the ComfyUI node
//! in `comfyui-imageapi-versions`) writes them into the library, so the
//! read-only rule for library roots holds. This module owns
//!
//! - the link rules at ingest ([`link_new_files`], [`sweep_name_links`]):
//!   the PNG's XMP `DerivedFrom` hash first, then the
//!   `<name>.<kind>[.<n>].<ext>` file-name convention;
//! - the HTTP surface (`GET /image/versions`, `PUT /image/versions/display`).
//!
//! The data model lives in [`crate::database::version_dao`]. The link is the
//! original's `content_hash`, never a path: ComfyUI writes through a UNC
//! share and ImageApi reads `/media/...`, and files get moved.

pub mod handlers;
pub mod xmp;

use std::path::{Path, PathBuf};
use std::sync::{Arc, LazyLock, Mutex};

use log::{debug, info, warn};
use regex::Regex;

use crate::database::version_dao::{
    KINDS, LinkOutcome, NewVersionLink, SOURCE_NAME, SOURCE_XMP, VersionDao, is_known_kind,
};
use crate::database::{ExifDao, version_dao};
use crate::libraries::Library;

pub use handlers::{get_versions, put_display};

/// Directory the producer saves versions into, next to the original.
pub const VERSIONS_SUBFOLDER: &str = "aiedits";

/// Cap on name-convention candidates examined per library per tick.
const NAME_SWEEP_LIMIT: i64 = 500;

static VERSION_NAME: LazyLock<Regex> = LazyLock::new(|| {
    let kinds = KINDS.join("|");
    Regex::new(&format!(
        r"(?i)^(?P<name>.+)\.(?P<kind>{kinds})(?:\.(?P<n>\d+))?\.(?P<ext>[a-z0-9]+)$"
    ))
    .unwrap()
});

/// A file name that follows the producer's `<original name>.<kind>[.<n>].<ext>`.
#[derive(Debug, Clone, PartialEq, Eq)]
pub struct VersionName {
    /// The original's full file name, extension included.
    pub original_name: String,
    pub kind: &'static str,
}

/// Parse a bare file name (no directory). The original's extension stays in
/// `original_name`, so `IMG_1.NEF` and `IMG_1.JPG` are told apart.
pub fn parse_version_name(file_name: &str) -> Option<VersionName> {
    let caps = VERSION_NAME.captures(file_name)?;
    let kind_text = caps["kind"].to_ascii_lowercase();
    let kind = KINDS.iter().find(|k| **k == kind_text)?;
    Some(VersionName {
        original_name: caps["name"].to_string(),
        kind,
    })
}

/// Rel paths where the original of a name-convention file may sit, best
/// first: the same directory, then the parent of an `aiedits/` directory.
pub fn original_candidates(rel_path: &str) -> Vec<String> {
    let (dir, file_name) = match rel_path.rsplit_once('/') {
        Some((d, f)) => (d, f),
        None => ("", rel_path),
    };
    let Some(parsed) = parse_version_name(file_name) else {
        return Vec::new();
    };
    let join = |d: &str| {
        if d.is_empty() {
            parsed.original_name.clone()
        } else {
            format!("{d}/{}", parsed.original_name)
        }
    };
    let mut out = vec![join(dir)];
    let (parent, leaf) = match dir.rsplit_once('/') {
        Some((p, l)) => (p, l),
        None => ("", dir),
    };
    if leaf.eq_ignore_ascii_case(VERSIONS_SUBFOLDER) {
        out.push(join(parent));
    }
    out
}

/// Resolve the name-convention link of `rel_path` (in `library`) against the
/// `image_exif` rows of that library. `None` when the name doesn't follow the
/// convention or no candidate original has a hash yet (the sweep retries).
pub fn name_link(
    library: &Library,
    rel_path: &str,
    version_hash: &str,
    exif_dao: &Arc<Mutex<Box<dyn ExifDao>>>,
) -> Option<NewVersionLink> {
    let file_name = rel_path.rsplit('/').next().unwrap_or(rel_path);
    let kind = parse_version_name(file_name)?.kind;
    let candidates = original_candidates(rel_path);
    let rows = exif_dao
        .lock()
        .expect("Unable to lock ExifDao")
        .get_exif_batch(Some(library.id), &candidates)
        .ok()?;
    // First candidate (in preference order) that has a hash of its own.
    let original_hash = candidates.iter().find_map(|c| {
        rows.iter()
            .find(|r| &r.file_path == c)
            .and_then(|r| r.content_hash.clone())
            .filter(|h| h != version_hash)
    })?;
    Some(NewVersionLink {
        version_hash: version_hash.to_string(),
        original_hash,
        kind: kind.to_string(),
        source: SOURCE_NAME.to_string(),
        library_id: library.id,
        rel_path: rel_path.to_string(),
        references: Vec::new(),
    })
}

/// Resolve the XMP link of a PNG on disk. `None` when the file isn't a PNG,
/// can't be read, or carries no `DerivedFrom` blake3 id.
pub fn xmp_link(
    library: &Library,
    rel_path: &str,
    abs_path: &Path,
    version_hash: &str,
) -> Option<NewVersionLink> {
    let is_png = abs_path
        .extension()
        .is_some_and(|e| e.eq_ignore_ascii_case("png"));
    if !is_png {
        return None;
    }
    let parsed = match xmp::read_version_xmp(abs_path) {
        Ok(p) => p?,
        Err(e) => {
            debug!(
                "versions: could not read XMP of {}: {}",
                abs_path.display(),
                e
            );
            return None;
        }
    };
    let file_name = rel_path.rsplit('/').next().unwrap_or(rel_path);
    let kind = parsed
        .kind
        .filter(|k| is_known_kind(k))
        .or_else(|| parse_version_name(file_name).map(|n| n.kind.to_string()))
        .unwrap_or_else(|| "edit".to_string());
    Some(NewVersionLink {
        version_hash: version_hash.to_string(),
        original_hash: parsed.original_hash,
        kind,
        source: SOURCE_XMP.to_string(),
        library_id: library.id,
        rel_path: rel_path.to_string(),
        references: parsed.references,
    })
}

/// Record a link, logging the outcome. Returns true if a new link was made.
pub fn record_link(version_dao: &Arc<Mutex<Box<dyn VersionDao>>>, link: &NewVersionLink) -> bool {
    let outcome = version_dao
        .lock()
        .expect("Unable to lock VersionDao")
        .link_version(link);
    match outcome {
        Ok(LinkOutcome::Linked) => {
            info!(
                "versions: linked {} -> original {}.. ({}, {})",
                link.rel_path,
                &link.original_hash[..12.min(link.original_hash.len())],
                link.kind,
                link.source
            );
            true
        }
        Ok(LinkOutcome::AlreadyLinked) => false,
        Ok(LinkOutcome::Rejected(why)) => {
            warn!("versions: refused link for {}: {}", link.rel_path, why);
            false
        }
        Err(e) => {
            warn!("versions: link failed for {}: {:?}", link.rel_path, e);
            false
        }
    }
}

/// Link freshly ingested files by their embedded XMP. `files` are the
/// `(absolute, rel_path)` pairs the watcher just registered; rows without a
/// hash yet are skipped (an XMP-less file is the name sweep's job anyway).
/// Returns the number of new links.
pub fn link_new_files(
    library: &Library,
    files: &[(PathBuf, String)],
    exif_dao: &Arc<Mutex<Box<dyn ExifDao>>>,
    version_dao: &Arc<Mutex<Box<dyn VersionDao>>>,
) -> usize {
    let pngs: Vec<&(PathBuf, String)> = files
        .iter()
        .filter(|(p, _)| p.extension().is_some_and(|e| e.eq_ignore_ascii_case("png")))
        .collect();
    if pngs.is_empty() {
        return 0;
    }
    let rels: Vec<String> = pngs.iter().map(|(_, r)| r.clone()).collect();
    let rows = exif_dao
        .lock()
        .expect("Unable to lock ExifDao")
        .get_exif_batch(Some(library.id), &rels)
        .unwrap_or_default();

    let mut linked = 0;
    for (abs, rel) in pngs {
        let Some(hash) = rows
            .iter()
            .find(|r| &r.file_path == rel)
            .and_then(|r| r.content_hash.clone())
        else {
            continue;
        };
        if let Some(link) = xmp_link(library, rel, abs, &hash)
            && record_link(version_dao, &link)
        {
            linked += 1;
        }
    }
    linked
}

/// Link name-convention files whose original is (now) known. Idempotent and
/// DB-only, so it also serves as the retry for "original not hashed yet" and
/// "original ingested after the version". Returns the number of new links.
pub fn sweep_name_links(
    library: &Library,
    exif_dao: &Arc<Mutex<Box<dyn ExifDao>>>,
    version_dao: &Arc<Mutex<Box<dyn VersionDao>>>,
) -> usize {
    let candidates = match version_dao
        .lock()
        .expect("Unable to lock VersionDao")
        .unlinked_name_candidates(library.id, NAME_SWEEP_LIMIT)
    {
        Ok(c) => c,
        Err(e) => {
            warn!(
                "versions: name sweep failed for '{}': {:?}",
                library.name, e
            );
            return 0;
        }
    };
    let mut linked = 0;
    for (rel, hash) in candidates {
        if let Some(link) = name_link(library, &rel, &hash, exif_dao)
            && record_link(version_dao, &link)
        {
            linked += 1;
        }
    }
    linked
}

/// Sweep names and invalidate what a new link makes stale. The watcher calls
/// this once per library per tick, after ingest.
pub fn link_pending(
    library: &Library,
    exif_dao: &Arc<Mutex<Box<dyn ExifDao>>>,
    version_dao: &Arc<Mutex<Box<dyn VersionDao>>>,
) {
    if sweep_name_links(library, exif_dao, version_dao) > 0 {
        crate::duplicates::invalidate_caches();
    }
}

/// The version facts a per-photo record (a memory, a search hit) carries.
/// Flattened into the record; absent for a photo without versions, so
/// responses stay byte-identical for libraries that have none.
///
/// `display_path` is `null` when the original is displayed, a rel_path when a
/// version is. `/photos` carries the same facts as the parallel arrays
/// `photo_display_paths` / `photo_version_counts` instead.
#[derive(Debug, Clone, Default, serde::Serialize)]
pub struct VersionFields {
    #[serde(skip_serializing_if = "Option::is_none")]
    pub display_path: Option<Option<String>>,
    #[serde(skip_serializing_if = "Option::is_none")]
    pub version_count: Option<u32>,
}

impl VersionFields {
    pub fn from_summary(summary: Option<&version_dao::PhotoVersionSummary>) -> Self {
        match summary {
            Some(s) if s.version_count > 0 => VersionFields {
                display_path: Some(s.display.as_ref().map(|(_, p)| p.clone())),
                version_count: Some(s.version_count),
            },
            _ => VersionFields::default(),
        }
    }
}

/// Load the per-photo version summaries once for a response. Empty (and
/// cheap) when nothing has versions; a DB failure degrades to "no versions"
/// rather than failing the listing.
pub fn load_summaries(
    version_dao: &Mutex<Box<dyn VersionDao>>,
) -> std::collections::HashMap<(i32, String), version_dao::PhotoVersionSummary> {
    version_dao
        .lock()
        .expect("Unable to lock VersionDao")
        .photo_summaries()
        .unwrap_or_else(|e| {
            warn!("versions: could not load summaries: {:?}", e);
            Default::default()
        })
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn parses_the_producer_naming_convention() {
        let n = parse_version_name("scan_0042.jpg.restored.png").unwrap();
        assert_eq!(n.original_name, "scan_0042.jpg");
        assert_eq!(n.kind, "restored");
        let n = parse_version_name("scan_0042.jpg.restored.2.png").unwrap();
        assert_eq!(n.original_name, "scan_0042.jpg");
        let n = parse_version_name("IMG_1.NEF.upscale.png").unwrap();
        assert_eq!(n.original_name, "IMG_1.NEF");
        assert_eq!(n.kind, "upscale");
        assert_eq!(
            parse_version_name("a.b.c.EDIT.12.PNG")
                .unwrap()
                .original_name,
            "a.b.c"
        );
    }

    #[test]
    fn rejects_names_outside_the_convention() {
        assert!(parse_version_name("holiday.png").is_none());
        assert!(parse_version_name("restored.png").is_none());
        assert!(parse_version_name("x.jpg.blurred.png").is_none());
        assert!(parse_version_name("x.jpg.restored").is_none());
        assert!(parse_version_name("x.restored.png.bak").is_none());
        assert!(parse_version_name(".restored.png").is_none());
    }

    #[test]
    fn candidates_prefer_the_same_dir_then_the_parent_of_aiedits() {
        assert_eq!(
            original_candidates("2004/aiedits/scan.jpg.restored.png"),
            vec!["2004/aiedits/scan.jpg", "2004/scan.jpg"]
        );
        assert_eq!(
            original_candidates("2004/scan.jpg.restored.png"),
            vec!["2004/scan.jpg"]
        );
        assert_eq!(
            original_candidates("aiedits/scan.jpg.edit.3.png"),
            vec!["aiedits/scan.jpg", "scan.jpg"]
        );
        assert_eq!(
            original_candidates("scan.jpg.upscale.png"),
            vec!["scan.jpg"]
        );
        assert!(original_candidates("2004/plain.png").is_empty());
    }
}