mnote/rust/crates/mnote-web/src/local_folder_watcher_registry.rs

use crate::routes::{local_workspace_id_from_root_uri, refresh_local_search_index_for_path};
use notify::event::ModifyKind;
use notify::{Config, Event, EventKind, RecommendedWatcher, RecursiveMode, Watcher};
use serde_json::{json, Value};
use std::collections::HashMap;
use std::path::{Path, PathBuf};
use std::sync::atomic::{AtomicUsize, Ordering};
use std::sync::{Arc, Mutex, Weak};
use std::time::{SystemTime, UNIX_EPOCH};
use tokio::sync::{broadcast, mpsc, oneshot};

#[derive(Clone)]
pub struct LocalFolderWatcherRegistry {
    inner: Arc<LocalFolderWatcherRegistryInner>,
}

impl std::fmt::Debug for LocalFolderWatcherRegistry {
    fn fmt(&self, formatter: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        formatter
            .debug_struct("LocalFolderWatcherRegistry")
            .finish_non_exhaustive()
    }
}

impl LocalFolderWatcherRegistry {
    pub fn new() -> Self {
        Self {
            inner: Arc::new(LocalFolderWatcherRegistryInner {
                entries: Mutex::new(HashMap::new()),
            }),
        }
    }

    pub(crate) fn subscribe(
        &self,
        canonical_root: &Path,
    ) -> Result<LocalFolderWatcherSubscription, String> {
        let key = canonical_root_uri(canonical_root);
        let channel = self.inner.get_or_create_channel(&key, canonical_root)?;
        channel.subscriber_count.fetch_add(1, Ordering::SeqCst);
        Ok(LocalFolderWatcherSubscription {
            receiver: channel.sender.subscribe(),
            guard: LocalFolderWatcherSubscriptionGuard {
                registry: Arc::downgrade(&self.inner),
                key,
                channel,
            },
        })
    }

    #[cfg(test)]
    pub fn active_watcher_count(&self) -> usize {
        self.inner.entries.lock().expect("registry lock").len()
    }
}

struct LocalFolderWatcherRegistryInner {
    entries: Mutex<HashMap<String, Arc<LocalFolderWatchChannel>>>,
}

impl LocalFolderWatcherRegistryInner {
    fn get_or_create_channel(
        &self,
        key: &str,
        canonical_root: &Path,
    ) -> Result<Arc<LocalFolderWatchChannel>, String> {
        if let Some(existing) = self
            .entries
            .lock()
            .expect("registry lock")
            .get(key)
            .cloned()
        {
            return Ok(existing);
        }

        let channel = Arc::new(LocalFolderWatchChannel::new(
            key.to_string(),
            spawn_local_folder_watcher(key, canonical_root.to_path_buf())?,
        ));

        let mut entries = self.entries.lock().expect("registry lock");
        if let Some(existing) = entries.get(key).cloned() {
            return Ok(existing);
        }
        entries.insert(key.to_string(), channel.clone());
        Ok(channel)
    }

    fn remove_if_idle(&self, key: &str, channel: &Arc<LocalFolderWatchChannel>) {
        let mut entries = self.entries.lock().expect("registry lock");
        let should_remove = entries
            .get(key)
            .map(|current| {
                Arc::ptr_eq(current, channel)
                    && channel.subscriber_count.load(Ordering::SeqCst) == 0
            })
            .unwrap_or(false);
        if should_remove {
            entries.remove(key);
            channel.shutdown();
        }
    }
}

struct LocalFolderWatchChannel {
    sender: broadcast::Sender<Value>,
    subscriber_count: AtomicUsize,
    shutdown_tx: Mutex<Option<oneshot::Sender<()>>>,
}

impl LocalFolderWatchChannel {
    fn new(_root_uri: String, parts: (broadcast::Sender<Value>, oneshot::Sender<()>)) -> Self {
        Self {
            sender: parts.0,
            subscriber_count: AtomicUsize::new(0),
            shutdown_tx: Mutex::new(Some(parts.1)),
        }
    }

    fn shutdown(&self) {
        if let Some(sender) = self.shutdown_tx.lock().expect("shutdown lock").take() {
            let _ = sender.send(());
        }
    }
}

pub(crate) struct LocalFolderWatcherSubscription {
    pub(crate) receiver: broadcast::Receiver<Value>,
    guard: LocalFolderWatcherSubscriptionGuard,
}

impl LocalFolderWatcherSubscription {
    pub(crate) fn root_uri(&self) -> &str {
        self.guard.key.as_str()
    }
}

struct LocalFolderWatcherSubscriptionGuard {
    registry: Weak<LocalFolderWatcherRegistryInner>,
    key: String,
    channel: Arc<LocalFolderWatchChannel>,
}

impl Drop for LocalFolderWatcherSubscriptionGuard {
    fn drop(&mut self) {
        let previous = self.channel.subscriber_count.fetch_sub(1, Ordering::SeqCst);
        if previous != 1 {
            return;
        }
        if let Some(registry) = self.registry.upgrade() {
            registry.remove_if_idle(&self.key, &self.channel);
        }
    }
}

fn spawn_local_folder_watcher(
    root_uri: &str,
    canonical_root: PathBuf,
) -> Result<(broadcast::Sender<Value>, oneshot::Sender<()>), String> {
    let (event_sender, mut event_receiver) = mpsc::unbounded_channel::<notify::Result<Event>>();
    let mut watcher = RecommendedWatcher::new(
        move |result| {
            let _ = event_sender.send(result);
        },
        Config::default(),
    )
    .map_err(|error| format!("本地文件事件监听启动失败: {error}"))?;
    watcher
        .watch(&canonical_root, RecursiveMode::Recursive)
        .map_err(|error| format!("本地文件夹监听失败: {error}"))?;

    let (sender, _) = broadcast::channel::<Value>(256);
    let (shutdown_tx, mut shutdown_rx) = oneshot::channel::<()>();
    let sender_for_task = sender.clone();
    let root_uri_for_task = root_uri.to_string();
    tokio::spawn(async move {
        let _watcher = watcher;
        loop {
            tokio::select! {
                _ = &mut shutdown_rx => {
                    return;
                }
                maybe_result = event_receiver.recv() => {
                    let Some(result) = maybe_result else {
                        return;
                    };
                    let Ok(event) = result else {
                        continue;
                    };
                    if !should_emit_event_kind(&event.kind) {
                        continue;
                    }
                    for path in event.paths {
                        if !is_local_search_index_path(&path) {
                            continue;
                        }
                        let Some(relative_path) = relative_path_string(&canonical_root, &path) else {
                            continue;
                        };
                        refresh_local_search_index_for_event(
                            &canonical_root,
                            &root_uri_for_task,
                            &relative_path,
                        );
                        if !is_markdown_path(&path) {
                            continue;
                        }
                        let payload = json!({
                            "sourceKind": "local_folder",
                            "rootUri": root_uri_for_task,
                            "relativePath": relative_path,
                            "documentId": format!("local-md:{}", encode_local_id_segment(&relative_path)),
                            "eventKind": format!("{:?}", event.kind),
                            "revision": event_revision(&path),
                        });
                        let _ = sender_for_task.send(payload);
                    }
                }
            }
        }
    });

    Ok((sender, shutdown_tx))
}

fn refresh_local_search_index_for_event(root: &Path, root_uri: &str, relative_path: &str) {
    let Ok(workspace_id) = local_workspace_id_from_root_uri(root_uri) else {
        return;
    };
    let _ = refresh_local_search_index_for_path(root, root_uri, &workspace_id, relative_path);
}

fn canonical_root_uri(root: &Path) -> String {
    format!("file://{}", root.display())
}

fn relative_path_string(root: &Path, path: &Path) -> Option<String> {
    path.strip_prefix(root)
        .ok()
        .map(|relative| relative.to_string_lossy().replace('\\', "/"))
        .filter(|relative| !relative.is_empty())
}

fn is_markdown_path(path: &Path) -> bool {
    path.extension()
        .and_then(|extension| extension.to_str())
        .map(|extension| {
            extension.eq_ignore_ascii_case("md") || extension.eq_ignore_ascii_case("markdown")
        })
        .unwrap_or(false)
}

fn is_local_search_index_path(path: &Path) -> bool {
    if is_markdown_path(path) {
        return true;
    }
    let file_name = path
        .file_name()
        .and_then(|value| value.to_str())
        .unwrap_or_default()
        .to_lowercase();
    if file_name.ends_with(".mindmap.json") {
        return true;
    }
    let extension = path
        .extension()
        .and_then(|extension| extension.to_str())
        .unwrap_or_default()
        .to_lowercase();
    matches!(
        extension.as_str(),
        "doc" | "docx" | "odt" | "ppt" | "pptx" | "odp" | "xls" | "xlsx" | "ods"
    )
}

fn should_emit_event_kind(kind: &EventKind) -> bool {
    match kind {
        EventKind::Create(_) | EventKind::Remove(_) => true,
        EventKind::Modify(modify_kind) => matches!(
            modify_kind,
            ModifyKind::Any | ModifyKind::Data(_) | ModifyKind::Name(_) | ModifyKind::Metadata(_)
        ),
        _ => false,
    }
}

fn encode_local_id_segment(value: &str) -> String {
    let mut encoded = String::with_capacity(value.len());
    for byte in value.as_bytes() {
        let character = *byte as char;
        if character.is_ascii_alphanumeric() || matches!(character, '.' | '_' | '-') {
            encoded.push(character);
        } else {
            encoded.push('~');
            encoded.push_str(&format!("{byte:02X}"));
        }
    }
    encoded
}

fn event_revision(path: &Path) -> u128 {
    std::fs::metadata(path)
        .ok()
        .and_then(|metadata| metadata.modified().ok())
        .map(system_time_ms)
        .unwrap_or_else(|| system_time_ms(SystemTime::now()))
}

fn system_time_ms(time: SystemTime) -> u128 {
    time.duration_since(UNIX_EPOCH)
        .map(|duration| duration.as_millis())
        .unwrap_or(0)
}

#[cfg(test)]
mod tests {
    use super::{
        is_local_search_index_path, refresh_local_search_index_for_event, should_emit_event_kind,
        LocalFolderWatcherRegistry,
    };
    use notify::event::{AccessKind, CreateKind, DataChange, ModifyKind};
    use notify::EventKind;

    fn test_root(name: &str) -> std::path::PathBuf {
        let root = std::env::temp_dir().join(format!(
            "mnote-local-folder-watcher-{name}-{}-{}",
            std::process::id(),
            std::time::SystemTime::now()
                .duration_since(std::time::UNIX_EPOCH)
                .map(|duration| duration.as_nanos())
                .unwrap_or(0)
        ));
        std::fs::create_dir_all(&root).expect("create temp root");
        root
    }

    #[tokio::test]
    async fn same_root_subscribers_share_single_watcher() {
        let registry = LocalFolderWatcherRegistry::new();
        let root = test_root("shared");

        let first = registry.subscribe(&root).expect("first subscription");
        let second = registry.subscribe(&root).expect("second subscription");

        assert_eq!(registry.active_watcher_count(), 1);

        drop(first);
        tokio::task::yield_now().await;
        assert_eq!(registry.active_watcher_count(), 1);

        drop(second);
        tokio::task::yield_now().await;
        assert_eq!(registry.active_watcher_count(), 0);

        let _ = std::fs::remove_dir_all(root);
    }

    #[tokio::test]
    async fn different_roots_create_independent_watchers() {
        let registry = LocalFolderWatcherRegistry::new();
        let first_root = test_root("first");
        let second_root = test_root("second");

        let first = registry
            .subscribe(&first_root)
            .expect("first root subscription");
        let second = registry
            .subscribe(&second_root)
            .expect("second root subscription");

        assert_eq!(registry.active_watcher_count(), 2);

        drop(first);
        drop(second);
        tokio::task::yield_now().await;
        assert_eq!(registry.active_watcher_count(), 0);

        let _ = std::fs::remove_dir_all(first_root);
        let _ = std::fs::remove_dir_all(second_root);
    }

    #[test]
    fn event_kind_filter_ignores_access_events() {
        assert!(should_emit_event_kind(&EventKind::Create(CreateKind::File)));
        assert!(should_emit_event_kind(&EventKind::Modify(
            ModifyKind::Data(DataChange::Content,)
        )));
        assert!(!should_emit_event_kind(&EventKind::Access(
            AccessKind::Read
        )));
    }

    #[test]
    fn watcher_index_path_filter_accepts_markdown_and_resources() {
        assert!(is_local_search_index_path(&std::path::Path::new(
            "docs/page.md"
        )));
        assert!(is_local_search_index_path(&std::path::Path::new(
            "maps/idea.mindmap.json"
        )));
        assert!(is_local_search_index_path(&std::path::Path::new(
            "office/report.xlsx"
        )));
        assert!(!is_local_search_index_path(&std::path::Path::new(
            "docs/image.png"
        )));
    }

    #[test]
    fn watcher_event_refreshes_local_search_index_path() {
        let root = test_root("search-index-event");
        std::fs::create_dir_all(root.join("docs")).expect("create docs");
        std::fs::write(root.join("README.md"), "# Home\n").expect("write home");
        std::fs::write(
            root.join("docs").join("watched.md"),
            "---\ntitle: Watched\n---\n# Watched\nWatcherToken\n",
        )
        .expect("write watched");

        let root_uri = format!("file://{}", root.display());
        refresh_local_search_index_for_event(&root, &root_uri, "docs/watched.md");

        let index_path = root.join(".mnote").join("index").join("search-index.json");
        let index = std::fs::read_to_string(&index_path).expect("index exists");
        assert!(index.contains("docs/watched.md"));
        assert!(index.contains("WatcherToken"));

        let _ = std::fs::remove_dir_all(root);
    }
}