/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use std::cell::RefCell;
use std::mem;
use std::path::PathBuf;
use std::str;
use std::sync::Arc;
use std::thread;
use std::thread::JoinHandle;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;
use cpython::*;
use cpython_ext::ExtractInner;
use cpython_ext::PyNone;
use cpython_ext::PyPath;
use cpython_ext::PyPathBuf;
use cpython_ext::ResultPyErrExt;
use cpython_ext::Str;
use crossbeam::channel::bounded;
use crossbeam::channel::Receiver;
use crossbeam::channel::Sender;
use pyrevisionstore::contentstore;
use pyrevisionstore::filescmstore;
use revisionstore::datastore::RemoteDataStore;
use revisionstore::localstore::LocalStore;
use revisionstore::redact_if_needed;
use revisionstore::HgIdDataStore;
use revisionstore::LegacyStore;
use revisionstore::StoreKey;
use revisionstore::StoreResult;
use types::HgId;
use types::Key;
use types::RepoPathBuf;
use vfs::UpdateFlag;
use vfs::VFS;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "worker"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<writerworker>(py)?;
    m.add_class::<removerworker>(py)?;
    Ok(m)
}

/// The `Worker` will batch up to `WORKER_BATCH_SIZE` work items before sending them to the
/// individual threads.
const WORKER_BATCH_SIZE: usize = 100;

struct Worker<Ret, Work> {
    threads: Vec<JoinHandle<Ret>>,
    sender: Sender<Vec<Work>>,
    pending: Vec<Work>,
}

impl<Ret: Send + 'static, Work: Sync + Send + 'static> Worker<Ret, Work> {
    /// Create a new worker.
    ///
    /// The internal channel size is bounded in size which makes the `push_work`
    /// method blocking. This is needed to prevent work from simply accumulating
    /// into it, having a bound allows for pushes to be blocking which allows for
    /// some good progress reporting to the user. The size of the channel was
    /// chosen to be twice the number of workers as a way to always keep the
    /// workers busy. As long as we can push one unit of work faster than the
    /// workers can process, the overall performance won't be bound to the Python
    /// threads.
    fn new<State: Send + Clone + 'static>(
        num_workers: usize,
        state: State,
        work: impl Fn(State, Receiver<Vec<Work>>) -> Ret + Send + 'static + Copy,
    ) -> Self {
        let (sender, receiver) = bounded(num_workers * 2);

        let mut threads = Vec::with_capacity(num_workers);
        for _ in 1..=num_workers {
            let chan = receiver.clone();
            let worker_state = state.clone();

            threads.push(thread::spawn(move || work(worker_state, chan)));
        }

        let pending = Vec::with_capacity(WORKER_BATCH_SIZE);

        Self {
            threads,
            sender,
            pending,
        }
    }

    /// Push work to the workers. The pushed work will be picked up by an idle
    /// worker, if the amount of pending work is already too high, this method will
    /// release the Python GIL and block.
    ///
    /// We need to batch work to avoid many pitfalls:
    ///  1) Sending data over the channel will involve locking/unlocking it,
    ///     and getting work from it will also involve similar locking. In order
    ///     to keep the workers busy, we need to minimize the amount of time that
    ///     all threads spend waiting for this lock.
    ///  2) Similarly, since the queue is bounded in size, we need to release
    ///     the GIL when pushing to it so that the other Python threads can run,
    ///     acquiring/releasing the GIL too often will also lead to reduced
    ///     performance.
    ///  3) Since pushing work to the worker is single threaded, the batched work
    ///     has a high chance of being on the same files, increasing locality.
    /// Caveat: the batch size was chosen semi-arbitrarily, and should be
    /// tweaked, too high of a value, and the progress reporting won't be
    /// good, too small of a value and the performance will suffer.
    fn push_work(&mut self, py: Python, work: Work) -> Result<()> {
        self.pending.push(work);

        if self.pending.len() == WORKER_BATCH_SIZE {
            // Release the GIL so other Python threads (the progress bar for instance) have a
            // chance to run while we're blocked in `send`.
            py.allow_threads(move || -> Result<()> {
                let pending =
                    mem::replace(&mut self.pending, Vec::with_capacity(WORKER_BATCH_SIZE));
                // This may block if the channel is full, see the comment in the
                // constructor for details.
                self.sender.send(pending)?;
                Ok(())
            })?;
        }
        Ok(())
    }

    /// Wait until all the previously pushed work have completed. Return an
    /// iterator over all the threads results.
    fn wait(mut self) -> Result<impl Iterator<Item = Ret>> {
        let pending = mem::take(&mut self.pending);
        self.sender.send(pending)?;
        drop(self.sender);

        let threads = mem::take(&mut self.threads);

        Ok(threads.into_iter().map(|thread| thread.join().unwrap()))
    }
}

/// Fetch the content of the passed in `key` and write it to it's path.
fn update(state: &WriterState, key: Key, flag: UpdateFlag) -> Result<usize> {
    let content = state
        .store
        .get_file_content(&key)?
        .ok_or_else(|| format_err!("Can't find key: {}", key))?;

    let meta = match state.store.get_meta(StoreKey::hgid(key.clone()))? {
        StoreResult::NotFound(key) => {
            return Err(format_err!("Can't find metadata for key: {:?}", key));
        }
        StoreResult::Found(meta) => meta,
    };

    if meta.is_lfs() {
        bail!("LFS pointers cannot be deserialized properly yet");
    }

    let content = redact_if_needed(content);

    state.working_copy.write(&key.path, &content, flag)
}

fn threaded_writer(
    state: WriterState,
    chan: Receiver<Vec<(Key, UpdateFlag)>>,
) -> (usize, Vec<(RepoPathBuf, UpdateFlag)>) {
    let mut failures = Vec::new();

    let mut written = 0;
    while let Ok(vec) = chan.recv() {
        let store_keys: Vec<_> = vec.iter().map(|(k, _)| StoreKey::hgid(k.clone())).collect();
        let missing = match state.store.get_missing(&store_keys) {
            Ok(missing) => missing,
            Err(e) => {
                tracing::warn!("{:?}", e);
                let failed_inputs: Vec<_> = vec.into_iter().map(|(k, f)| (k.path, f)).collect();
                failures.extend_from_slice(&failed_inputs);
                continue;
            }
        };
        if !missing.is_empty() {
            // Any errors will get reported below.
            let _ = state.store.prefetch(&missing);
        }

        for (key, flag) in vec.into_iter() {
            let res = update(&state, key.clone(), flag)
                .with_context(|| format!("Can't update {} at {}", key.path, key.hgid));

            match res {
                Ok(count) => written += count,
                Err(e) => {
                    tracing::warn!("{:?}", e);
                    failures.push((key.path, flag));
                }
            };
        }
    }

    (written, failures)
}

#[derive(Clone)]
struct WriterState {
    store: Arc<dyn LegacyStore>,
    working_copy: VFS,
}

impl WriterState {
    pub fn new(root: PathBuf, store: Arc<dyn LegacyStore>) -> Result<Self> {
        let working_copy = VFS::new(root)?;
        Ok(Self {
            store,
            working_copy,
        })
    }
}

py_class!(class writerworker |py| {
    data inner: RefCell<Option<Worker<(usize, Vec<(RepoPathBuf, UpdateFlag)>), (Key, UpdateFlag)>>>;

    def __new__(_cls, store: PyObject, root: &PyPath, numthreads: usize) -> PyResult<writerworker> {
        let store = contentstore::downcast_from(py, store.clone_ref(py)).map(|s| s.extract_inner(py) as Arc<dyn LegacyStore>)
            .or_else(|_| filescmstore::downcast_from(py, store).map(|s|  s.extract_inner(py) as Arc<dyn LegacyStore>))?;

        let root = root.to_path_buf();
        let writer_state = WriterState::new(root, store).map_pyerr(py)?;

        let inner = Worker::new(numthreads, writer_state, threaded_writer);

        writerworker::create_instance(py, RefCell::new(Some(inner)))
    }

    /// Issue an asynchronous write call. The request will be picked up by a
    /// background thread which will write the content corresponding to `node` to
    /// the file `name`. May block and release the GIL when too much work is
    /// pending.
    def write(&self, name: &PyPath, node: &PyBytes, flags: &str) -> PyResult<PyNone> {
        let path = name.to_repo_path_buf().map_pyerr(py)?;
        let node = HgId::from_slice(node.data(py)).map_pyerr(py)?;

        let flags = if flags == "l" {
            UpdateFlag::Symlink
        } else if flags == "x" {
            UpdateFlag::Executable
        } else if flags == "" {
            UpdateFlag::Regular
        } else {
            return Err(format_err!("Unknown flags: {}", flags)).map_pyerr(py);
        };

        self.inner(py).borrow_mut().as_mut().unwrap().push_work(py, (Key::new(path, node), flags)).map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Wait for all the pending `write` calls to complete.
    def wait(&self) -> PyResult<(usize, Vec<(PyPathBuf, Str)>)> {
        let inner = self.inner(py).borrow_mut().take().unwrap();

        py.allow_threads(move || -> Result<(usize, Vec<(PyPathBuf, Str)>)> {
            let mut written = 0;
            let mut failures = Vec::new();

            for (count, fail) in inner.wait()? {
                written += count;
                failures.extend(fail.into_iter().map(|(path, flag)| {
                    let path = PyPathBuf::from(path);

                    let flags = match flag {
                        UpdateFlag::Regular => Str::from("".to_string()),
                        UpdateFlag::Symlink => Str::from("l".to_string()),
                        UpdateFlag::Executable => Str::from("x".to_string()),
                    };

                    (path, flags)
                }));
            }

            Ok((written, failures))
        }).map_pyerr(py)
    }
});

fn threaded_remover(state: RemoverState, chan: Receiver<Vec<RepoPathBuf>>) -> Vec<RepoPathBuf> {
    let mut failures = Vec::new();

    while let Ok(vec) = chan.recv() {
        for path in vec.into_iter() {
            if let Err(e) = state.working_copy.remove(&path) {
                tracing::warn!("{:?}", e);
                failures.push(path);
            }
        }
    }

    failures
}

#[derive(Clone)]
struct RemoverState {
    working_copy: VFS,
}

impl RemoverState {
    pub fn new(root: PathBuf) -> Result<Self> {
        let working_copy = VFS::new(root)?;
        Ok(Self { working_copy })
    }
}

py_class!(class removerworker |py| {
    data inner: RefCell<Option<Worker<Vec<RepoPathBuf>, RepoPathBuf>>>;

    def __new__(_cls, root: &PyPath, numthreads: usize) -> PyResult<removerworker> {
        let root = root.to_path_buf();
        let remover_state = RemoverState::new(root).map_pyerr(py)?;

        let inner = Worker::new(numthreads, remover_state, threaded_remover);

        removerworker::create_instance(py, RefCell::new(Some(inner)))
    }

    /// Issue an asynchronous remove call. The request will be processed by a
    /// background thread. May block and release the GIL when too much work is pending.
    def remove(&self, name: &PyPath) -> PyResult<PyNone> {
        let path = name.to_repo_path_buf().map_pyerr(py)?;
        self.inner(py).borrow_mut().as_mut().unwrap().push_work(py, path).map_pyerr(py)?;
        Ok(PyNone)
    }

    /// Wait for all the pending `remove` calls to complete
    def wait(&self) -> PyResult<Vec<PyPathBuf>> {
        let inner = self.inner(py).borrow_mut().take().unwrap();

        let failures = py.allow_threads(move || -> Result<Vec<_>> {
            Ok(inner.wait()?.flatten().map(|path| PyPathBuf::from(path)).collect::<Vec<PyPathBuf>>())
        }).map_pyerr(py)?;

        Ok(failures)
    }
});

#[cfg(test)]
mod tests {
    use std::collections::HashSet;
    use std::fs::create_dir_all;
    use std::fs::metadata;
    use std::fs::read_dir;
    use std::fs::read_to_string;
    use std::fs::symlink_metadata;
    use std::fs::File;
    #[cfg(windows)]
    use std::fs::OpenOptions;
    use std::io::Write;
    #[cfg(unix)]
    use std::os::unix::fs::PermissionsExt;
    #[cfg(windows)]
    use std::os::windows::fs::OpenOptionsExt;

    use anyhow::ensure;
    use memmap2::MmapOptions;
    use minibytes::Bytes;
    use quickcheck::quickcheck;
    use quickcheck::TestResult;
    use revisionstore::datastore::Delta;
    use revisionstore::datastore::HgIdMutableDeltaStore;
    use revisionstore::testutil::make_config;
    use revisionstore::ContentStore;
    use tempfile::TempDir;
    use types::testutil::key;
    use types::RepoPath;

    use super::*;

    #[test]
    fn test_update_basic() -> Result<()> {
        let workingdir = TempDir::new()?;
        let cachedir = TempDir::new()?;
        let localdir = TempDir::new()?;
        let config = make_config(&cachedir);
        let store = Arc::new(ContentStore::new(&localdir, &config)?);

        let k = key("a", "2");
        let delta = Delta {
            data: Bytes::from("CONTENT"),
            base: None,
            key: k.clone(),
        };
        store.add(&delta, &Default::default())?;

        let root = workingdir.as_ref().to_path_buf();
        let state = WriterState::new(root, store)?;
        let written = update(&state, k, UpdateFlag::Regular)?;

        assert_eq!(written, 7);

        Ok(())
    }

    #[test]
    fn test_executable() -> Result<()> {
        let workingdir = TempDir::new()?;
        let cachedir = TempDir::new()?;
        let localdir = TempDir::new()?;
        let config = make_config(&cachedir);
        let store = Arc::new(ContentStore::new(&localdir, &config)?);

        let k = key("a", "2");
        let delta = Delta {
            data: Bytes::from("CONTENT"),
            base: None,
            key: k.clone(),
        };
        store.add(&delta, &Default::default())?;

        let root = workingdir.as_ref().to_path_buf();
        let state = WriterState::new(root, store)?;
        update(&state, k, UpdateFlag::Executable)?;

        let mut file = workingdir.as_ref().to_path_buf();
        file.push("a");

        let perms = metadata(&file)?.permissions();

        assert!(!perms.readonly());

        #[cfg(not(windows))]
        assert_eq!(perms.mode() & 0o755, 0o755);

        Ok(())
    }

    #[test]
    fn test_symlink() -> Result<()> {
        let workingdir = TempDir::new()?;
        let cachedir = TempDir::new()?;
        let localdir = TempDir::new()?;
        let config = make_config(&cachedir);
        let store = Arc::new(ContentStore::new(&localdir, &config)?);

        let k = key("a", "2");
        let delta = Delta {
            data: Bytes::from("b"),
            base: None,
            key: k.clone(),
        };
        store.add(&delta, &Default::default())?;

        let root = workingdir.as_ref().to_path_buf();
        let state = WriterState::new(root, store)?;
        update(&state, k, UpdateFlag::Symlink)?;

        let mut file = workingdir.as_ref().to_path_buf();
        file.push("a");

        let file_type = symlink_metadata(&file)?.file_type();

        if cfg!(windows) {
            assert!(file_type.is_file());
        } else {
            assert!(file_type.is_symlink());
        }

        Ok(())
    }

    #[test]
    fn test_update_nested() -> Result<()> {
        let workingdir = TempDir::new()?;
        let cachedir = TempDir::new()?;
        let localdir = TempDir::new()?;
        let config = make_config(&cachedir);
        let store = Arc::new(ContentStore::new(&localdir, &config)?);

        let k = key("a/b/c/d/e", "2");
        let delta = Delta {
            data: Bytes::from("CONTENT"),
            base: None,
            key: k.clone(),
        };
        store.add(&delta, &Default::default())?;

        let root = workingdir.as_ref().to_path_buf();
        let state = WriterState::new(root, store)?;
        let written = update(&state, k, UpdateFlag::Regular)?;
        assert_eq!(written, 7);

        Ok(())
    }

    #[test]
    fn test_update_replace_file_with_dir() -> Result<()> {
        let workingdir = TempDir::new()?;
        let cachedir = TempDir::new()?;
        let localdir = TempDir::new()?;
        let config = make_config(&cachedir);
        let store = Arc::new(ContentStore::new(&localdir, &config)?);

        let k = key("a/b/c/d/e", "2");
        let delta = Delta {
            data: Bytes::from("CONTENT"),
            base: None,
            key: k.clone(),
        };
        store.add(&delta, &Default::default())?;

        let mut path = workingdir.as_ref().to_path_buf();
        path.extend(&["a", "b"]);
        create_dir_all(&path)?;
        path.push("c");
        File::create(&path)?;

        let root = workingdir.as_ref().to_path_buf();
        let state = WriterState::new(root, store)?;
        let written = update(&state, k, UpdateFlag::Regular)?;
        assert_eq!(written, 7);

        let mut path = workingdir.as_ref().to_path_buf();
        path.extend(&["a", "b", "c", "d", "e"]);
        assert!(metadata(path)?.is_file());

        Ok(())
    }

    #[cfg(not(windows))]
    #[test]
    fn test_update_replace_symlink_with_dir() -> Result<()> {
        let workingdir = TempDir::new()?;
        let cachedir = TempDir::new()?;
        let localdir = TempDir::new()?;
        let config = make_config(&cachedir);
        let store = Arc::new(ContentStore::new(&localdir, &config)?);

        let k = key("a/b/c/d/e", "2");
        let delta = Delta {
            data: Bytes::from("CONTENT"),
            base: None,
            key: k.clone(),
        };
        store.add(&delta, &Default::default())?;

        let mut path = workingdir.as_ref().to_path_buf();
        path.extend(&["a", "b"]);
        create_dir_all(&path)?;
        path.push("c");

        std::os::unix::fs::symlink("foo", &path)?;

        let root = workingdir.as_ref().to_path_buf();
        let state = WriterState::new(root, store)?;
        update(&state, k, UpdateFlag::Regular)?;

        path.extend(&["d", "e"]);
        assert!(metadata(path)?.is_file());

        Ok(())
    }

    #[test]
    fn test_update_replace_dir_with_file() -> Result<()> {
        let workingdir = TempDir::new()?;
        let cachedir = TempDir::new()?;
        let localdir = TempDir::new()?;
        let config = make_config(&cachedir);
        let store = Arc::new(ContentStore::new(&localdir, &config)?);

        let k = key("a/b/c/d/e", "2");
        let delta = Delta {
            data: Bytes::from("CONTENT"),
            base: None,
            key: k.clone(),
        };
        store.add(&delta, &Default::default())?;

        let mut path = workingdir.as_ref().to_path_buf();
        path.extend(&["a", "b", "c", "d", "e", "f"]);
        create_dir_all(&path)?;

        let root = workingdir.as_ref().to_path_buf();
        let state = WriterState::new(root, store)?;
        let written = update(&state, k, UpdateFlag::Regular)?;
        assert_eq!(written, 7);

        let mut path = workingdir.as_ref().to_path_buf();
        path.extend(&["a", "b", "c", "d", "e"]);
        assert!(metadata(path)?.is_file());

        Ok(())
    }

    #[test]
    fn test_remove_basic() -> Result<()> {
        let workingdir = TempDir::new()?;

        let root = workingdir.as_ref().to_path_buf();
        let mut path = root.clone();
        path.push("TEST");

        File::create(&path)?;

        let state = RemoverState::new(root)?;
        state.working_copy.remove(RepoPath::from_str("TEST")?)?;

        assert_eq!(read_dir(&workingdir)?.count(), 0);

        Ok(())
    }

    #[test]
    fn test_remove_nested() -> Result<()> {
        let workingdir = TempDir::new()?;

        let root = workingdir.as_ref().to_path_buf();
        let mut path = root.clone();
        path.extend(&["THESE", "ARE", "DIRECTORIES"]);

        create_dir_all(&path)?;
        path.push("FILE");
        File::create(&path)?;

        let state = RemoverState::new(root)?;
        state
            .working_copy
            .remove(RepoPath::from_str("THESE/ARE/DIRECTORIES/FILE")?)?;
        assert_eq!(read_dir(&workingdir)?.count(), 0);

        Ok(())
    }

    #[test]
    fn test_remove_nested_partial() -> Result<()> {
        let workingdir = TempDir::new()?;

        let root = workingdir.as_ref().to_path_buf();
        let mut path = root.clone();
        path.extend(&["THESE", "ARE", "DIRECTORIES"]);

        create_dir_all(&path)?;
        path.push("FILE");
        File::create(&path)?;

        let mut path = root.clone();
        path.extend(&["OTHER"]);

        create_dir_all(&path)?;
        path.push("FILE");
        File::create(&path)?;

        let state = RemoverState::new(root)?;
        state
            .working_copy
            .remove(RepoPath::from_str("THESE/ARE/DIRECTORIES/FILE")?)?;
        assert_eq!(read_dir(&workingdir)?.count(), 1);

        Ok(())
    }

    #[test]
    fn test_remove_while_open() -> Result<()> {
        let workingdir = TempDir::new()?;

        let root = workingdir.as_ref().to_path_buf();
        let path = root.join("TEST");

        let f = File::create(&path)?;

        let state = RemoverState::new(root)?;
        state.working_copy.remove(RepoPath::from_str("TEST")?)?;

        drop(f);

        assert_eq!(read_dir(&workingdir)?.count(), 0);

        Ok(())
    }

    #[test]
    fn test_remove_while_mapped() -> Result<()> {
        let workingdir = TempDir::new()?;

        let root = workingdir.as_ref().to_path_buf();
        let path = root.join("TEST");

        File::create(&path)?.write_all(b"CONTENT")?;
        let f = File::open(&path)?;
        let map = unsafe { MmapOptions::new().map(&f)? };

        let state = RemoverState::new(root)?;
        state.working_copy.remove(RepoPath::from_str("TEST")?)?;

        drop(map);

        if cfg!(windows) {
            assert_eq!(read_dir(&workingdir)?.count(), 1);

            // The file must have been removed
            state.working_copy.remove(RepoPath::from_str("TEST")?)?;
            assert_eq!(read_dir(&workingdir)?.count(), 1);
        } else {
            assert_eq!(read_dir(&workingdir)?.count(), 0);
        }

        Ok(())
    }

    fn validate_paths<'a>(paths: impl Iterator<Item = &'a RepoPath>) -> bool {
        let mut files = HashSet::new();
        let mut directories = HashSet::new();

        for path in paths {
            // Keys for an empty path are meaningless.
            if path.as_str().trim() == "" {
                return false;
            }

            // We cannot have a file also be a directory.
            if directories.contains(path) {
                return false;
            }

            // Files have to be unique
            if files.contains(path) {
                return false;
            }

            // Make sure we do no have paths with directories and files with the same name.
            for parent in path.parents().skip(1) {
                if files.contains(parent) {
                    return false;
                }

                directories.insert(parent);
            }

            files.insert(path);
        }

        true
    }

    quickcheck! {
        fn update_only(keys: Vec<Key>) -> Result<TestResult> {
            let workingdir = TempDir::new()?;
            let cachedir = TempDir::new()?;
            let localdir = TempDir::new()?;
            let config = make_config(&cachedir);
            let store = Arc::new(ContentStore::new(&localdir, &config)?);

            if !validate_paths(keys.iter().map(|k| k.path.as_ref())) {
                return Ok(TestResult::discard());
            }

            let mut expected_size = 0;
            for key in keys.iter() {
                let data = Bytes::from(format!("{}", key));
                expected_size += data.len();
                let delta = Delta {
                    data,
                    base: None,
                    key: key.clone(),
                };

                store.add(&delta, &Default::default())?;
            }

            let root = workingdir.as_ref().to_path_buf();
            let state = WriterState::new(root, store)?;

            let mut written_size = 0;
            for key in keys.iter() {
                written_size += update(&state, key.clone(), UpdateFlag::Regular)?;
            }

            for key in keys.iter() {
                let mut fullpath = workingdir.as_ref().to_path_buf();
                fullpath.push(key.path.as_str());

                let ondisk = read_to_string(&fullpath)?;
                let expected = format!("{}", key);
                ensure!(ondisk == expected, format!("Got: {}, expected: {}", ondisk, expected));
            }

            Ok(TestResult::from_bool(expected_size == written_size))
        }

        fn remove_only(paths: Vec<RepoPathBuf>) -> Result<TestResult> {
            let workingdir = TempDir::new()?;

            if !validate_paths(paths.iter().map(|path| path.as_ref())) {
                return Ok(TestResult::discard());
            }

            for path in paths.iter() {
                let mut fullpath = workingdir.as_ref().to_path_buf();
                fullpath.push(path.as_str());

                create_dir_all(fullpath.parent().unwrap())?;
                File::create(&fullpath)?;
            }

            let root = workingdir.as_ref().to_path_buf();
            let state = RemoverState::new(root)?;
            for path in paths.iter() {
                state.working_copy.remove(&path)?;
            }

            Ok(TestResult::from_bool(read_dir(&workingdir)?.count() == 0))
        }

        fn update_then_remove(keys: Vec<Key>) -> Result<TestResult> {
            let workingdir = TempDir::new()?;
            let cachedir = TempDir::new()?;
            let localdir = TempDir::new()?;
            let config = make_config(&cachedir);
            let store = Arc::new(ContentStore::new(&localdir, &config)?);

            if !validate_paths(keys.iter().map(|k| k.path.as_ref())) {
                return Ok(TestResult::discard());
            }

            for key in keys.iter() {
                let delta = Delta {
                    data: Bytes::from(format!("{}", key)),
                    base: None,
                    key: key.clone(),
                };

                store.add(&delta, &Default::default())?;
            }

            let root = workingdir.as_ref().to_path_buf();
            let state = WriterState::new(root, store)?;

            for key in keys.iter() {
                update(&state, key.clone(), UpdateFlag::Regular)?;
            }

            let root = workingdir.as_ref().to_path_buf();
            let state = RemoverState::new(root)?;
            for key in keys.iter() {
                state.working_copy.remove(&key.path)?;
            }

            Ok(TestResult::from_bool(read_dir(&workingdir)?.count() == 0))
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

extern crate workingcopy as rsworkingcopy;

use std::cell::RefCell;
use std::sync::Arc;
use std::time::Duration;
use std::time::SystemTime;

use anyhow::anyhow;
use anyhow::Error;
use cpython::*;
use cpython_ext::error::ResultPyErrExt;
use cpython_ext::PyPathBuf;
use io::IO;
use parking_lot::RwLock;
use pathmatcher::Matcher;
use pyconfigloader::config;
use pypathmatcher::extract_matcher;
use pypathmatcher::extract_option_matcher;
use pytreestate::treestate;
use rsworkingcopy::walker::WalkError;
use rsworkingcopy::walker::Walker;
use rsworkingcopy::workingcopy::WorkingCopy;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "workingcopy"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<walker>(py)?;
    m.add_class::<workingcopy>(py)?;

    m.add(
        py,
        "parsegitsubmodules",
        py_fn!(py, parse_git_submodules(data: &PyBytes)),
    )?;

    Ok(m)
}

py_class!(class walker |py| {
    data inner: RefCell<Walker<Arc<dyn Matcher + Sync + Send>>>;
    data _errors: RefCell<Vec<Error>>;
    def __new__(
        _cls,
        root: PyPathBuf,
        dot_dir: String,
        pymatcher: PyObject,
        include_directories: bool,
    ) -> PyResult<walker> {
        let matcher = extract_matcher(py, pymatcher)?;
        let walker = Walker::new(
            root.to_path_buf(),
            dot_dir,
            Vec::new(),
            matcher,
            include_directories,
        ).map_pyerr(py)?;
        walker::create_instance(py, RefCell::new(walker), RefCell::new(Vec::new()))
    }

    def __iter__(&self) -> PyResult<Self> {
        Ok(self.clone_ref(py))
    }

    def __next__(&self) -> PyResult<Option<PyPathBuf>> {
        loop {
            match self.inner(py).borrow_mut().next() {
                Some(Ok(path)) => return Ok(Some(PyPathBuf::from(path.as_ref()))),
                Some(Err(e)) => self._errors(py).borrow_mut().push(e),
                None => return Ok(None),
            };
        }
    }

    def errors(&self) -> PyResult<Vec<(cpython_ext::Str, cpython_ext::Str)>> {
        Ok(self._errors(py).borrow().iter().map(|e| match e.downcast_ref::<WalkError>() {
            Some(e) => (e.filename().into(), e.message().into()),
            None => ("unknown".to_string().into(), e.to_string().into()),
        }).collect::<Vec<(cpython_ext::Str, cpython_ext::Str)>>())
    }

});

py_class!(pub class workingcopy |py| {
    data inner: Arc<RwLock<WorkingCopy>>;

    def treestate(&self) -> PyResult<treestate> {
        treestate::create_instance(py, self.inner(py).read().treestate())
    }

    def status(
        &self,
        pymatcher: Option<PyObject>,
        lastwrite: u32,
        config: &config,
    ) -> PyResult<PyObject> {
        let wc = self.inner(py).write();
        let matcher = extract_option_matcher(py, pymatcher)?;
        let last_write = SystemTime::UNIX_EPOCH.checked_add(
            Duration::from_secs(lastwrite.into())).ok_or_else(|| anyhow!("Failed to convert {} to SystemTime", lastwrite)
        ).map_pyerr(py)?;
        let io = IO::main().map_pyerr(py)?;
        let config = config.get_cfg(py);
        pystatus::to_python_status(py,
            &py.allow_threads(|| {
                wc.status(matcher, last_write, &config, &io)
            }).map_pyerr(py)?
        )
    }
});

fn parse_git_submodules(py: Python, data: &PyBytes) -> PyResult<Vec<(String, String, String)>> {
    Ok(rsworkingcopy::git::parse_submodules(data.data(py))
        .map_pyerr(py)?
        .into_iter()
        .map(|sm| (sm.name, sm.url, sm.path))
        .collect())
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::borrow::Cow;

use cpython::*;
use cpython_ext::convert::Serde;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "xdiff"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(py, "blocks", py_fn!(py, blocks(a: PyObject, b: PyObject)))?;
    m.add(
        py,
        "edit_cost",
        py_fn!(py, edit_cost(a: PyObject, b: PyObject, max_edit_cost: u64)),
    )?;
    Ok(m)
}

// (a: bytes | str, b: bytes | str) -> List[(a1, a2, b1, b2)].
// Yield matching blocks.
fn blocks(py: Python, a: PyObject, b: PyObject) -> PyResult<Serde<Vec<(u64, u64, u64, u64)>>> {
    let a = a.extract::<BytesOrStr>(py)?;
    let b = b.extract::<BytesOrStr>(py)?;
    let a_data = a.as_bytes(py)?;
    let b_data = b.as_bytes(py)?;
    let result = py.allow_threads(|| xdiff::blocks(&a_data, &b_data));
    Ok(Serde(result))
}

// (a: bytes | str, b: bytes | str, max_edit_cost) -> edit_cost
fn edit_cost(py: Python, a: PyObject, b: PyObject, max_edit_cost: u64) -> PyResult<u64> {
    let a = a.extract::<BytesOrStr>(py)?;
    let b = b.extract::<BytesOrStr>(py)?;
    let a_data = a.as_bytes(py)?;
    let b_data = b.as_bytes(py)?;
    let result = py.allow_threads(|| xdiff::edit_cost(&a_data, &b_data, max_edit_cost));
    Ok(result)
}

enum BytesOrStr {
    Bytes(PyBytes),
    Str(PyString),
}

impl<'s> FromPyObject<'s> for BytesOrStr {
    fn extract(py: Python, obj: &'s PyObject) -> PyResult<Self> {
        if let Ok(bytes_obj) = obj.extract::<PyBytes>(py) {
            Ok(Self::Bytes(bytes_obj))
        } else {
            let str_obj = obj.extract::<PyString>(py)?;
            Ok(Self::Str(str_obj))
        }
    }
}

impl BytesOrStr {
    fn as_bytes<'a>(&'a self, py: Python) -> PyResult<Cow<'a, [u8]>> {
        match self {
            Self::Bytes(obj) => Ok(Cow::Borrowed(obj.data(py))),
            Self::Str(obj) => {
                let data = obj.data(py).to_string(py)?;
                let data = match data {
                    Cow::Borrowed(data) => Cow::Borrowed(data.as_bytes()),
                    Cow::Owned(data) => Cow::Owned(data.into_bytes()),
                };
                Ok(data)
            }
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use std::cell::RefCell;
use std::io;

use cpython::*;
use cpython_ext::ResultPyErrExt;
use cpython_ext::SimplePyBuf;
use zstd::stream::decode_all;
use zstd::stream::encode_all;
use zstd::stream::raw::Decoder;
use zstd::stream::raw::InBuffer;
use zstd::stream::raw::Operation;
use zstd::stream::raw::OutBuffer;
use zstdelta::apply;
use zstdelta::diff;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "zstd"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add(
        py,
        "apply",
        py_fn!(py, apply_py(base: &PyObject, delta: &PyObject)),
    )?;
    m.add(
        py,
        "diff",
        py_fn!(py, diff_py(base: &PyObject, data: &PyObject)),
    )?;
    m.add(py, "decode_all", py_fn!(py, decode_all_py(data: &PyObject)))?;
    m.add(
        py,
        "encode_all",
        py_fn!(py, encode_all_py(data: &PyObject, level: i32)),
    )?;
    m.add_class::<zstream>(py)?;
    Ok(m)
}

/// Convert `io::Result<Vec<u8>>` to a `PyResult<PyBytes>`.
fn convert(py: Python, result: io::Result<Vec<u8>>) -> PyResult<PyBytes> {
    result
        .map_err(|e| PyErr::new::<exc::RuntimeError, _>(py, format!("{}", e)))
        .map(|buf| PyBytes::new(py, &buf))
}

fn diff_py(py: Python, base: &PyObject, data: &PyObject) -> PyResult<PyBytes> {
    let base = SimplePyBuf::new(py, base);
    let data = SimplePyBuf::new(py, data);
    convert(py, diff(base.as_ref(), data.as_ref()))
}

fn apply_py(py: Python, base: &PyObject, delta: &PyObject) -> PyResult<PyBytes> {
    let base = SimplePyBuf::new(py, base);
    let delta = SimplePyBuf::new(py, delta);
    convert(py, apply(base.as_ref(), delta.as_ref()))
}

fn decode_all_py(py: Python, data: &PyObject) -> PyResult<PyBytes> {
    let data = SimplePyBuf::new(py, data);
    convert(py, decode_all(io::Cursor::new(data.as_ref())))
}

fn encode_all_py(py: Python, data: &PyObject, level: i32) -> PyResult<PyBytes> {
    let data = SimplePyBuf::new(py, data);
    convert(py, encode_all(io::Cursor::new(data.as_ref()), level))
}

py_class!(pub class zstream |py| {
    data decoder: RefCell<Decoder<'static>>;

    def __new__(_cls) -> PyResult<zstream> {
        let decoder = Decoder::new().map_pyerr(py)?;

        zstream::create_instance(py, RefCell::new(decoder))
    }

    def decompress_buffer(&self, data: &PyObject) -> PyResult<PyBytes> {
        let data = SimplePyBuf::new(py, data);
        let mut decoder = self.decoder(py).borrow_mut();
        let mut src = InBuffer::around(data.as_ref());
        let mut dst = vec![0u8; zstd_safe::dstream_out_size()];
        let mut dst = OutBuffer::around(&mut dst);

        while src.pos < src.src.len() {
            decoder.run(&mut src, &mut dst).map_pyerr(py)?;
        }
        Ok(PyBytes::new(py, dst.as_slice()))
    }
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use std::cell::RefCell;

use ::zstore::Id20;
use ::zstore::Repair;
use ::zstore::Zstore;
use cpython::*;
use cpython_ext::ResultPyErrExt;
use cpython_ext::Str;

pub fn init_module(py: Python, package: &str) -> PyResult<PyModule> {
    let name = [package, "zstore"].join(".");
    let m = PyModule::new(py, &name)?;
    m.add_class::<zstore>(py)?;
    Ok(m)
}

py_class!(class zstore |py| {
    data store: RefCell<Zstore>;

    def __new__(_cls, path: String) -> PyResult<Self> {
        let log = Zstore::open(&path).map_pyerr(py)?;
        Self::create_instance(py, RefCell::new(log))
    }

    /// Lookup a blob by id. Return None if the id is unknown.
    def get(&self, id: PyBytes) -> PyResult<Option<PyBytes>> {
        let id = Id20::from_slice(id.data(py)).map_pyerr(py)?;
        let store = self.store(py).borrow();
        let data = store.get(id).map_pyerr(py)?;
        Ok(data.map(|data| PyBytes::new(py, &data)))
    }

    /// Insert a blob. Return its Id.
    def insert(&self, data: PyBytes, delta_bases: Vec<PyBytes> = Vec::new()) -> PyResult<PyBytes> {
        let mut store = self.store(py).borrow_mut();
        let delta_bases = delta_bases.into_iter()
            .map(|id| Id20::from_slice(id.data(py)))
            .collect::<Result<Vec<_>, _>>()
            .map_pyerr(py)?;
        let id = store.insert(data.data(py), &delta_bases).map_pyerr(py)?;
        Ok(PyBytes::new(py, id.as_ref()))
    }

    /// Test if the store contains an id.
    def contains(&self, id: PyBytes) -> PyResult<bool> {
        let id = Id20::from_slice(id.data(py)).map_pyerr(py)?;
        let store = self.store(py).borrow();
        store.contains(id).map_pyerr(py)
    }

    /// Write pending data to disk. Raise if race condition is detected.
    def flush(&self) -> PyResult<u64> {
        self.store(py).borrow_mut().flush().map_pyerr(py)
    }

    def __getitem__(&self, id: PyBytes) -> PyResult<Option<PyBytes>> {
        self.get(py, id)
    }

    def __contains__(&self, id: PyBytes) -> PyResult<bool> {
        self.contains(py, id)
    }

    @staticmethod
    def repair(path: &str) -> PyResult<Str> {
        py.allow_threads(|| Zstore::repair(path)).map_pyerr(py).map(Into::into)
    }
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! Lightweight utilities. To optimize compile time, do not put complex
//! functions or functions depending on other crates here.

use cpython::py_fn;
use cpython::PyModule;
use cpython::PyResult;
use cpython::Python;
use cpython_ext::PyNone;

/// Populate an existing module so it contains utilities.
pub(crate) fn populate_module(py: Python<'_>, module: &PyModule) -> PyResult<PyNone> {
    module.add(
        py,
        "sleep",
        py_fn!(py, sleep(seconds: f64, gil: bool = false)),
    )?;
    Ok(PyNone)
}

fn sleep(py: Python, seconds: f64, gil: bool) -> PyResult<PyNone> {
    let duration = std::time::Duration::from_micros((seconds * 1e6) as u64);
    if gil {
        std::thread::sleep(duration);
    } else {
        py.allow_threads(|| {
            std::thread::sleep(duration);
        });
    }
    Ok(PyNone)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

#![allow(non_camel_case_types)]

use cpython::PyModule;
use cpython::PyResult;
use cpython::Python;
use cpython_ext::PyNone;

mod functions;
mod modules;

/// Populate an existing empty module so it contains utilities.
pub fn populate_module(py: Python<'_>, module: &PyModule) -> PyResult<PyNone> {
    functions::populate_module(py, module)?;
    modules::populate_module(py, module)?;
    Ok(PyNone)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::PyModule;
use cpython::PyResult;
use cpython::Python;
use cpython_ext::PyNone;
use paste::paste;

macro_rules! add_modules {
    ( $py:ident, $m:ident, [ $( $name:ident, )* ] ) => {
        let name = $m.get($py, "__name__")?.extract::<String>($py)?;
        $(
            paste! {
                $m.add($py, stringify!($name), ::[<py $name>]::init_module($py, &name)?)?;
            }
         )*
    }
}

/// Populate an existing empty module so it contains utilities.
pub(crate) fn populate_module(py: Python<'_>, module: &PyModule) -> PyResult<PyNone> {
    let m = module;
    m.add(py, "__doc__", "Mercurial Rust Bindings")?;
    add_modules!(
        py,
        m,
        [
            // see update_modules.sh
            // [[[cog
            // import cog, glob, os
            // for path in sorted(glob.glob('modules/py*/TARGETS')):
            //     name = os.path.basename(os.path.dirname(path))
            //     cog.outl(f'{name[2:]},')
            // ]]]
            atexit,
            auth,
            blackbox,
            bytes,
            cats,
            checkout,
            clientinfo,
            cliparser,
            conchparser,
            configloader,
            copytrace,
            dag,
            diffhelpers,
            dirs,
            doctor,
            drawdag,
            eagerepo,
            edenapi,
            error,
            exchange,
            fail,
            fs,
            gitstore,
            hgmetrics,
            hgtime,
            identity,
            indexedlog,
            io,
            lock,
            lz4,
            manifest,
            metalog,
            mutationstore,
            nodeipc,
            nodemap,
            pathhistory,
            pathmatcher,
            pprint,
            process,
            progress,
            refencode,
            regex,
            renderdag,
            repo,
            revisionstore,
            revlogindex,
            sptui,
            status,
            threading,
            tracing,
            treestate,
            vlq,
            worker,
            workingcopy,
            xdiff,
            zstd,
            zstore,
            // [[[end]]]
        ]
    );

    Ok(PyNone)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

// conch_parser is separate from edenscmnative/bindings intentionally
// so it can be used standalone without coupling with the rest of
// edenscm logic.

use cpython::serde::to_py_object;
use cpython::*;
use third_party_conch_parser::lexer::Lexer;
use third_party_conch_parser::parse::DefaultParser;

// rustfmt turns "|py, m|" into "|py, m,|" and breaks compile
#[rustfmt::skip::macros(py_module_initializer)]
py_module_initializer!(conch_parser, initconch_parser, PyInit_conch_parser, |py, m| {
    m.add(py, "parse", py_fn!(py, parse(code: &str)))?;
    Ok(())
});

fn parse(py: Python, code: &str) -> PyResult<PyObject> {
    let lex = Lexer::new(code.chars());
    let mut parser = DefaultParser::new(lex);
    let mut commands = Vec::new();
    while let Some(command) = parser
        .complete_command()
        .map_err(|e| PyErr::new::<exc::ValueError, _>(py, e.to_string()))?
    {
        commands.push(command)
    }
    to_py_object(py, &commands)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use cpython::py_module_initializer;
use hgcommands::prepare_builtin_modules;

py_module_initializer!(bindings, initbindings, PyInit_bindings, |py, m| {
    m.add(py, "__doc__", "Bootstraps the hg python environment")?;
    prepare_builtin_modules(py, m)?;
    Ok(())
});


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

//! This is a little macOS specific utility that is intended to be installed setuid root
//! so that it can mount scratch volumes into a portion of the filesytem
//! owned by a non-privileged user.
//! It is intended to be used together with edenfs, but may also be
//! useful for non-virtualized repos as a way to move IO out of a recursive
//! watch.
use std::os::unix::fs::MetadataExt;
use std::path::Path;
use std::path::PathBuf;
use std::process::Command;
use std::process::Stdio;
use std::str;
use std::time::Duration;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use eden_apfs::*;
use once_cell::sync::Lazy;
#[cfg(target_os = "macos")]
use serde::*;
use structopt::StructOpt;

#[cfg(feature = "fb")]
mod facebook;

const MOUNT_APFS: &'static str = "/sbin/mount_apfs";
const MAX_ADDVOLUME_RETRY: u64 = 3;

static MOUNT: Lazy<SystemCommandImpl> = Lazy::new(|| SystemCommandImpl(PathBuf::from(MOUNT_PATH)));

#[derive(StructOpt, Debug)]
enum Opt {
    /// List APFS volumes
    #[structopt(name = "list")]
    List {
        #[structopt(long = "all")]
        all: bool,
    },

    /// List APFS volumes that are not mounted and not used by any of the active checkouts.
    /// The intent is that `all_checkouts` is produced by `edenfsctl list`.
    #[structopt(name = "list-stale-volumes")]
    ListStaleVolumes {
        all_checkouts: Vec<String>,
        #[structopt(long = "json")]
        json: bool,
    },

    /// Mount some space at the specified path.
    /// You must be the owner of the path.
    #[structopt(name = "mount")]
    Mount { mount_point: String },

    /// Unmount the eden space from a specific path.
    /// This will only allow unmounting volumes that were created
    /// by this utility.
    #[structopt(name = "unmount")]
    UnMount {
        /// The mounted path that you wish to unmount
        mount_point: String,
        /// Force the unmount, even if files are open and busy
        #[structopt(long = "force")]
        force: bool,
    },

    /// Unmount and delete a volume associated with a specific path.
    /// This will only allow deleting volumes that were created
    /// by this utility
    #[structopt(name = "delete")]
    Delete {
        /// The mounted path that you wish to unmount
        mount_point: String,
    },

    /// Unmount and delete all APFS volumes created by this utility
    #[structopt(name = "delete-all")]
    DeleteAll {
        #[structopt(long = "kill_dependent_processes")]
        kill_dependent_processes: bool,
    },

    /// Unmount and delete a volume.
    /// This will only allow deleting volumes that were created
    /// by this utility
    #[structopt(name = "delete-volume")]
    DeleteVolume {
        /// The volume that you wish to delete
        volume: String,
    },
}

#[cfg(target_os = "macos")]
#[derive(Deserialize)]
#[serde(rename_all = "PascalCase")]
struct PartitionInfo {
    parent_whole_disk: String,
}

fn kill_active_pids_in_mounts(mut mount_points: Vec<String>) -> Result<()> {
    if mount_points.is_empty() {
        println!("Not killing anything as there are no mounts");
        return Ok(());
    }

    mount_points.push("-t".to_owned());
    // -t has to come before the mount point list
    let last_index = mount_points.len() - 1;
    mount_points.swap(0, last_index);
    let lsof: &str = "/usr/sbin/lsof";
    println!(
        "Listing dependent processes with: `{} {}`",
        lsof,
        mount_points.join(" ")
    );
    let mut active_pids = new_cmd_with_best_available_privs(lsof)
        .args(&mount_points)
        .stdout(Stdio::piped())
        .spawn()?;
    let active_pids_output = (active_pids.stdout.take()).context("lsof stdout not available")?;

    let xargs = "/usr/bin/xargs";
    let kill_args: Vec<&str> = vec!["-t", "/bin/kill", "-9"];
    println!(
        "and then killing them with : `{} {}`",
        xargs,
        kill_args.join(" ")
    );

    let output = new_cmd_with_best_available_privs(xargs)
        .args(&kill_args)
        .stdin(active_pids_output)
        .output()?;
    if !output.status.success() {
        return Err(anyhow!(
            "failed to execute lsof {} | xargs kill -9 \n {:#?}",
            mount_points.join(" "),
            output
        ));
    }
    println!("result: {:?}", output);
    Ok(())
}

fn list_mount_points(containers: &Vec<ApfsContainer>, mounts: &MountTable) -> Result<Vec<String>> {
    let mut mount_points: Vec<String> = vec![];

    for container in containers {
        for vol in &container.volumes {
            if vol.is_edenfs_managed_volume() {
                if let Some(mount_point) = vol.get_current_mount_point(&*MOUNT, Some(mounts)) {
                    mount_points.push(mount_point);
                }
            }
        }
    }

    Ok(mount_points)
}

/// Prepare a command to be run with root privs.
/// The path must be absolute to avoid being fooled into running something
/// unexpected.
/// The caller must already have root privs, otherwise this will fail.
fn new_cmd_with_root_privs(path: &str) -> Command {
    let path: PathBuf = path.into();
    assert!(path.is_absolute());
    assert!(
        geteuid() == 0,
        "root privs are required to run {}",
        path.display()
    );
    Command::new(path)
}

/// Prepare a command that will be run as root if we have root privileges, and
/// unprivileged if we do not. Use sparingly, running as root or non root is
/// better if it suits your needs.
/// This should be used for commands that can run both privileged and
/// unprivileged. Where privileged where privilege may give them more
/// capabilities, but they are best effort any ways, so we still want to try
/// to run them when we don't have privilege.
fn new_cmd_with_best_available_privs(path: &str) -> Command {
    let path: PathBuf = path.into();
    assert!(path.is_absolute());
    Command::new(path)
}

/// Create a new subvolume with the specified name.
/// Note that this does NOT require any special privilege on macOS.
///
/// Note that this code tries to create the subvolume multiple times to workaround a bug where the
/// `diskutil apfs addVolume` command succeeds but the subvolume isn't created. Apple claims that
/// this is fixed in macOS 11.5 but Sandcastle isn't on 11.5 yet.
fn make_new_volume(apfs_util: &ApfsUtil, name: &str, disk: &str) -> Result<ApfsVolume> {
    let mut tried = 0;
    loop {
        let output = new_cmd_unprivileged(DISKUTIL_PATH)
            .args(&["apfs", "addVolume", disk, "apfs", name, "-nomount"])
            .output()?;
        if !output.status.success() {
            anyhow::bail!("failed to execute diskutil addVolume: {:?}", output);
        }
        let containers = apfs_util.list_containers()?;

        if let Some(volume) = find_existing_volume(&containers, name) {
            return Ok(volume.clone());
        } else {
            tried += 1;
            if tried == MAX_ADDVOLUME_RETRY {
                return Err(anyhow!("failed to create volume `{}`: {:#?}", name, output));
            } else {
                println!(
                    "APFS subvolume created, but not found in `diskutil apfs list`, retrying."
                );
                // Let's sleep a bit before retrying in case this is timing related.
                std::thread::sleep(Duration::from_secs(1));

                // Maybe the volume wasn't available immediately, let's see if it appeared.
                let containers = apfs_util.list_containers()?;
                if let Some(volume) = find_existing_volume(&containers, name) {
                    return Ok(volume.clone());
                }

                // Nope, let's just loop
            }
        }
    }
}

fn get_real_uid() -> Result<u32> {
    let uid = getuid();

    if uid != 0 {
        return Ok(uid);
    }

    // We're really root (not just setuid root).  We may actually be
    // running under sudo so let's see what sudo says about the UID
    match std::env::var("SUDO_UID") {
        Ok(uid) => Ok(uid.parse().context(format!(
            "parsing the SUDO_UID={} env var as an integer",
            uid
        ))?),
        Err(std::env::VarError::NotPresent) => Ok(uid),
        Err(std::env::VarError::NotUnicode(_)) => bail!("the SUDO_UID env var is not unicode"),
    }
}

/// Canonicalize a path and return the canonical path in string form.
#[cfg(target_os = "macos")]
fn find_disk_for_eden_mount(mount_point: &str) -> Result<String> {
    let mut client_link = PathBuf::from(mount_point);
    client_link.push(".eden");
    client_link.push("client");

    let mut stat: libc::statfs = unsafe { std::mem::zeroed() };

    let client_link_cstr = std::ffi::CString::new(
        client_link
            .to_str()
            .ok_or_else(|| anyhow!("not a valid UTF-8 path somehow"))?,
    )?;
    let rv = unsafe { libc::statfs(client_link_cstr.as_ptr(), &mut stat) };
    if -1 == rv {
        return Err(std::io::Error::last_os_error().into());
    }

    let fstype = unsafe { std::ffi::CStr::from_ptr(stat.f_fstypename.as_ptr()).to_str()? };
    if "apfs" != fstype {
        bail!("disk at {} must be apfs", mount_point);
    }
    let partition = unsafe { std::ffi::CStr::from_ptr(stat.f_mntfromname.as_ptr()).to_str()? };
    let output = new_cmd_unprivileged(DISKUTIL_PATH)
        .args(&["info", "-plist", &partition])
        .output()?;
    if !output.status.success() {
        anyhow::bail!("failed to execute diskutil info: {:?}", output);
    }

    Ok(parse_plist::<PartitionInfo>(&String::from_utf8(output.stdout)?)?.parent_whole_disk)
}

#[cfg(not(target_os = "macos"))]
fn find_disk_for_eden_mount(_mount_point: &str) -> Result<String> {
    Err(anyhow!("only supported on macOS"))
}

fn mount_scratch_space_on(apfs_util: &ApfsUtil, input_mount_point: &str) -> Result<()> {
    let mount_point = canonicalize_mount_point_path(input_mount_point)?;
    println!("want to mount at {:?}", mount_point);

    // First, let's ensure that mounting at this location makes sense.
    // Inspect the directory and ensure that it is owned by us.
    let metadata = std::fs::metadata(&mount_point)
        .context(format!("Obtaining filesystem metadata for {}", mount_point))?;
    let my_uid = get_real_uid()?;
    if metadata.uid() != my_uid {
        bail!(
            "Refusing to set up a volume for {} because the owned uid {} doesn't match your uid {}",
            mount_point,
            metadata.uid(),
            my_uid
        );
    }

    println!("my real uid is {}, effective is {}", my_uid, unsafe {
        libc::geteuid()
    });

    let containers = apfs_util.list_containers()?;
    let name = encode_mount_point_as_volume_name(&mount_point);
    let volume = match find_existing_volume(&containers, &name) {
        Some(existing) => {
            let mount_table = MountTable::parse_system_mount_table(&*MOUNT)?;
            if let Some(current_mount_point) =
                existing.get_current_mount_point(&*MOUNT, Some(&mount_table))
            {
                if !existing.is_preferred_location(&current_mount_point)? {
                    // macOS will automatically mount volumes at system boot,
                    // but mount them under /Volumes.  That will block our attempt
                    // to mount the scratch space below, so if we see that this
                    // volume is mounted and not where we want it, we simply unmount
                    // it here now: this should be fine because we own these volumes
                    // and where they get mounted.  No one else should have a legit
                    // reason for mounting it elsewhere.
                    apfs_util.unmount_scratch(&mount_point, true, &mount_table)?;
                }
            }
            existing.clone()
        }
        None => make_new_volume(apfs_util, &name, &find_disk_for_eden_mount(&mount_point)?)?,
    };

    // Mount the volume at the desired mount point.
    // This is the only part of this utility that requires root privs.
    let output = new_cmd_with_root_privs(MOUNT_APFS)
        .args(&[
            "-onobrowse,nodev,nosuid",
            "-u",
            &format!("{}", metadata.uid()),
            "-g",
            &format!("{}", metadata.gid()),
            &format!("/dev/{}", volume.device_identifier),
            &mount_point,
        ])
        .output()?;
    if !output.status.success() {
        // See [`crate::facebook::write_apfs_issue_marker`] for detail
        #[cfg(feature = "fb")]
        if let Some(code) = output.status.code() {
            // This is the error code we get when we failed to mount the APFS subvolumes.
            if code == 66 {
                crate::facebook::write_apfs_issue_marker();
            }
        }
        anyhow::bail!(
            "failed to execute mount_apfs /dev/{} {}: {:#?}",
            volume.device_identifier,
            mount_point,
            output
        );
    }
    println!("output: {:?}", output);

    // Make sure that we own the mounted directory; the default is mounted
    // with root:wheel ownership, and that isn't desirable
    chown(&mount_point, metadata.uid(), metadata.gid())?;

    disable_spotlight(&mount_point).ok();
    disable_fsevents(&mount_point).ok();
    disable_trashcan(&mount_point).ok();

    Ok(())
}

fn chown(path: &str, uid: u32, gid: u32) -> Result<()> {
    let cstr = std::ffi::CString::new(path)
        .with_context(|| format!("creating a C string from path `{}`", path))?;
    let rc = unsafe { libc::chown(cstr.as_ptr(), uid, gid) };
    if rc != 0 {
        let err = std::io::Error::last_os_error();
        Err(err).with_context(|| format!("failed to chown {} to uid={}, gid={}", path, uid, gid))
    } else {
        Ok(())
    }
}

/// Don't bother indexing an artifact dir.  It's just a waste of resources
/// to build an index for something managed entirely by the machine.
fn disable_spotlight(mount_point: &str) -> Result<()> {
    let output = new_cmd_with_root_privs("/usr/bin/mdutil")
        .args(&["-Ed", "-i", "off", mount_point])
        .output()?;
    if !output.status.success() {
        eprintln!(
            "failed to disable spotlight on {}: {:#?}",
            mount_point, output
        );
    }

    let spotlight = Path::new(mount_point).join(".Spotlight-V100");
    std::fs::remove_dir_all(&spotlight).ok();

    Ok(())
}

/// Disable fsevents logging for the artifact dirs: this is for performance
/// reasons; we don't need/want fseventsd to run here.
fn disable_fsevents(mount_point: &str) -> Result<()> {
    // See https://developer.apple.com/library/archive/documentation/Darwin/Conceptual/FSEvents_ProgGuide/FileSystemEventSecurity/FileSystemEventSecurity.html#//apple_ref/doc/uid/TP40005289-CH6-SW5
    // Those docs say that we should recreate the directory and touch a control
    // file to disable logging data, but the presence of the directory can
    // confuse some tools, so we simply delete it at mount time; that should
    // be good enough in most cases.

    let fseventsd = Path::new(mount_point).join(".fseventsd");
    std::fs::remove_dir_all(&fseventsd).ok();

    Ok(())
}

/// The .Trashes directory has root permissions by default, which makes it
/// awkward for users to clean up the contents of the mount point when it
/// is used in place of an artifact directory.
fn disable_trashcan(mount_point: &str) -> Result<()> {
    let trashes = Path::new(mount_point).join(".Trashes");
    std::fs::remove_dir_all(trashes)?;

    // There's some thought that touching a regular file named
    // `.Trashes` is a good idea to prevent the trash dir from
    // coming back (which is what happens when using Finder to
    // send something to the trash).
    // For now we're just removing it at mount time.

    Ok(())
}

fn main() -> Result<()> {
    let opts = Opt::from_args();

    let apfs_util = ApfsUtil::new(DISKUTIL_PATH, MOUNT_PATH);

    match opts {
        Opt::List { all } => {
            let containers = apfs_util.list_containers()?;
            let mounts = MountTable::parse_system_mount_table(&*MOUNT)?;
            for container in containers {
                for vol in container.volumes {
                    if all || vol.is_edenfs_managed_volume() {
                        let name = vol.name.as_ref().map(String::as_str).unwrap_or("");
                        if let Some(mount_point) =
                            vol.get_current_mount_point(&*MOUNT, Some(&mounts))
                        {
                            println!("{}\t{}\t{}", vol.device_identifier, name, mount_point);
                        } else {
                            println!("{}\t{}", vol.device_identifier, name);
                        }
                    }
                }
            }
            Ok(())
        }

        Opt::ListStaleVolumes {
            all_checkouts,
            json,
        } => {
            let all_checkouts = all_checkouts
                .iter()
                .map(|v| canonicalize_mount_point_path(v.as_ref()))
                .collect::<Result<Vec<_>>>()?;

            let mut stale_volume_names = vec![];
            for vol in apfs_util.list_stale_volumes_unsafe(&all_checkouts)? {
                stale_volume_names.push(vol.name.context("Volume has no name")?);
            }
            if json {
                println!("{}", serde_json::to_string(&stale_volume_names)?);
            } else {
                for name in stale_volume_names.iter() {
                    println!("{}", name);
                }
            }
            Ok(())
        }

        Opt::Mount { mount_point } => mount_scratch_space_on(&apfs_util, &mount_point),

        Opt::UnMount { mount_point, force } => {
            apfs_util.unmount_scratch(
                &mount_point,
                force,
                &MountTable::parse_system_mount_table(&*MOUNT)?,
            )?;
            Ok(())
        }

        Opt::Delete { mount_point } => {
            apfs_util.delete_scratch(&mount_point)?;
            Ok(())
        }

        Opt::DeleteVolume { volume } => {
            apfs_util.delete_volume(&volume)?;
            Ok(())
        }

        Opt::DeleteAll {
            kill_dependent_processes,
        } => {
            let containers = apfs_util.list_containers()?;
            let mounts = MountTable::parse_system_mount_table(&*MOUNT)?;

            if kill_dependent_processes {
                let mount_points: Vec<String> = list_mount_points(&containers, &mounts)?;
                kill_active_pids_in_mounts(mount_points)?;
            }

            let mut was_failure = false;
            for container in containers {
                for vol in container.volumes {
                    if vol.is_edenfs_managed_volume() {
                        let mut try_delete = true;

                        if let Some(mount_point) =
                            vol.get_current_mount_point(&*MOUNT, Some(&mounts))
                        {
                            // In the context of deleting all volumes, we want to
                            // force the unmount--we know it is safe.
                            let force = true;
                            if let Err(err) =
                                apfs_util.unmount_scratch(&mount_point, force, &mounts)
                            {
                                eprintln!("Failed to unmount: {}", err);
                                try_delete = false;
                                was_failure = true;
                            }
                        }

                        if try_delete {
                            let mount_point = vol.preferred_mount_point().unwrap();
                            if let Err(err) = apfs_util.delete_scratch(&mount_point) {
                                eprintln!("Failed to delete {:#?}: {}", vol, err);
                                was_failure = true
                            } else {
                                println!("Deleted {}", mount_point);
                            }
                        }
                    }
                }
            }
            if was_failure {
                Err(anyhow!("Failed to unmount or delete one or more volumes."))
            } else {
                Ok(())
            }
        }
    }
}