use std::{
    collections::BTreeSet,
    fs::File,
    io::Read,
    path::{Path, PathBuf},
};

use anyhow::{Context, Error};
use shared_buffer::OwnedBuffer;

use crate::{
    compat::Metadata,
    v2::write::{DirEntry, Directory, FileEntry},
    wasmer_package::Strictness,
    PathSegment, PathSegments,
};

/// A lazily loaded volume in a Wasmer package.
///
/// Note that it is the package resolver's role to interpret a package's
/// [`crate::metadata::annotations::FileSystemMappings`]. A [`Volume`] contains
/// directories as they were when the package was published.
#[derive(Debug, Clone, PartialEq)]
pub struct Volume {
    /// A pre-computed set of intermediate directories that are needed to allow
    /// access to the whitelisted files and directories.
    intermediate_directories: BTreeSet<PathBuf>,
    /// Specific files that this volume has access to.
    whitelisted_files: BTreeSet<PathBuf>,
    /// Directories that allow the user to access anything inside them.
    whitelisted_directories: BTreeSet<PathBuf>,
    /// The base directory all [`PathSegments`] will be resolved relative to.
    base_dir: PathBuf,
}

impl Volume {
    /// The name of the volume used to store metadata files.
    pub(crate) const METADATA: &str = "metadata";
    /// The name of the volume used to store files that will be available at
    /// runtime.
    pub(crate) const ASSET: &str = "atom";

    /// Create a new metadata volume.
    pub(crate) fn new_metadata(
        manifest: &wasmer_toml::Manifest,
        base_dir: impl Into<PathBuf>,
    ) -> Result<Self, Error> {
        let base_dir = base_dir.into();
        let mut files = BTreeSet::new();

        if let Some(license_file) = &manifest.package.license_file {
            files.insert(base_dir.join(license_file));
        }

        if let Some(readme) = &manifest.package.readme {
            files.insert(base_dir.join(readme));
        }

        for module in &manifest.modules {
            if let Some(bindings) = &module.bindings {
                let bindings_files = bindings.referenced_files(&base_dir)?;
                files.extend(bindings_files);
            }
        }

        Ok(Volume::new(base_dir, files, BTreeSet::new()))
    }

    /// Create a new volume for the assets.
    pub(crate) fn new_asset(
        manifest: &wasmer_toml::Manifest,
        base_dir: impl Into<PathBuf>,
    ) -> Result<Self, Error> {
        let base_dir = base_dir.into();
        let dirs: BTreeSet<_> = manifest
            .fs
            .values()
            .map(|path| base_dir.join(path))
            .collect();

        for path in &dirs {
            // Perform a basic sanity check to make sure the directories exist.
            let _ = std::fs::metadata(path).with_context(|| {
                format!("Unable to get the metadata for \"{}\"", path.display())
            })?;
        }

        Ok(Volume::new(base_dir, BTreeSet::new(), dirs))
    }

    fn new(
        base_dir: PathBuf,
        whitelisted_files: BTreeSet<PathBuf>,
        whitelisted_directories: BTreeSet<PathBuf>,
    ) -> Self {
        let mut intermediate_directories: BTreeSet<PathBuf> = whitelisted_files
            .iter()
            .filter_map(|p| p.parent())
            .chain(whitelisted_directories.iter().map(|p| p.as_path()))
            .flat_map(|dir| dir.ancestors())
            .filter(|dir| dir.starts_with(&base_dir))
            .map(|dir| dir.to_path_buf())
            .collect();

        // The base directory is always accessible (even if its contents isn't)
        intermediate_directories.insert(base_dir.clone());

        Volume {
            intermediate_directories,
            whitelisted_files,
            whitelisted_directories,
            base_dir,
        }
    }

    fn is_accessible(&self, path: &Path) -> bool {
        self.intermediate_directories.contains(path)
            || self.whitelisted_files.contains(path)
            || self
                .whitelisted_directories
                .iter()
                .any(|dir| path.starts_with(dir))
    }

    fn resolve(&self, path: &PathSegments) -> Option<PathBuf> {
        let resolved = resolve(&self.base_dir, path);
        let accessible = self.is_accessible(&resolved);
        accessible.then_some(resolved)
    }

    /// Read a file from the volume.
    pub fn read_file(&self, path: &PathSegments) -> Option<OwnedBuffer> {
        let path = self.resolve(path)?;
        let mut f = File::open(path).ok()?;

        // First we try to mmap it
        if let Ok(mmapped) = OwnedBuffer::from_file(&f) {
            return Some(mmapped);
        }

        // otherwise, fall back to reading the file's contents into memory
        let mut buffer = Vec::new();
        f.read_to_end(&mut buffer).ok()?;
        Some(OwnedBuffer::from_bytes(buffer))
    }

    /// Read the contents of a directory.
    pub fn read_dir(&self, path: &PathSegments) -> Option<Vec<(PathSegment, Metadata)>> {
        let resolved = self.resolve(path)?;
        let mut entries = Vec::new();

        for entry in resolved.read_dir().ok()? {
            let entry = entry.ok()?.path();

            if !self.is_accessible(&entry) {
                continue;
            }

            let segment: PathSegment = entry.file_name()?.to_str()?.parse().ok()?;

            let path = path.join(segment.clone());
            let metadata = self.metadata(&path)?;
            entries.push((segment, metadata));
        }

        entries.sort_by_key(|k| k.0.clone());

        Some(entries)
    }

    /// Get the metadata for a particular item.
    pub fn metadata(&self, path: &PathSegments) -> Option<Metadata> {
        let path = self.resolve(path)?;
        let meta = path.metadata().ok()?;

        if meta.is_dir() {
            Some(Metadata::Dir)
        } else if meta.is_file() {
            Some(Metadata::File {
                length: meta.len().try_into().ok()?,
            })
        } else {
            None
        }
    }

    pub(crate) fn as_directory_tree(&self, strictness: Strictness) -> Result<Directory<'_>, Error> {
        let mut paths = Vec::new();

        let asset_files = all_asset_files(
            self.whitelisted_directories.iter().map(|p| p.as_path()),
            &self.base_dir,
            strictness,
        )?;
        paths.extend(asset_files);

        paths.extend(self.whitelisted_files.iter().cloned());

        directory_tree(paths, &self.base_dir, strictness)
    }
}

/// Resolve a [`PathSegments`] to its equivalent path on disk.
fn resolve(base_dir: &Path, path: &PathSegments) -> PathBuf {
    let mut resolved = base_dir.to_path_buf();
    for segment in path.iter() {
        resolved.push(segment.as_str());
    }

    resolved
}

/// Recursively walk a set of directories, collecting the paths for all files
/// and directories into one big set.
fn all_asset_files<'a>(
    directories: impl IntoIterator<Item = &'a Path>,
    base_dir: &Path,
    strictness: Strictness,
) -> Result<BTreeSet<PathBuf>, Error> {
    let mut paths = BTreeSet::new();

    fn add_directory(
        dir: &Path,
        paths: &mut BTreeSet<PathBuf>,
        strictness: Strictness,
    ) -> Result<(), Error> {
        let entries = dir
            .read_dir()
            .with_context(|| format!("Unable to read the \"{}\" directory", dir.display()))?;

        for result in entries {
            let entry = result.with_context(|| {
                format!(
                    "Unable to read an item in the \"{}\" directory",
                    dir.display()
                )
            })?;

            let path = entry.path();
            let metadata = entry.metadata().with_context(|| {
                format!("Unable to get the metadata for \"{}\"", path.display())
            })?;

            if metadata.is_file() {
                paths.insert(path);
            } else if metadata.is_dir() {
                if let Err(e) = add_directory(&path, paths, strictness) {
                    strictness.on_error(&path, e)?;
                }

                paths.insert(path);
            }
        }

        Ok(())
    }

    for host_path in directories {
        let path = base_dir.join(host_path);
        add_directory(&path, &mut paths, strictness)?;
        paths.insert(path);
    }

    Ok(paths)
}

/// Given a list of absolute paths, create a directory tree relative to some
/// base directory.
fn directory_tree(
    paths: impl IntoIterator<Item = PathBuf>,
    base_dir: &Path,
    strictness: Strictness,
) -> Result<Directory<'static>, Error> {
    let mut root = Directory::default();

    for path in paths {
        if let Err(e) = insert_item_from_host(&mut root, &path, base_dir) {
            let error = e.context(format!(
                "Unable to add \"{}\" to the directory tree",
                path.display()
            ));
            strictness.on_error(&path, error)?;
        }
    }

    Ok(root)
}

fn insert_item_from_host(
    mut dir: &mut Directory<'_>,
    absolute: &Path,
    base_dir: &Path,
) -> Result<(), Error> {
    let path_in_volume = absolute.strip_prefix(base_dir).with_context(|| {
        format!(
            "Unable to add \"{}\" because it is outside the base directory ({})",
            absolute.display(),
            base_dir.display()
        )
    })?;

    if path_in_volume == Path::new("") {
        // We're trying to insert the base directory.
        return Ok(());
    }

    // Do the equivalent of std::fs::create_dir_all() and make "dir" point to
    // the parent directory.
    if let Some(parent) = path_in_volume.parent() {
        for component in parent.components() {
            match component {
                std::path::Component::Normal(component) => {
                    let segment = component
                        .to_str()
                        .and_then(|s| PathSegment::parse(s).ok())
                        .unwrap();
                    match dir
                        .children
                        .entry(segment)
                        .or_insert_with(|| DirEntry::Dir(Directory::default()))
                    {
                        DirEntry::Dir(d) => dir = d,
                        DirEntry::File(_) => {
                            anyhow::bail!(
                                "Can't nest a \"{}\" inside a file",
                                path_in_volume.display()
                            );
                        }
                    }
                }
                _ => unreachable!(
                    "The path should be fully resolved and relative to the base directory"
                ),
            }
        }
    }

    let filename = path_in_volume.file_name().and_then(|s| s.to_str()).unwrap();
    let segment = PathSegment::parse(filename)?;

    if absolute.is_file() {
        let file = FileEntry::from_path(absolute)
            .with_context(|| format!("Unable to open \"{}\"", absolute.display()))?;
        dir.children.insert(segment, file.into());
    } else if absolute.is_dir() {
        dir.children
            .insert(segment, DirEntry::Dir(Directory::default()));
    }

    Ok(())
}

#[cfg(test)]
mod tests {
    use tempfile::TempDir;

    use super::*;

    #[test]
    fn metadata_volume() {
        let temp = TempDir::new().unwrap();
        let wasmer_toml = r#"
            [package]
            name = "some/package"
            version = "0.0.0"
            description = ""
            license-file = "./path/to/LICENSE"
            readme = "README.md"

            [[module]]
            name = "asdf"
            source = "asdf.wasm"
            abi = "none"
            bindings = { wai-version = "0.2.0", exports = "asdf.wai", imports = ["browser.wai"] }
        "#;
        let license_dir = temp.path().join("path").join("to");
        std::fs::create_dir_all(&license_dir).unwrap();
        std::fs::write(license_dir.join("LICENSE"), "license").unwrap();
        std::fs::write(temp.path().join("README.md"), "readme").unwrap();
        std::fs::write(temp.path().join("asdf.wai"), "exports").unwrap();
        std::fs::write(temp.path().join("browser.wai"), "imports").unwrap();
        let manifest: wasmer_toml::Manifest = toml::from_str(wasmer_toml).unwrap();

        let volume = Volume::new_metadata(&manifest, temp.path().to_path_buf()).unwrap();

        let entries = volume.read_dir(&PathSegments::ROOT).unwrap();
        assert_eq!(
            entries,
            vec![
                (
                    PathSegment::parse("README.md").unwrap(),
                    Metadata::File { length: 6 },
                ),
                (
                    PathSegment::parse("asdf.wai").unwrap(),
                    Metadata::File { length: 7 },
                ),
                (
                    PathSegment::parse("browser.wai").unwrap(),
                    Metadata::File { length: 7 },
                ),
                (PathSegment::parse("path").unwrap(), Metadata::Dir),
            ],
        );
        let license: PathSegments = "/path/to/LICENSE".parse().unwrap();
        assert_eq!(
            String::from_utf8(volume.read_file(&license).unwrap().into()).unwrap(),
            "license"
        );
    }

    #[test]
    fn asset_volume() {
        let temp = TempDir::new().unwrap();
        let wasmer_toml = r#"
            [package]
            name = "some/package"
            version = "0.0.0"
            description = ""
            license_file = "./path/to/LICENSE"
            readme = "README.md"

            [[module]]
            name = "asdf"
            source = "asdf.wasm"
            abi = "none"
            bindings = { wai-version = "0.2.0", exports = "asdf.wai", imports = ["browser.wai"] }

            [fs]
            "/etc" = "./etc"
        "#;
        let license_dir = temp.path().join("path").join("to");
        std::fs::create_dir_all(&license_dir).unwrap();
        std::fs::write(license_dir.join("LICENSE"), "license").unwrap();
        std::fs::write(temp.path().join("README.md"), "readme").unwrap();
        std::fs::write(temp.path().join("asdf.wai"), "exports").unwrap();
        std::fs::write(temp.path().join("browser.wai"), "imports").unwrap();
        let share = temp.path().join("etc").join("share");
        std::fs::create_dir_all(&share).unwrap();
        std::fs::write(share.join("package.1"), "man page").unwrap();

        let manifest: wasmer_toml::Manifest = toml::from_str(wasmer_toml).unwrap();

        let volume = Volume::new_asset(&manifest, temp.path().to_path_buf()).unwrap();

        let entries = volume.read_dir(&PathSegments::ROOT).unwrap();
        assert_eq!(
            entries,
            vec![(PathSegment::parse("etc").unwrap(), Metadata::Dir)],
        );
        let man_page: PathSegments = "/etc/share/package.1".parse().unwrap();
        assert_eq!(
            String::from_utf8(volume.read_file(&man_page).unwrap().into()).unwrap(),
            "man page"
        );
    }
}