// Copyright 2019 The Fuchsia Authors. All rights reserved.
// Use of this source code is governed by a BSD-style license that can be
// found in the LICENSE file.

//! Generate a random directory tree that can be commited to disk all at once. The general formula
//! goes like this -
//!
//! ```
//! use {rand::{Rng, SeedableRng, rngs::StdRng}, static_tree::{EntryDistribution, DirectoryEntry}};
//! let mut rng = StdRng::seed_from_u64(seed);
//! let dist = EntryDistribution::new(depth);
//! let tree: DirectoryEntry = rng.sample(&dist);
//! tree.write_tree_at(root).expect("failed to write tree");
//! ```

use {
    anyhow::Error,
    fidl_fuchsia_io as fio,
    rand::{distributions, Rng},
};

/// A random distribution specialized to generation of random directory trees. This distribution
/// decreases the likelyhood of a directory being generated linearly relative to the depth of the
/// subset of the tree being generated, until it's 0% at the maximum depth.
#[derive(Clone, Copy, Debug, Eq, PartialEq)]
pub struct EntryDistribution {
    depth: u32,
    max_depth: u32,
}

impl EntryDistribution {
    /// Create a new `EntryDistribution` with a maximum depth of `max_depth`. This distribution is
    /// used for generating [`DirectoryEntry`]s and [`Entry`]s.
    pub fn new(max_depth: u32) -> EntryDistribution {
        EntryDistribution { depth: 1, max_depth }
    }

    /// get the entry distribution for the next level down on the directory tree.
    fn next_level(&self) -> EntryDistribution {
        debug_assert!(
            self.depth <= self.max_depth,
            "directory tree exceeded max depth ({} vs max of {}). programming error.",
            self.depth,
            self.max_depth
        );
        EntryDistribution { depth: self.depth + 1, max_depth: self.max_depth }
    }

    /// creates a bernoulli distribution where the likelyhood is progressively decreased at further
    /// depths until it's 0% at max_depth.
    fn directory_distribution(&self) -> distributions::Bernoulli {
        distributions::Bernoulli::from_ratio(self.max_depth - self.depth, self.max_depth).unwrap()
    }
}

/// A file entry in the generated tree. Contains a randomly generated u64 for a name and randomly
/// generated byte contents. The file size is random, in a range of 1 byte to 2^16 bytes (~65kb).
#[derive(Debug, Clone, Eq, PartialEq)]
pub struct FileEntry {
    name: u64,
    contents: Vec<u8>,
}

impl distributions::Distribution<FileEntry> for distributions::Standard {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> FileEntry {
        let size = rng.gen_range(1..1 << 16);
        // unfortunately, we can't use sample_iter to generate the content here. the trait
        // definition for distribution requires the Rng type parameter to have ?Sized (or "maybe
        // sized"), but the sample_iter function requires the provided Rng be Sized, either through
        // the explicit Self: Sized on the Rng sample_iter function, or the implicit Sized present
        // on type parameters for the equivalent Distribution function.
        let mut contents = vec![0; size];
        rng.fill(contents.as_mut_slice());
        FileEntry { name: rng.gen(), contents }
    }
}

impl FileEntry {
    async fn write_file_at(self, root: &fio::DirectoryProxy) -> Result<(), Error> {
        let file = fuchsia_fs::directory::open_file(
            root,
            &self.name.to_string(),
            fio::OpenFlags::CREATE | fio::OpenFlags::RIGHT_WRITABLE,
        )
        .await?;
        fuchsia_fs::file::write(&file, &self.contents).await?;
        Ok(())
    }
}

/// A directory entry in the generated tree. Contains a randomly generated u64 for a name and a
/// vector of randomly generated directory entries, with a number of entries somewhere in the range
/// of [0, 6). The sample function generates the entire subtree depth-first before returning, and is
/// mutually recursive with the [`Entry`] sample function.
#[derive(Clone, Debug, Eq, PartialEq)]
pub struct DirectoryEntry {
    name: u64,
    entries: Vec<Entry>,
}

impl distributions::Distribution<DirectoryEntry> for EntryDistribution {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> DirectoryEntry {
        // each directory has a random number of entries in the range [0, 6)
        let num_entries = rng.gen_range(0..6);
        let mut entries = vec![];
        let entry_dist = self.next_level();
        for _ in 0..num_entries {
            entries.push(rng.sample(&entry_dist));
        }
        DirectoryEntry { name: rng.gen(), entries }
    }
}

impl DirectoryEntry {
    /// get the path the directory entry will be written to relative to a given root.
    pub fn get_name(&self) -> String {
        self.name.to_string()
    }

    /// Take the entire randomly generated tree and commit it to disk.
    pub fn write_tree_at<'a>(
        self,
        root: &'a fio::DirectoryProxy,
    ) -> futures::future::BoxFuture<'a, Result<(), Error>> {
        Box::pin(async {
            let this = fuchsia_fs::directory::create_directory(
                root,
                &self.get_name(),
                fio::OpenFlags::RIGHT_READABLE | fio::OpenFlags::RIGHT_WRITABLE,
            )
            .await?;
            for entry in self.entries {
                match entry {
                    Entry::File(file_entry) => file_entry.write_file_at(&this).await?,
                    Entry::Directory(dir_entry) => dir_entry.write_tree_at(&this).await?,
                }
            }
            Ok(())
        })
    }
}

/// An entry of either File or Directory. The chance of it being a directory is relative to the depth
/// recorded in EntryDistribution. The associated entry is also then generated.
#[derive(Clone, Debug, Eq, PartialEq)]
pub enum Entry {
    /// This entry is a file.
    File(FileEntry),
    /// This entry is a directory
    Directory(DirectoryEntry),
}

impl distributions::Distribution<Entry> for EntryDistribution {
    fn sample<R: Rng + ?Sized>(&self, rng: &mut R) -> Entry {
        if rng.sample(self.directory_distribution()) {
            Entry::Directory(rng.sample(self))
        } else {
            Entry::File(rng.sample(distributions::Standard))
        }
    }
}

#[cfg(test)]
mod tests {
    use {
        super::{DirectoryEntry, Entry, EntryDistribution, FileEntry},
        fs_management::Minfs,
        fuchsia_async as fasync,
        ramdevice_client::RamdiskClient,
        rand::{rngs::mock::StepRng, Rng as _},
    };

    // this fixture will have to get updated any time the generation logic changes. depending on the
    // nature of the change, the comparison logic may have to be changed as well.
    fn get_fixture(initial: u64, increment: u64, depth: u32) -> DirectoryEntry {
        // confirm that the rng and depth used to generate this tree are the same.
        assert_eq!(initial, 0xdeadbeef, "initial value changed - update fixture");
        assert_eq!(increment, 1, "increment value changed - update fixture");
        assert_eq!(depth, 2, "depth value changed - update fixture");
        DirectoryEntry {
            name: 3735928580,
            entries: vec![
                Entry::File(FileEntry { name: 3735928563, contents: vec![242] }),
                Entry::File(FileEntry { name: 3735928567, contents: vec![246] }),
                Entry::File(FileEntry { name: 3735928571, contents: vec![250] }),
                Entry::File(FileEntry { name: 3735928575, contents: vec![254] }),
                Entry::File(FileEntry { name: 3735928579, contents: vec![2] }),
            ],
        }
    }

    #[test]
    fn gen_tree() {
        let initial = 0xdeadbeef;
        let increment = 1;
        let depth = 2;
        // make sure we are generating a tree at all.
        let mut rng = StepRng::new(initial, increment);
        let dist = EntryDistribution::new(depth);
        let tree: DirectoryEntry = rng.sample(dist);
        // this skips the contents for the files, because they are huge, so we do our own manual
        // comparison.
        let fixture = get_fixture(initial, increment, depth);
        assert_eq!(fixture, tree);
        for i in 0..fixture.entries.len() {
            match &fixture.entries[i] {
                Entry::File(fixture_file_entry) => match &tree.entries[i] {
                    Entry::File(tree_file_entry) => {
                        assert_eq!(fixture_file_entry.name, tree_file_entry.name)
                    }
                    _ => panic!("expected a file in generated tree"),
                },
                _ => panic!("expected a file in fixture tree"),
            }
        }
    }

    #[test]
    fn same_rng_same_tree() {
        // this confirms an assumption about the properties of tree generation that we rely on for
        // consistency checking.
        let dist = EntryDistribution::new(5);

        let tree1: DirectoryEntry = StepRng::new(1337, 1).sample(dist);
        let tree2: DirectoryEntry = StepRng::new(1337, 1).sample(dist);

        assert_eq!(tree1, tree2);
    }

    #[fasync::run_singlethreaded(test)]
    async fn write_tree() {
        let root = "/test-root";
        let initial = 0xdeadbeef;
        let increment = 1;
        let depth = 2;

        let mut rng = StepRng::new(initial, increment);
        let dist = EntryDistribution::new(depth);
        let tree: DirectoryEntry = rng.sample(dist);

        let mut ramdisk =
            RamdiskClient::create(512, 1 << 16).await.expect("failed to make ramdisk");
        let controller = ramdisk.take_controller().expect("invalid controller");
        let mut minfs = Minfs::new(controller);

        minfs.format().await.expect("failed to format minfs");
        let mut minfs = minfs.serve().await.expect("failed to mount minfs");
        minfs.bind_to_path(root).expect("failed to bind path");

        tree.write_tree_at(minfs.root()).await.expect("failed to write tree");

        let fixture = get_fixture(initial, increment, depth);
        let path = std::path::PathBuf::from(format!("{}/{}", root, fixture.name));
        assert!(path.is_dir(), "{}", path.display());
        for (i, entry) in std::fs::read_dir(&path).expect("failed to read directory").enumerate() {
            let entry = entry.expect("failed to read entry");
            let file_type = entry.file_type().expect("failed to get file type");
            assert!(file_type.is_file());
            let file_name =
                entry.file_name().into_string().expect("failed to convert file name to string");
            let expected_name = match &fixture.entries[i] {
                Entry::File(f) => f.name,
                Entry::Directory(_) => panic!("expected a file in the fixture tree"),
            };
            assert_eq!(file_name, expected_name.to_string());
        }

        minfs.shutdown().await.expect("failed to unmount minfs");
        ramdisk.destroy().await.expect("failed to destroy ramdisk");
    }
}