// 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,
    rand::{distributions, Rng},
    std::{fmt, fs::File, io::Write},
};

/// 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)
    }
}

/// 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(Clone, Eq, PartialEq)]
pub struct FileEntry {
    name: u64,
    contents: Vec<u8>,
}

impl fmt::Debug for FileEntry {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        write!(f, "FileEntry {{ name: {:?} }}", self.name)
    }
}

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 {
    fn write_file_at(self, root: &str) -> Result<(), Error> {
        let file_path = format!("{}/{}", root, self.name);
        let mut file = File::create(&file_path)?;
        file.write_all(&self.contents)?;
        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(self, root: &str) -> Result<(), Error> {
        let path = format!("{}/{}", root, self.get_name());
        std::fs::create_dir_all(&path)?;
        for entry in self.entries {
            match entry {
                Entry::File(file_entry) => file_entry.write_file_at(&path)?,
                Entry::Directory(dir_entry) => dir_entry.write_tree_at(&path)?,
            }
        }
        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,
        ramdevice_client::RamdiskClient,
        rand::{rngs::StdRng, Rng, SeedableRng},
    };

    // 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(seed: u64, depth: u32) -> DirectoryEntry {
        // confirm that the seed and depth used to generate this tree are the same.
        assert_eq!(seed, 1234, "seed value changed - update fixture");
        assert_eq!(depth, 2, "depth value changed - update fixture");
        DirectoryEntry {
            name: 6998163126151612998,
            entries: vec![
                Entry::File(FileEntry { name: 15265247264153552834, contents: vec![] }),
                Entry::File(FileEntry { name: 8405537650559459291, contents: vec![] }),
                Entry::File(FileEntry { name: 12776892210595361501, contents: vec![] }),
            ],
        }
    }

    #[test]
    fn gen_tree() {
        let seed = 1234;
        let depth = 2;
        // make sure we are generating a tree at all.
        let mut rng = StdRng::seed_from_u64(seed);
        let dist = EntryDistribution::new(depth);
        let tree: DirectoryEntry = rng.sample(dist);
        println!("generated tree: {:?}", tree);
        // this skips the contents for the files, because they are huge, so we do our own manual
        // comparison.
        let fixture = get_fixture(seed, depth);
        assert_eq!(fixture.name, tree.name);
        assert_eq!(fixture.entries.len(), tree.entries.len());
        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_seed_same_tree() {
        // this confirms an assumption about the properties of tree generation that we rely on for
        // consistency checking.
        let seed = 1234;
        let gen_tree: fn(u64) -> DirectoryEntry = |seed| {
            let mut rng = StdRng::seed_from_u64(seed);
            let dist = EntryDistribution::new(5);
            rng.sample(dist)
        };

        let tree1 = gen_tree(seed);
        let tree2 = gen_tree(seed);

        assert_eq!(tree1, tree2);
    }

    #[test]
    fn write_tree() {
        let root = "/test-root";
        let seed = 1234;
        let depth = 2;

        let mut rng = StdRng::seed_from_u64(seed);
        let dist = EntryDistribution::new(depth);
        let tree: DirectoryEntry = rng.sample(dist);

        let ramdisk = RamdiskClient::create(512, 1 << 16).expect("failed to make ramdisk");
        let device_path = ramdisk.get_path();

        let mut minfs = Minfs::new(device_path).expect("failed to make new minfs");
        minfs.format().expect("failed to format minfs");
        minfs.mount(root).expect("failed to mount minfs");

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

        let fixture = get_fixture(seed, depth);
        let path = std::path::PathBuf::from(format!("{}/{}", root, fixture.name));
        assert!(path.is_dir());
        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.unmount().expect("failed to unmount minfs");
        ramdisk.destroy().expect("failed to destroy ramdisk");
    }
}