//! Observes successful namespace and open-file mutations, even if their caller is cancelled.
use crate::*;
use std::{
    collections::BTreeMap,
    sync::{Mutex, Weak},
};
use tokio::sync::{OwnedRwLockReadGuard, OwnedRwLockWriteGuard, RwLock};

type Namespace = Mutex<BTreeMap<String, Weak<RwLock<()>>>>;
enum Held {
    Read { _guard: OwnedRwLockReadGuard<()> },
    Write { _guard: OwnedRwLockWriteGuard<()> },
}

type TrackedPath = Mutex<Option<String>>;
type Path = Arc<TrackedPath>;
type Callback = dyn Fn(&str) + Send + Sync;
/// Shared observation for mount namespaces. Paths supplied to the callback are virtual paths.
/// Callbacks run after filesystem locks have been released; they must not block.
#[derive(Clone)]
pub struct Observer {
    changed: Arc<Callback>,
    paths: Arc<Mutex<Vec<Weak<TrackedPath>>>>,
    // Only open-registration versus rename/remove needs ordering. Lock the affected
    // subtrees, not all IO: a blocked open must not stall unrelated namespaces.
    namespace: Arc<Namespace>,
}
impl Observer {
    pub fn new(changed: impl Fn(&str) + Send + Sync + 'static) -> Self {
        Self {
            changed: Arc::new(changed),
            paths: Arc::default(),
            namespace: Arc::default(),
        }
    }
    /// Wrap a layer at its canonical namespace prefix (e.g. `/run`).
    pub fn layer(&self, layer: Arc<dyn Layer>, prefix: &str, process: Process) -> Arc<dyn Layer> {
        Arc::new(Observed {
            layer,
            prefix: prefix.trim_end_matches('/').into(),
            observer: self.clone(),
            process,
        })
    }
    /// Opens share their path and ancestors; rename/remove exclusively hold their
    /// endpoints, sharing ancestors. Sorted acquisition avoids cross-rename deadlocks.
    /// No descriptor IO or notification holds these guards.
    async fn order(&self, endpoints: &[(&str, bool)]) -> Vec<Held> {
        let mut modes = BTreeMap::new();
        for &(path, exclusive) in endpoints {
            modes
                .entry(path.to_owned())
                .and_modify(|mode| *mode |= exclusive)
                .or_insert(exclusive);
            let mut ancestor = path;
            while let Some((parent, _)) = ancestor.rsplit_once('/') {
                let parent = if parent.is_empty() { "/" } else { parent };
                modes.entry(parent.to_owned()).or_insert(false);
                if parent == "/" {
                    break;
                }
                ancestor = parent;
            }
        }
        let locks: Vec<_> = {
            let mut namespace = self.namespace.lock().unwrap();
            namespace.retain(|_, lock| lock.strong_count() != 0);
            modes
                .into_iter()
                .map(|(path, exclusive)| {
                    let lock = namespace.entry(path).or_default();
                    let strong = lock.upgrade().unwrap_or_else(|| {
                        let strong = Arc::new(RwLock::new(()));
                        *lock = Arc::downgrade(&strong);
                        strong
                    });
                    (strong, exclusive)
                })
                .collect()
        };
        let mut held = Vec::new();
        for (lock, exclusive) in locks {
            held.push(if exclusive {
                Held::Write {
                    _guard: lock.write_owned().await,
                }
            } else {
                Held::Read {
                    _guard: lock.read_owned().await,
                }
            });
        }
        held
    }
    fn path(&self, path: String) -> Path {
        let path = Arc::new(Mutex::new(Some(path)));
        let mut paths = self.paths.lock().unwrap();
        paths.retain(|p| p.strong_count() != 0);
        paths.push(Arc::downgrade(&path));
        path
    }
    fn moved(&self, from: &str, to: Option<&str>) {
        let mut paths = self.paths.lock().unwrap();
        paths.retain(|p| {
            let Some(p) = p.upgrade() else { return false };
            let mut path = p.lock().unwrap();
            if let Some(old) = path.as_deref() {
                if under(old, from) {
                    *path = to.map(|to| format!("{to}{}", &old[from.len()..]));
                } else if to.is_some_and(|to| under(old, to)) {
                    // A replaced destination's old descriptor no longer names a visible file.
                    *path = None;
                }
            }
            true
        });
    }
}
/// Component-aware descendant matching, including the path itself.
pub fn under(path: &str, base: &str) -> bool {
    base == "/" || path == base || path.strip_prefix(base).is_some_and(|s| s.starts_with('/'))
}
struct Observed {
    layer: Arc<dyn Layer>,
    prefix: String,
    observer: Observer,
    process: Process,
}
impl Observed {
    fn path(&self, path: &str) -> String {
        let path = format!("{}{}", self.prefix, path.trim_end_matches('/'));
        if path.is_empty() { "/".into() } else { path }
    }
}
struct Handle {
    file: Arc<dyn File>,
    path: Path,
    observer: Observer,
    process: Process,
}
impl Handle {
    async fn mutate<T: Send + 'static>(
        &self,
        f: impl Future<Output = Result<T>> + Send + 'static,
    ) -> Result<T> {
        let (path, observer) = (self.path.clone(), self.observer.clone());
        self.process
            .operation(async move {
                let result = f.await?;
                let path = path.lock().unwrap().clone();
                if let Some(path) = path {
                    (observer.changed)(&path);
                }
                Ok(result)
            })
            .await
    }
}
#[async_trait]
impl Layer for Observed {
    async fn stat(&self, path: &str, follow: bool) -> Result<Stat> {
        self.layer.stat(path, follow).await
    }
    async fn readlink(&self, path: &str) -> Result<String> {
        self.layer.readlink(path).await
    }
    async fn list(&self, path: &str) -> Result<Vec<Entry>> {
        self.layer.list(path).await
    }
    async fn open(&self, path: &str, options: Open) -> Result<Arc<dyn File>> {
        let (layer, local, path, observer, process) = (
            self.layer.clone(),
            path.to_owned(),
            self.path(path),
            self.observer.clone(),
            self.process.clone(),
        );
        self.process
            .operation(async move {
                let order = observer.order(&[(&path, false)]).await;
                let file = layer.open(&local, options).await?;
                let handle = Arc::new(Handle {
                    file,
                    path: observer.path(path.clone()),
                    observer: observer.clone(),
                    process,
                });
                drop(order);
                if options.create || options.truncate {
                    (observer.changed)(&path);
                }
                Ok(handle as Arc<dyn File>)
            })
            .await
    }
    async fn mkdir(&self, path: &str) -> Result<()> {
        let (layer, local, path, observer) = (
            self.layer.clone(),
            path.to_owned(),
            self.path(path),
            self.observer.clone(),
        );
        self.process
            .operation(async move {
                layer.mkdir(&local).await?;
                (observer.changed)(&path);
                Ok(())
            })
            .await
    }
    async fn remove_file(&self, path: &str) -> Result<()> {
        self.remove(path, false).await
    }
    async fn remove_dir(&self, path: &str) -> Result<()> {
        self.remove(path, true).await
    }
    async fn rename(&self, from: &str, to: &str) -> Result<()> {
        let (layer, a, b, from, to, observer) = (
            self.layer.clone(),
            from.to_owned(),
            to.to_owned(),
            self.path(from),
            self.path(to),
            self.observer.clone(),
        );
        self.process
            .operation(async move {
                let order = observer.order(&[(&from, true), (&to, true)]).await;
                layer.rename(&a, &b).await?;
                if from != to {
                    observer.moved(&from, Some(&to));
                }
                drop(order);
                (observer.changed)(&from);
                (observer.changed)(&to);
                Ok(())
            })
            .await
    }
}
impl Observed {
    async fn remove(&self, path: &str, directory: bool) -> Result<()> {
        let (layer, local, path, observer) = (
            self.layer.clone(),
            path.to_owned(),
            self.path(path),
            self.observer.clone(),
        );
        self.process
            .operation(async move {
                let order = observer.order(&[(&path, true)]).await;
                if directory {
                    layer.remove_dir(&local).await?;
                } else {
                    layer.remove_file(&local).await?;
                }
                observer.moved(&path, None);
                drop(order);
                (observer.changed)(&path);
                Ok(())
            })
            .await
    }
}
#[async_trait]
impl File for Handle {
    async fn stat(&self) -> Result<Stat> {
        self.file.stat().await
    }
    async fn read(&self, offset: u64, length: usize) -> Result<Vec<u8>> {
        self.file.read(offset, length).await
    }
    async fn write(&self, offset: u64, bytes: Vec<u8>) -> Result<usize> {
        let f = self.file.clone();
        self.mutate(async move { f.write(offset, bytes).await })
            .await
    }
    async fn append(&self, bytes: Vec<u8>) -> Result<usize> {
        let f = self.file.clone();
        self.mutate(async move { f.append(bytes).await }).await
    }
    async fn set_size(&self, size: u64) -> Result<()> {
        let f = self.file.clone();
        self.mutate(async move { f.set_size(size).await }).await
    }
    async fn sync(&self) -> Result<()> {
        self.file.sync().await
    }
    async fn set_times(
        &self,
        accessed: Option<SystemTime>,
        modified: Option<SystemTime>,
    ) -> Result<()> {
        let f = self.file.clone();
        self.mutate(async move { f.set_times(accessed, modified).await })
            .await
    }
}

#[cfg(test)]
mod tests {
    use super::*;
    use tokio::sync::mpsc;
    #[test]
    fn component_matching() {
        assert!(under("/mailbox/a", "/mailbox"));
        assert!(under("/mailbox", "/mailbox"));
        assert!(under("/anything", "/"));
        assert!(!under("/mailboxes/a", "/mailbox"));
    }
    /// Every visible mutation is delivered after the layer call, including rebased
    /// open handles in another process's namespace. Removed handles stay usable.
    #[tokio::test]
    async fn handles_follow_rename_and_unlink() {
        for host in [false, true] {
            let temp = tempfile::tempdir().unwrap();
            let process = Process::default();
            let layer: Arc<dyn Layer> = if host {
                Arc::new(Host::new(temp.path(), process.clone()).unwrap())
            } else {
                Arc::new(Memory::default())
            };
            let (tx, mut rx) = mpsc::unbounded_channel();
            let observer = Observer::new(move |p| {
                tx.send(p.to_owned()).unwrap();
            });
            let a = observer.layer(layer.clone(), "/run", process.clone());
            let b = observer.layer(layer, "/run", process);
            a.mkdir("/dir").await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), "/run/dir");
            let f = a
                .open(
                    "/dir/a",
                    Open {
                        create: true,
                        write: true,
                        ..Open::default()
                    },
                )
                .await
                .unwrap();
            assert_eq!(rx.recv().await.unwrap(), "/run/dir/a");
            b.rename("/dir", "/moved").await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), "/run/dir");
            assert_eq!(rx.recv().await.unwrap(), "/run/moved");
            f.write(0, b"new".to_vec()).await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), "/run/moved/a");
            assert_eq!(
                b.open("/moved/a", Open::default())
                    .await
                    .unwrap()
                    .read(0, 3)
                    .await
                    .unwrap(),
                b"new"
            );
            let old = b
                .open(
                    "/old",
                    Open {
                        create: true,
                        write: true,
                        ..Open::default()
                    },
                )
                .await
                .unwrap();
            rx.recv().await.unwrap();
            b.rename("/moved/a", "/old").await.unwrap();
            rx.recv().await.unwrap();
            rx.recv().await.unwrap();
            old.write(0, b"invisible".to_vec()).await.unwrap();
            assert!(rx.try_recv().is_err());
            f.append(b"!".to_vec()).await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), "/run/old");
            b.remove_file("/old").await.unwrap();
            assert_eq!(rx.recv().await.unwrap(), "/run/old");
            f.set_size(0).await.unwrap();
            assert!(rx.try_recv().is_err());
            assert!(b.remove_file("/missing").await.is_err());
            assert!(rx.try_recv().is_err());
        }
    }
    /// A pending exclusive namespace change must not block a sibling's opens,
    /// even though Tokio's RwLock prefers queued writers over new readers.
    #[tokio::test]
    async fn namespace_order_is_scoped_and_shared() {
        async fn ready<T>(future: impl Future<Output = T>) -> T {
            let mut future = std::pin::pin!(future);
            std::future::poll_fn(|cx| match future.as_mut().poll(cx) {
                std::task::Poll::Ready(value) => std::task::Poll::Ready(value),
                std::task::Poll::Pending => panic!("unrelated namespace operation blocked"),
            })
            .await
        }
        let observer = Observer::new(|_| {});
        let open = observer.order(&[("/run/slow/file", false)]).await;
        let concurrent_open = ready(observer.order(&[("/run/slow/file", false)])).await;
        let mut rename =
            std::pin::pin!(observer.order(&[("/run/slow", true), ("/run/moved", true)]));
        std::future::poll_fn(|cx| {
            assert!(rename.as_mut().poll(cx).is_pending());
            std::task::Poll::Ready(())
        })
        .await;
        // Same mounted layer, different subtree; also a different mounted layer.
        let sibling_open = ready(observer.order(&[("/run/sibling/file", false)])).await;
        let host_open = ready(observer.order(&[("/state/file", false)])).await;
        drop(sibling_open);
        let sibling_rename =
            ready(observer.order(&[("/run/sibling", true), ("/run/other", true)])).await;
        drop((sibling_rename, host_open, concurrent_open, open));
        let renamed = ready(rename).await;
        // A descendant open waits for its ancestor's rename until rebasing is done.
        let mut descendant = std::pin::pin!(observer.order(&[("/run/moved/file", false)]));
        std::future::poll_fn(|cx| {
            assert!(descendant.as_mut().poll(cx).is_pending());
            std::task::Poll::Ready(())
        })
        .await;
        drop(renamed);
        drop(ready(descendant).await);
        drop(
            observer
                .order(&[("/run/moved/file", true), ("/run/moved", true)])
                .await,
        );
        drop(
            observer
                .order(&[("/run/moved", true), ("/run/moved", true)])
                .await,
        );
        // Dead weak entries are pruned, rather than retaining every path ever opened.
        observer.order(&[("/", false)]).await;
        assert_eq!(observer.namespace.lock().unwrap().len(), 1);
    }

    /// Pause an open after obtaining its real identity, before registering its
    /// observed handle, like a slow host syscall or /self/log query.
    struct SlowOpen {
        memory: Arc<Memory>,
        gate: Mutex<
            Option<(
                tokio::sync::oneshot::Sender<()>,
                tokio::sync::oneshot::Receiver<()>,
            )>,
        >,
    }
    #[async_trait]
    impl Layer for SlowOpen {
        async fn stat(&self, p: &str, follow: bool) -> Result<Stat> {
            self.memory.stat(p, follow).await
        }
        async fn list(&self, p: &str) -> Result<Vec<Entry>> {
            self.memory.list(p).await
        }
        async fn open(&self, p: &str, options: Open) -> Result<Arc<dyn File>> {
            let f = self.memory.open(p, options).await?;
            let gate = self.gate.lock().unwrap().take();
            if let Some((started, release)) = gate {
                started.send(()).unwrap();
                release.await.unwrap();
            }
            Ok(f)
        }
        async fn mkdir(&self, p: &str) -> Result<()> {
            self.memory.mkdir(p).await
        }
        async fn remove_file(&self, p: &str) -> Result<()> {
            self.memory.remove_file(p).await
        }
        async fn remove_dir(&self, p: &str) -> Result<()> {
            self.memory.remove_dir(p).await
        }
        async fn rename(&self, a: &str, b: &str) -> Result<()> {
            self.memory.rename(a, b).await
        }
    }
    #[tokio::test]
    async fn blocked_io_does_not_stall_other_processes() {
        let memory = Arc::new(Memory::default());
        memory.mkdir("/slow").await.unwrap();
        let options = Open {
            create: true,
            write: true,
            ..Open::default()
        };
        let raw = memory.open("/slow/file", options).await.unwrap();
        let (started, opening) = tokio::sync::oneshot::channel();
        let (release, released) = tokio::sync::oneshot::channel();
        let (tx, mut rx) = mpsc::unbounded_channel();
        let observer = Observer::new(move |p| {
            tx.send(p.to_owned()).unwrap();
        });
        let slow = observer.layer(
            Arc::new(SlowOpen {
                memory: memory.clone(),
                gate: Mutex::new(Some((started, released))),
            }),
            "/run",
            Process::default(),
        );
        let other = observer.layer(memory, "/run", Process::default());
        let opening_task = tokio::spawn(async move {
            slow.open(
                "/slow/file",
                Open {
                    write: true,
                    ..Open::default()
                },
            )
            .await
        });
        opening.await.unwrap();
        let renamer = other.clone();
        let renamed = tokio::spawn(async move { renamer.rename("/slow", "/moved").await });
        // These must complete while the first process's open is still blocked,
        // and a conflicting rename is waiting to rebase its registered handle.
        other
            .open(
                "/",
                Open {
                    directory: true,
                    ..Open::default()
                },
            )
            .await
            .unwrap();
        let f = other.open("/other", options).await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), "/run/other");
        f.write(0, b"visible".to_vec()).await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), "/run/other");
        other.mkdir("/sibling").await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), "/run/sibling");
        other.rename("/other", "/sibling/file").await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), "/run/other");
        assert_eq!(rx.recv().await.unwrap(), "/run/sibling/file");
        other.remove_file("/sibling/file").await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), "/run/sibling/file");
        assert!(!renamed.is_finished());
        release.send(()).unwrap();
        let opened = opening_task.await.unwrap().unwrap();
        renamed.await.unwrap().unwrap();
        rx.recv().await.unwrap();
        rx.recv().await.unwrap();
        opened.set_size(0).await.unwrap();
        assert_eq!(rx.recv().await.unwrap(), "/run/moved/file");

        // No namespace guard surrounds descriptor IO: a rename completes while
        // this write is held in flight. Its notification uses the rebased path.
        let handle = Handle {
            file: raw.clone(),
            path: observer.path("/run/moved/file".into()),
            observer,
            process: Process::default(),
        };
        let (started, writing) = tokio::sync::oneshot::channel();
        let (release, released) = tokio::sync::oneshot::channel();
        let write = tokio::spawn(async move {
            handle
                .mutate(async move {
                    started.send(()).unwrap();
                    released.await.unwrap();
                    raw.write(0, b"after rename".to_vec()).await
                })
                .await
        });
        writing.await.unwrap();
        other.rename("/moved", "/again").await.unwrap();
        rx.recv().await.unwrap();
        rx.recv().await.unwrap();
        other.open("/again/file", Open::default()).await.unwrap();
        release.send(()).unwrap();
        write.await.unwrap().unwrap();
        assert_eq!(rx.recv().await.unwrap(), "/run/again/file");
    }

    /// Killing a waiting writer does not discard observation of a completed mutation;
    /// the process's idle barrier includes the callback, not only the system call.
    #[tokio::test]
    async fn cancellation_keeps_mutation_and_observation() {
        let (tx, mut rx) = mpsc::unbounded_channel();
        let observer = Observer::new(move |p| {
            tx.send(p.to_owned()).unwrap();
        });
        let process = Process::default();
        let layer = observer.layer(Arc::new(Memory::default()), "/", process.clone());
        let order = observer.order(&[("/published", true)]).await;
        let (started, ready) = tokio::sync::oneshot::channel();
        let p = process.clone();
        let task = tokio::spawn(async move {
            // Signal from inside the tracked operation, not before its guard exists.
            p.operation(async move {
                started.send(()).unwrap();
                layer
                    .open(
                        "/published",
                        Open {
                            create: true,
                            write: true,
                            ..Open::default()
                        },
                    )
                    .await
            })
            .await
        });
        ready.await.unwrap();
        task.abort();
        let _ = task.await;
        drop(order);
        process.idle().await;
        assert_eq!(rx.try_recv().unwrap(), "/published");
    }
}