ImageApi/src/enhance/path_map.rs

//! Library path <-> ComfyUI path mapping (#192).
//!
//! ImageApi sees `/media/nas/Media/Pictures/a.jpg`; ComfyUI (on Windows)
//! sees `Z:\Media\Pictures\a.jpg` or `\\nas\Media\Pictures\a.jpg`. Each
//! mapping is a prefix pair; the longest matching prefix wins and the
//! separators are converted. Used in both directions: library -> ComfyUI on
//! submit, and ComfyUI -> library to recognise the path `approve` returns.
//! With no mappings configured, paths pass through unchanged.

use std::path::{Path, PathBuf};

use anyhow::{Result, bail};

use super::config::PathMapping;

#[derive(Debug, Clone)]
struct Entry {
    /// Library side, `/`-separated, always ending in `/`.
    from: String,
    /// ComfyUI side without a trailing separator.
    to: String,
    /// Separator ComfyUI's side uses.
    to_sep: char,
}

#[derive(Debug, Clone, Default)]
pub struct PathMap {
    entries: Vec<Entry>,
}

impl PathMap {
    pub fn new(mappings: &[PathMapping]) -> Result<Self> {
        let mut entries = Vec::new();
        for m in mappings {
            let from = m.from.replace('\\', "/");
            let to = m.to.trim();
            if from.trim_matches('/').is_empty() || to.is_empty() {
                bail!("path_map entries need non-empty from and to ({m:?})");
            }
            let to_sep = if to.contains('\\') { '\\' } else { '/' };
            entries.push(Entry {
                from: format!("{}/", from.trim_end_matches('/')),
                to: to.trim_end_matches(['/', '\\']).to_string(),
                to_sep,
            });
        }
        Ok(PathMap { entries })
    }

    pub fn is_identity(&self) -> bool {
        self.entries.is_empty()
    }

    /// The path ComfyUI should open for a library file. `None` when mappings
    /// exist but none covers the file.
    pub fn to_comfy(&self, abs: &Path) -> Option<String> {
        let text = abs.to_str()?.replace('\\', "/");
        if self.entries.is_empty() {
            return Some(abs.to_string_lossy().into_owned());
        }
        let best = self
            .entries
            .iter()
            .filter(|e| text.starts_with(e.from.as_str()))
            .max_by_key(|e| e.from.len())?;
        let rest = &text[best.from.len()..];
        let sep = best.to_sep.to_string();
        Some(format!("{}{}{}", best.to, sep, rest.replace('/', &sep)))
    }

    /// The library-side absolute path of a path ComfyUI reported. Matching
    /// is case-insensitive (Windows). `None` when no mapping covers it or it
    /// tries to climb out with `..`.
    pub fn from_comfy(&self, comfy: &str) -> Option<PathBuf> {
        let text = comfy.trim().replace('\\', "/");
        if self.entries.is_empty() {
            return Some(PathBuf::from(text));
        }
        let best = self
            .entries
            .iter()
            .filter_map(|e| {
                let prefix = format!("{}/", e.to.replace('\\', "/"));
                let head = text.get(..prefix.len())?;
                head.eq_ignore_ascii_case(&prefix)
                    .then_some((e, prefix.len()))
            })
            .max_by_key(|(_, len)| *len)?;
        let rest = &text[best.1..];
        if rest.split('/').any(|c| c == "..") {
            return None;
        }
        Some(PathBuf::from(format!("{}{}", best.0.from, rest)))
    }
}

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

    fn map(pairs: &[(&str, &str)]) -> PathMap {
        let m: Vec<PathMapping> = pairs
            .iter()
            .map(|(f, t)| PathMapping {
                from: f.to_string(),
                to: t.to_string(),
            })
            .collect();
        PathMap::new(&m).unwrap()
    }

    #[test]
    fn drive_letter_round_trip() {
        let m = map(&[("/media/nas/", "Z:\\")]);
        let c = m
            .to_comfy(Path::new("/media/nas/Media/Pictures/a b.jpg"))
            .unwrap();
        assert_eq!(c, "Z:\\Media\\Pictures\\a b.jpg");
        assert_eq!(
            m.from_comfy(&c).unwrap(),
            PathBuf::from("/media/nas/Media/Pictures/a b.jpg")
        );
        // Drive letters are case-insensitive.
        assert_eq!(
            m.from_comfy("z:\\x\\y.png").unwrap(),
            PathBuf::from("/media/nas/x/y.png")
        );
    }

    #[test]
    fn unc_target_and_missing_trailing_slash() {
        let m = map(&[("/media/nas", "\\\\nas\\")]);
        assert_eq!(
            m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(),
            "\\\\nas\\Media\\a.jpg"
        );
        assert_eq!(
            m.from_comfy("\\\\nas\\Media\\a.jpg").unwrap(),
            PathBuf::from("/media/nas/Media/a.jpg")
        );
    }

    #[test]
    fn longest_prefix_wins_both_ways() {
        let m = map(&[("/media/nas/", "Z:\\"), ("/media/nas/Media/", "Y:\\")]);
        assert_eq!(
            m.to_comfy(Path::new("/media/nas/Media/a.jpg")).unwrap(),
            "Y:\\a.jpg"
        );
        assert_eq!(
            m.to_comfy(Path::new("/media/nas/Other/a.jpg")).unwrap(),
            "Z:\\Other\\a.jpg"
        );
        assert_eq!(
            m.from_comfy("Y:\\a.jpg").unwrap(),
            PathBuf::from("/media/nas/Media/a.jpg")
        );
    }

    #[test]
    fn uncovered_and_escaping_paths_are_none() {
        let m = map(&[("/media/nas/", "Z:\\")]);
        assert!(m.to_comfy(Path::new("/elsewhere/a.jpg")).is_none());
        assert!(m.to_comfy(Path::new("/media/nasty/a.jpg")).is_none());
        assert!(m.from_comfy("C:\\a.png").is_none());
        assert!(m.from_comfy("Z:\\..\\etc\\passwd").is_none());
    }

    #[test]
    fn no_mappings_is_identity() {
        let m = PathMap::default();
        assert_eq!(m.to_comfy(Path::new("/a/b.jpg")).unwrap(), "/a/b.jpg");
        assert_eq!(m.from_comfy("/a/b.png").unwrap(), PathBuf::from("/a/b.png"));
    }
}