/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::fmt::Debug;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context as _;
use buck2_common::cas_digest::CasDigestConfig;
use buck2_common::cas_digest::DigestAlgorithm;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_common::external_symlink::ExternalSymlink;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_core::directory::find;
use buck2_core::directory::unordered_entry_walk;
use buck2_core::directory::DashMapDirectoryInterner;
use buck2_core::directory::Directory;
use buck2_core::directory::DirectoryBuilder;
use buck2_core::directory::DirectoryEntry;
use buck2_core::directory::DirectoryHasher;
use buck2_core::directory::DirectoryIterator;
use buck2_core::directory::DirectorySelector;
use buck2_core::directory::FingerprintedDirectory;
use buck2_core::directory::ImmutableDirectory;
use buck2_core::directory::SharedDirectory;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::paths::RelativePath;
use buck2_core::fs::paths::RelativePathBuf;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use chrono::DateTime;
use chrono::Utc;
use derive_more::Display;
use dupe::Dupe;
use once_cell::sync::Lazy;
use ref_cast::RefCast;
use remote_execution as RE;
use starlark_map::small_map::SmallMap;
use thiserror::Error;

use crate::artifact_value::ArtifactValue;
use crate::digest::CasDigestFromReExt;
use crate::digest::CasDigestToReExt;
use crate::digest_config::DigestConfig;
use crate::re::manager::ManagedRemoteExecutionClient;

#[allocative::root]
pub static INTERNER: Lazy<DashMapDirectoryInterner<ActionDirectoryMember, TrackedFileDigest>> =
    Lazy::new(DashMapDirectoryInterner::new);

/// Represents a relative symlink, and stores the symlink's target path.
#[derive(Debug, Display, Eq, PartialEq, Allocative)]
pub struct Symlink(RelativePathBuf);

#[derive(Clone, Debug, Dupe, PartialEq, Eq, Display, Allocative)]
pub enum ActionDirectoryMember {
    File(FileMetadata),
    Symlink(Arc<Symlink>),
    ExternalSymlink(Arc<ExternalSymlink>),
}

pub type ActionDirectoryEntry<D> = DirectoryEntry<D, ActionDirectoryMember>;

pub type ActionImmutableDirectory = ImmutableDirectory<ActionDirectoryMember, TrackedFileDigest>;

pub type ActionSharedDirectory = SharedDirectory<ActionDirectoryMember, TrackedFileDigest>;

pub type ActionDirectoryBuilder = DirectoryBuilder<ActionDirectoryMember, TrackedFileDigest>;

pub trait ActionDirectory = Directory<ActionDirectoryMember, TrackedFileDigest>;

pub trait ActionFingerprintedDirectory =
    FingerprintedDirectory<ActionDirectoryMember, TrackedFileDigest>;

#[derive(Allocative, RefCast)]
#[repr(transparent)]
pub struct ReDirectorySerializer {
    pub cas_digest_config: CasDigestConfig,
}

impl ReDirectorySerializer {
    fn create_re_directory<'a, D, I>(entries: I) -> RE::Directory
    where
        I: IntoIterator<
            Item = (
                &'a FileName,
                DirectoryEntry<&'a D, &'a ActionDirectoryMember>,
            ),
        >,
        D: ActionFingerprintedDirectory + ?Sized + 'a,
    {
        let mut files: Vec<RE::FileNode> = Vec::new();
        let mut directories: Vec<RE::DirectoryNode> = Vec::new();
        let mut symlinks: Vec<RE::SymlinkNode> = Vec::new();

        for (name, entry) in entries {
            let name = name.as_str().into();

            match entry {
                DirectoryEntry::Dir(d) => {
                    directories.push(RE::DirectoryNode {
                        name,
                        digest: Some(d.fingerprint().to_grpc()),
                    });
                }
                DirectoryEntry::Leaf(ActionDirectoryMember::File(f)) => {
                    files.push(RE::FileNode {
                        name,
                        digest: Some(f.digest.to_grpc()),
                        is_executable: f.is_executable,
                        ..Default::default()
                    });
                }
                DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(s)) => {
                    symlinks.push(RE::SymlinkNode {
                        name,
                        target: s.to_string(),
                        ..Default::default()
                    });
                }
                DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(s)) => {
                    symlinks.push(RE::SymlinkNode {
                        name,
                        target: s.target_str().to_owned(),
                        ..Default::default()
                    });
                }
            }
        }

        files.sort_by(|a, b| a.name.cmp(&b.name));
        directories.sort_by(|a, b| a.name.cmp(&b.name));
        symlinks.sort_by(|a, b| a.name.cmp(&b.name));

        RE::Directory {
            files,
            directories,
            symlinks,
            ..Default::default()
        }
    }

    pub fn serialize_entries<'a, D, I>(entries: I) -> Vec<u8>
    where
        I: IntoIterator<
            Item = (
                &'a FileName,
                DirectoryEntry<&'a D, &'a ActionDirectoryMember>,
            ),
        >,
        D: ActionFingerprintedDirectory + ?Sized + 'a,
    {
        proto_serialize(&Self::create_re_directory(entries))
    }
}

fn proto_serialize<M: prost::Message>(m: &M) -> Vec<u8> {
    let mut serialized_buf = Vec::new();
    m.encode(&mut serialized_buf).unwrap();
    serialized_buf
}

impl DirectoryHasher<ActionDirectoryMember, TrackedFileDigest> for ReDirectorySerializer {
    fn hash_entries<'a, D, I>(&self, entries: I) -> TrackedFileDigest
    where
        I: IntoIterator<
            Item = (
                &'a FileName,
                DirectoryEntry<&'a D, &'a ActionDirectoryMember>,
            ),
        >,
        D: ActionFingerprintedDirectory + 'a,
    {
        TrackedFileDigest::from_content(&Self::serialize_entries(entries), self.cas_digest_config)
    }
}

impl Symlink {
    pub fn new(target: RelativePathBuf) -> Self {
        Self(target)
    }

    /// Returns the path the symlink points to.
    pub fn target(&self) -> &RelativePath {
        self.0.as_relative_path()
    }

    /// Creates a new `Symlink` from `self`, such that its target is
    /// `src_relative_to_dest/self.target()`.
    ///
    /// This solves a specific problem: symlink at path `src` points to an
    /// artifact at path `src/target`. We move the symlink to `dest`, but we
    /// want to keep linking to the same artifact. How can we adjust the target
    /// in order to achieve that? We need to know how to get to `src` from
    /// `dest`, which is what `src_relative_to_dest` tells us.
    pub fn relativized<P: AsRef<RelativePath>>(&self, src_relative_to_dest: P) -> Self {
        // FIXME(rafaelc): we don't need to normalize the target anymore!
        let relativized_t = src_relative_to_dest.as_ref().join_normalized(&self.0);
        Self(relativized_t)
    }
}

pub fn new_symlink<T: AsRef<Path>>(target: T) -> anyhow::Result<ActionDirectoryMember> {
    let target = target.as_ref();
    if target.is_absolute() {
        Ok(ActionDirectoryMember::ExternalSymlink(Arc::new(
            ExternalSymlink::new(target.to_path_buf(), None)?,
        )))
    } else {
        Ok(ActionDirectoryMember::Symlink(Arc::new(Symlink::new(
            RelativePathBuf::from_path(target).unwrap(),
        ))))
    }
}

pub fn directory_to_re_tree(directory: &impl ActionFingerprintedDirectory) -> RE::Tree {
    let children = directory
        .fingerprinted_ordered_walk()
        .without_paths()
        .filter_map(|entry| match entry {
            DirectoryEntry::Dir(d) => Some(d),
            DirectoryEntry::Leaf(..) => None,
        })
        .map(|d| ReDirectorySerializer::create_re_directory(d.fingerprinted_entries()))
        .collect();

    let root = ReDirectorySerializer::create_re_directory(directory.fingerprinted_entries());

    RE::Tree {
        root: Some(root),
        children,
    }
}

pub async fn re_directory_to_re_tree(
    directory: RE::Directory,
    client: &ManagedRemoteExecutionClient,
    use_case: RemoteExecutorUseCase,
) -> anyhow::Result<RE::Tree> {
    let mut children: Vec<RE::Directory> = vec![];
    let mut frontier = directory.directories.clone();
    while !frontier.is_empty() {
        let digests: Vec<remote_execution::TDigest> = frontier
            .into_iter()
            .filter_map(|d| d.digest)
            .map(|digest| remote_execution::TDigest {
                hash: digest.hash.clone(),
                size_in_bytes: digest.size_bytes,
                ..Default::default()
            })
            .collect();
        let mut retrieved = client
            .download_typed_blobs::<RE::Directory>(digests, use_case)
            .await?;
        frontier = retrieved
            .iter()
            .flat_map(|d| d.directories.clone())
            .collect();
        children.append(&mut retrieved);
    }
    Ok(RE::Tree {
        root: Some(directory),
        children,
    })
}

/// Constructs a `Directory` from an `RE::Tree`. As long as the
/// `RE::Tree` is valid (i.e. nothing is broken in the RE side), this
/// should always succeed.
#[allow(clippy::trivially_copy_pass_by_ref)] // SystemTime is a different size on Windows
pub fn re_tree_to_directory(
    tree: &RE::Tree,
    leaf_expires: &DateTime<Utc>,
    digest_config: DigestConfig,
) -> anyhow::Result<ActionDirectoryBuilder> {
    /// A map of digests to directories, populated lazily when we access it based on the hash we
    /// use. We need this because in a RE tree, the directories in the tree don't carry their hash,
    /// but the pointers are hashes, so we need to first see a hash before we can work out what
    /// hashing mechanism to use here.
    struct DirMap<'a> {
        by_kind: SmallMap<DigestAlgorithm, HashMap<FileDigest, &'a RE::Directory>>,
        directories: &'a [RE::Directory],
    }

    impl<'a> DirMap<'a> {
        fn new(directories: &'a [RE::Directory]) -> Self {
            Self {
                by_kind: SmallMap::new(),
                directories,
            }
        }

        fn get(&mut self, algo: DigestAlgorithm, digest: &FileDigest) -> Option<&'a RE::Directory> {
            let map = self.by_kind.entry(algo).or_insert_with(|| {
                self.directories
                    .iter()
                    .map(|d| (from_proto_message(d, algo), d))
                    .collect()
            });

            map.get(digest).copied()
        }
    }

    fn from_proto_message<M: prost::Message>(m: &M, algo: DigestAlgorithm) -> FileDigest {
        let mut m_encoded = Vec::new();
        m.encode(&mut m_encoded)
            .unwrap_or_else(|e| unreachable!("Protobuf messages are always encodeable: {}", e));
        FileDigest::from_content_for_algorithm(m_encoded.as_slice(), algo)
    }

    // Recursively builds the directory
    fn dfs_build<'a>(
        re_dir: &'_ RE::Directory,
        re_dir_name: &'_ (impl fmt::Display + ?Sized),
        dirmap: &'_ mut DirMap<'a>,
        leaf_expires: &DateTime<Utc>,
        digest_config: DigestConfig,
    ) -> anyhow::Result<ActionDirectoryBuilder> {
        let mut builder = ActionDirectoryBuilder::empty();
        for node in &re_dir.files {
            let name = FileNameBuf::try_from(node.name.clone()).with_context(|| {
                DirectoryReConversionError::IncorrectFileName {
                    name: node.name.clone(),
                    dir: re_dir_name.to_string(),
                }
            })?;

            let digest = node.digest.as_ref().with_context(|| {
                DirectoryReConversionError::NodeWithDigestNone {
                    name: node.name.clone(),
                    dir: re_dir_name.to_string(),
                }
            })?;
            let digest = FileDigest::from_grpc(digest, digest_config)?;
            let digest = TrackedFileDigest::new_expires(
                digest,
                *leaf_expires,
                digest_config.cas_digest_config(),
            );

            let member = ActionDirectoryMember::File(FileMetadata {
                digest,
                is_executable: node.is_executable,
            });

            builder.insert(name, DirectoryEntry::Leaf(member))?;
        }
        for node in &re_dir.symlinks {
            builder.insert(
                FileNameBuf::try_from(node.name.clone()).map_err(|_| {
                    DirectoryReConversionError::IncorrectFileName {
                        name: node.name.clone(),
                        dir: re_dir_name.to_string(),
                    }
                })?,
                DirectoryEntry::Leaf(node.try_into()?),
            )?;
        }
        for dir_node in &re_dir.directories {
            let child_digest = match &dir_node.digest {
                None => {
                    return Err(DirectoryReConversionError::NodeWithDigestNone {
                        name: dir_node.name.clone(),
                        dir: re_dir_name.to_string(),
                    }
                    .into());
                }
                Some(d) => d,
            };
            let (child_digest, child_digest_algo) =
                FileDigest::from_grpc_with_algo(child_digest, digest_config)?;
            let child_re_dir = match dirmap.get(child_digest_algo, &child_digest) {
                None => {
                    return Err(DirectoryReConversionError::IncompleteTreeChildrenList {
                        name: dir_node.name.clone(),
                        dir: re_dir_name.to_string(),
                    }
                    .into());
                }
                Some(d) => d,
            };
            let dir = dfs_build(
                child_re_dir,
                &child_digest,
                dirmap,
                leaf_expires,
                digest_config,
            )?;
            builder.insert(
                FileNameBuf::try_from(dir_node.name.clone()).map_err(|_| {
                    DirectoryReConversionError::IncorrectFileName {
                        name: dir_node.name.clone(),
                        dir: re_dir_name.to_string(),
                    }
                })?,
                DirectoryEntry::Dir(dir),
            )?;
        }

        // NOTE: We re-digest the directories we just received here (instead of trusting RE with
        // the hashes). But, since output directories tend to be small, that doesn't actually
        // matter.
        Ok(builder)
    }

    let root_dir = match &tree.root {
        Some(d) => d,
        None => return Ok(ActionDirectoryBuilder::empty()),
    };

    dfs_build(
        root_dir,
        "root directory",
        &mut DirMap::new(&tree.children),
        leaf_expires,
        digest_config,
    )
}

#[derive(Debug, Error)]
pub enum DirectoryReConversionError {
    // Conversion from RE::Tree errors (these shouldn't happen unless something is broken on RE side)
    #[error("Converting RE::Tree to Directory, dir `{dir}` has child `{name}` with digest=None.")]
    NodeWithDigestNone { name: String, dir: String },
    #[error(
        "Converting RE::Tree to Directory, dir `{dir}` has child `{name}` not in tree.children."
    )]
    IncompleteTreeChildrenList { name: String, dir: String },
    #[error("Converting RE::Tree to Directory, dir `{dir}` has incorrect file name `{name}`")]
    IncorrectFileName { name: String, dir: String },
}

impl<'a> TryFrom<&'a RE::SymlinkNode> for ActionDirectoryMember {
    type Error = anyhow::Error;

    fn try_from(node: &'a RE::SymlinkNode) -> Result<Self, Self::Error> {
        let symlink = if node.target.starts_with('/') {
            ActionDirectoryMember::ExternalSymlink(Arc::new(ExternalSymlink::new(
                PathBuf::from(node.target.as_str()),
                None,
            )?))
        } else {
            ActionDirectoryMember::Symlink(Arc::new(Symlink(RelativePathBuf::from(
                node.target.clone(),
            ))))
        };
        Ok(symlink)
    }
}

pub fn relativize_directory(
    builder: &mut ActionDirectoryBuilder,
    orig_root: &ProjectRelativePath,
    new_root: &ProjectRelativePath,
) -> anyhow::Result<()> {
    let mut replacements = ActionDirectoryBuilder::empty();

    {
        let mut walk = builder.unordered_walk();
        while let Some((path, entry)) = walk.next() {
            let link = match entry {
                DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(link)) => link,
                _ => continue,
            };

            let path = path.get();

            let orig_path = orig_root.join_normalized(&path)?;
            let new_path = new_root.join_normalized(&path)?;

            let orig_dest = orig_path
                .parent()
                .context("Symlink has no dir parent")?
                .join_normalized(link.target())?;

            let new_dest = new_path
                .parent()
                .context("Symlink has no dir parent")?
                .as_forward_relative_path()
                .as_relative_path()
                .relative(orig_dest);
            let new_link = Arc::new(Symlink::new(new_dest));

            replacements.insert(
                &path,
                DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(new_link)),
            )?;
        }
    }

    builder.merge(replacements)?;

    Ok(())
}

pub fn insert_entry(
    builder: &mut ActionDirectoryBuilder,
    path: &ProjectRelativePath,
    entry: ActionDirectoryEntry<ActionDirectoryBuilder>,
) -> anyhow::Result<()> {
    if let DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(s)) = entry {
        // ExternalSymlink is a bit of an odd concept. The way it works is that when you try to
        // read `foo/bar` and it turns out `foo` is an ExternalSymlink, we tell you the value of
        // `foo/bar` is an `ExternalSymlink` to some location with the "rest of the path" being
        // `bar`.

        // Later, when you try to add this to an action directory (this code), we need to put a
        // symlink at `foo` but all we have is the path this `ExternalSymlink` piece of data exists
        // at (`foo/bar`), and the "rest" of the path (`bar`), so we strip the rest of the path and
        // conclude the symlink must be at `foo` and should point at the external symlink's path
        // without the "rest" component.

        // This is kinda wonky though it sort of works. However, it does not work when we try to
        // read a symlink that points through an external symlink, e.g. reading `bar` in those
        // circumstances:

        // ```
        // bar -> foo/xxx
        // foo -> /home/blah
        // ```

        // When this happens, the "value" of `bar` is `ExternalSymlink(..., rest = "xxx")`, but the
        // path is `bar`, and we can't infer that the path must actually be `bar` and that the
        // symlink should be `/home/blah/xxx` (so in this case we keep the rest).

        // NOTE (torozco): I think we *could* perhaps make this work by using `path` as-is if we
        // can't "fix" it, and using the full symlink destination instead of
        // `s.without_remaining_path()`, but considering this doesn't seem to be something very
        // widespread and I might be missing something, I did not do it.
        let fixed_source_path = s.fix_source_path(path.as_ref()).with_context(|| {
            format!("Error locating source path for symlink at {}: {}", path, s)
        })?;
        let path = ProjectRelativePath::unchecked_new(fixed_source_path.as_str());
        let entry = DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(
            s.without_remaining_path(),
        ));
        builder.insert(path, entry)?;
    } else {
        builder.insert(path, entry)?;
    }

    Ok(())
}

/// Inserts an input to the tree, which will be required when following
/// symlinks to calculate the `deps` of the `ArtifactValue`.
pub fn insert_artifact(
    builder: &mut ActionDirectoryBuilder,
    path: &ProjectRelativePath,
    value: &ArtifactValue,
) -> anyhow::Result<()> {
    insert_entry(
        builder,
        path,
        value.entry().dupe().map_dir(|d| d.into_builder()),
    )?;
    // add input's deps
    if let Some(deps) = value.deps() {
        builder.merge(deps.dupe().into_builder())?;
    }
    Ok(())
}

pub fn insert_file(
    builder: &mut ActionDirectoryBuilder,
    path: &ProjectRelativePath,
    meta: FileMetadata,
) -> anyhow::Result<()> {
    insert_entry(
        builder,
        path,
        DirectoryEntry::Leaf(ActionDirectoryMember::File(meta)),
    )
}

#[cfg(test)]
pub fn insert_symlink(
    builder: &mut ActionDirectoryBuilder,
    path: &ProjectRelativePath,
    symlink: Arc<Symlink>,
) -> anyhow::Result<()> {
    insert_entry(
        builder,
        path,
        DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(symlink)),
    )
}

pub fn expand_selector_for_dependencies(
    builder: &ActionDirectoryBuilder,
    paths_to_take: &mut DirectorySelector,
) {
    // Note that we receive the list of paths to take as an input and start from that for the paths
    // to visit. Rhis ensures that if we add symlinks that point into the paths we are about to
    // take out, we don't bother visiting them individually, since we'll start by visiting the whole
    // thing.
    let mut paths_to_visit = paths_to_take.clone();

    let mut all_known_symlinks = HashSet::new();

    while !paths_to_visit.is_empty() {
        let mut next_paths_to_visit = DirectorySelector::empty();

        let mut handle_link_at_path = |link_path: &ForwardRelativePath, link: &Symlink| {
            let link_dst = match link_path.parent() {
                Some(parent) => parent.join_system(&link.target().to_path("")),
                None => ForwardRelativePath::new(link.target().as_str()).map(ToOwned::to_owned),
            };

            let dest = match link_dst {
                Ok(d) => d,
                Err(e) => {
                    // If we get here the symlink we found points to something that is not
                    // relative to this directory so we bail.
                    tracing::trace!(
                        "Invalid symlink: {} -> {}: {:#}",
                        link_path,
                        link.target(),
                        e
                    );
                    if cfg!(test) {
                        panic!("Invalid symlink found")
                    }
                    return;
                }
            };

            if all_known_symlinks.insert(dest.clone()) {
                tracing::trace!("New symlink: {} -> {}", link_path, dest);
                next_paths_to_visit.select(&dest);
                paths_to_take.select(&dest);
            }
        };

        let mut search = paths_to_visit.unordered_search(builder);

        while let Some((entry_path, entry)) = search.next() {
            // We're searching for all paths that are symlink destinations to find _new_ symlinks.
            // We might encounter paths that don't actually work while we do this. If that happens,
            // then we might be in a situation where we were trying to traverse *through* a
            // symlink. We check for that and handle the newly-discovered symlink if that happens.
            let mut entry_walk = match entry {
                Ok(entry) => unordered_entry_walk(entry),
                Err(e) => {
                    let link = match e.into_leaf() {
                        ActionDirectoryMember::Symlink(link) => link,
                        e => {
                            tracing::trace!(
                                "Ignoring entry to visit at: {}: cannot traverse through: {:?}",
                                entry_path,
                                e,
                            );

                            continue;
                        }
                    };

                    handle_link_at_path(&entry_path.get(), link);
                    continue;
                }
            };

            while let Some((sub_entry_path, sub_entry)) = entry_walk.next() {
                let link = match sub_entry {
                    DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(link)) => link,
                    _ => continue,
                };

                let link_path = entry_path.get().join(sub_entry_path.get());
                handle_link_at_path(&link_path, link);
            }
        }

        paths_to_visit = next_paths_to_visit;
    }
}

/// Given a builder and a Path, return it as an ArtifactValue. Dependencies are produced by
/// extracting any symlinks found in the builder.
pub fn extract_artifact_value(
    builder: &ActionDirectoryBuilder,
    path: &ProjectRelativePath,
    digest_config: DigestConfig,
) -> anyhow::Result<Option<ArtifactValue>> {
    let entry = match find(builder, path.as_forward_relative_path())? {
        Some(entry) => entry,
        _ => return Ok(None),
    };

    // This will be the set of all paths referenced by symlinks in our directory.
    let mut paths_to_take = DirectorySelector::empty();
    paths_to_take.select(path);

    expand_selector_for_dependencies(builder, &mut paths_to_take);

    let mut deps = ActionDirectoryBuilder::empty();
    let mut has_deps = false;

    for (entry_path, entry) in paths_to_take.unordered_search(builder).with_paths() {
        let entry = match entry {
            Ok(entry) => entry,
            Err(err) => {
                // We'll get here if we have a symlink that points at a path that traverses another
                // symlink (the my_genrule10 testcase has an example of this).
                tracing::trace!(
                    "Taking entry as a symlink traverses through it: {}",
                    entry_path
                );
                DirectoryEntry::Leaf(err.into_leaf())
            }
        };

        if entry_path == path.as_forward_relative_path() {
            continue;
        }

        has_deps = true;

        // TODO (@torozco): We could improve this a lot by sharing() those entries in place when we
        // extract them.
        deps.insert(
            &entry_path,
            entry.map_leaf(|l| l.dupe()).map_dir(|d| d.to_builder()),
        )?;
    }

    let entry = entry.map_leaf(|l| l.dupe()).map_dir(|d| {
        d.to_builder()
            .fingerprint(digest_config.as_directory_serializer())
            .shared(&*INTERNER)
    });

    let deps = if has_deps {
        Some(
            deps.fingerprint(digest_config.as_directory_serializer())
                .shared(&*INTERNER),
        )
    } else {
        None
    };

    Ok(Some(ArtifactValue::new(entry, deps)))
}

#[cfg(test)]
mod tests {
    use anyhow::Context as _;

    use super::*;

    fn path<'a>(s: &'a str) -> &'a ProjectRelativePath {
        ProjectRelativePath::new(s).unwrap()
    }

    fn assert_dirs_eq(d1: &impl ActionDirectory, d2: &impl ActionDirectory) {
        let mut w1 = d1.ordered_walk();
        let mut w2 = d2.ordered_walk();

        loop {
            let n1 = w1.next();
            let n2 = w2.next();

            let ((p1, e1), (p2, e2)) = match (n1, n2) {
                (Some(n1), Some(n2)) => (n1, n2),
                (Some((p1, _)), None) => {
                    panic!("Walks differ, first difference at: {}", p1);
                }
                (None, Some((p2, _))) => {
                    panic!("Walks differ, first difference at: {}", p2);
                }
                (None, None) => break,
            };

            // Compare paths
            assert_eq!(p1.get(), p2.get(), "Walks produced different paths");

            // Compare leaves
            assert_eq!(
                e1.map_dir(|_| ()),
                e2.map_dir(|_| ()),
                "Walks differ at path {}",
                p1
            );
        }
    }

    #[test]
    fn directory_relativized() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();

        let mut dir = {
            let mut builder = ActionDirectoryBuilder::empty();
            insert_file(
                &mut builder,
                path("inter/f"),
                FileMetadata::empty(digest_config.cas_digest_config()),
            )?;
            insert_symlink(
                &mut builder,
                path("sym"),
                Arc::new(Symlink::new("inter/f".into())),
            )?;
            insert_symlink(
                &mut builder,
                path("esym"),
                Arc::new(Symlink::new("../f".into())),
            )?;
            builder
        };

        let expected_dir = {
            let mut builder = ActionDirectoryBuilder::empty();
            insert_file(
                &mut builder,
                path("inter/f"),
                FileMetadata::empty(digest_config.cas_digest_config()),
            )?;
            insert_symlink(
                &mut builder,
                path("sym"),
                Arc::new(Symlink::new("../a/d0/inter/f".into())),
            )?;
            insert_symlink(
                &mut builder,
                path("esym"),
                Arc::new(Symlink::new("../a/f".into())),
            )?;
            builder
        };

        // Move directory from a/d0 to b.
        relativize_directory(&mut dir, path("a/d0"), path("b"))?;

        assert_dirs_eq(&dir, &expected_dir);

        Ok(())
    }

    fn build_test_dir() -> anyhow::Result<ActionDirectoryBuilder> {
        let digest_config = DigestConfig::testing_default();

        let mut builder = ActionDirectoryBuilder::empty();
        // /
        // |-f1
        // |-d1/
        // | |-d2/
        // | | |-d3/
        // | | | |-s1 -> ../../../d6/s4          d3_deps:[s4, f4]
        // | | | |-s2 -> ../d4                   d3_deps:[d4/*]
        // | | | |-s3 -> ../d5/../../../f1       d3_deps:[f1]
        // | | |-d4/
        // | | | |-f2
        // | | |-d5/
        // | | | |-f3
        // |-d6/
        // | |-s4 -> f4
        // | |-f4
        // | |-e1 -> /mnt/gvfs
        // |-f5

        for file in &["f1", "d1/d2/d4/f2", "d1/d2/d5/f3", "d6/f4", "f5"] {
            insert_file(
                &mut builder,
                path(file),
                FileMetadata::empty(digest_config.cas_digest_config()),
            )?;
        }

        for (sym, target) in &[
            ("d1/d2/d3/s1", "../../../d6/s4"),
            ("d1/d2/d3/s2", "../d4"),
            ("d1/d2/d3/s3", "../d5/../../../f1"),
            ("d6/s4", "f4"),
            ("d6/e1", "/mnt/gvfs"),
        ] {
            insert_symlink(
                &mut builder,
                path(sym),
                Arc::new(Symlink::new(target.into())),
            )?;
        }

        Ok(builder)
    }

    /// All deps of d6 are internal to it, so we don't return any deps.
    #[test]
    fn test_extract_no_deps() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();
        let root = build_test_dir()?;
        let value =
            extract_artifact_value(&root, path("d6"), digest_config)?.context("Not value!")?;
        assert!(value.deps().is_none());
        Ok(())
    }

    #[test]
    fn test_extract_has_deps_dir() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();

        let root = build_test_dir()?;
        let value = extract_artifact_value(&root, path("d1/d2/d3"), digest_config)?
            .context("Not value!")?;

        let expected = {
            let mut builder = ActionDirectoryBuilder::empty();

            for p in &["d6/s4", "d6/f4", "d1/d2/d4", "f1"] {
                let path = path(p);
                let entry = find(&root, path.as_forward_relative_path())?
                    .with_context(|| format!("Missing {}", path))?
                    .map_dir(|d| d.to_builder())
                    .map_leaf(|l| l.dupe());
                insert_entry(&mut builder, path, entry)?;
            }

            builder
        };

        assert_dirs_eq(value.deps().context("No deps!")?, &expected);

        Ok(())
    }

    #[test]
    fn test_extract_has_deps_leaf() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();

        let root = build_test_dir()?;
        let value = extract_artifact_value(&root, path("d1/d2/d3/s3"), digest_config)?
            .context("Not value!")?;

        let expected = {
            let mut builder = ActionDirectoryBuilder::empty();
            insert_file(
                &mut builder,
                path("f1"),
                FileMetadata::empty(digest_config.cas_digest_config()),
            )?;
            builder
        };

        assert_dirs_eq(value.deps().context("No deps!")?, &expected);

        Ok(())
    }

    #[test]
    fn test_extract_cycle() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();

        let mut builder = ActionDirectoryBuilder::empty();

        for (sym, target) in &[("d1/f1", "../d2/f2"), ("d2/f2", "../d1/f1")] {
            insert_symlink(
                &mut builder,
                path(sym),
                Arc::new(Symlink::new(target.into())),
            )?;
        }

        let expected = {
            let mut builder = ActionDirectoryBuilder::empty();
            insert_symlink(
                &mut builder,
                path("d2/f2"),
                Arc::new(Symlink::new("../d1/f1".into())),
            )?;
            builder
        };

        let value = extract_artifact_value(&builder, path("d1/f1"), digest_config)?
            .context("Not value!")?;

        assert_dirs_eq(value.deps().context("No deps!")?, &expected);

        Ok(())
    }

    #[test]
    fn test_extract_symlink_chain() -> anyhow::Result<()> {
        // Crank up the difficulty: l1 points through d3/f, but through l2. We need all of those in
        // the deps! In practice, this tends to not happen in Buck 2 because we always dereference
        // symlinks and traverse them, but might a well support it properly.
        let digest_config = DigestConfig::testing_default();

        let mut builder = ActionDirectoryBuilder::empty();
        insert_symlink(
            &mut builder,
            path("l1"),
            Arc::new(Symlink::new("l2/f".into())),
        )?;

        insert_symlink(
            &mut builder,
            path("l2"),
            Arc::new(Symlink::new("d3".into())),
        )?;

        insert_file(
            &mut builder,
            path("d3/f"),
            FileMetadata::empty(digest_config.cas_digest_config()),
        )?;

        let value =
            extract_artifact_value(&builder, path("l1"), digest_config)?.context("Not value!")?;

        let expected = {
            let mut builder = ActionDirectoryBuilder::empty();

            insert_symlink(
                &mut builder,
                path("l2"),
                Arc::new(Symlink::new("d3".into())),
            )?;

            insert_file(
                &mut builder,
                path("d3/f"),
                FileMetadata::empty(digest_config.cas_digest_config()),
            )?;

            builder
        };

        assert_dirs_eq(value.deps().context("No deps!")?, &expected);

        Ok(())
    }

    #[test]
    fn test_re_tree_roundtrip() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();

        let mut builder = ActionDirectoryBuilder::empty();
        insert_file(
            &mut builder,
            path("a/aa"),
            FileMetadata::empty(digest_config.cas_digest_config()),
        )?;
        insert_file(
            &mut builder,
            path("a/bb"),
            FileMetadata::empty(digest_config.cas_digest_config()),
        )?;
        insert_file(
            &mut builder,
            path("b/b"),
            FileMetadata::empty(digest_config.cas_digest_config()),
        )?;
        let dir = builder.fingerprint(digest_config.as_directory_serializer());

        let tree = directory_to_re_tree(&dir);
        let dir2 = re_tree_to_directory(&tree, &Utc::now(), digest_config)?;

        assert_dirs_eq(&dir, &dir2);

        Ok(())
    }

    /// Ensure that we serialize trees the same way RE does. The expected hash was obtained by
    /// running:
    ///
    /// ```sh
    /// buck2 run fbcode//remote_execution/rust/recli:recli -- \
    ///     exec command --out-dir test -- sh -c \
    ///     'mkdir -p test/a/aa test/a/aaa test/b/bb test/d && touch test/a/aa/f test/a/aaa/f test/b/bb/f test/d/f'
    /// ```
    #[test]
    fn test_re_tree_compatibility() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();

        let mut builder = ActionDirectoryBuilder::empty();
        for p in &["a/aa/f", "a/aaa/f", "b/bb/f", "d/f"] {
            insert_file(
                &mut builder,
                path(p),
                FileMetadata::empty(digest_config.cas_digest_config()),
            )?;
        }
        let dir = builder.fingerprint(digest_config.as_directory_serializer());
        let tree = directory_to_re_tree(&dir);
        let tree = proto_serialize(&tree);
        let digest = FileDigest::from_content(&tree, digest_config.cas_digest_config());
        assert_eq!(
            digest.to_string(),
            "28e2632abbdc186ddd3ce26d2faf2da40c4c2695:521"
        );

        Ok(())
    }

    #[cfg(target_os = "windows")]
    #[test]
    //Test that a symlink created with a windows path doesn't get interpreted as an invalid sylink
    //TODO(lmvasquezg) Update symlinks to store a normalized, OS-independent path
    fn test_unnormalized_symlinks() -> anyhow::Result<()> {
        let digest_config = DigestConfig::testing_default();

        let mut builder = ActionDirectoryBuilder::empty();
        insert_symlink(
            &mut builder,
            path("d1/f1"),
            Arc::new(Symlink::new("..\\d2\\f2".into())),
        )?;
        insert_file(
            &mut builder,
            path("d2/f2"),
            FileMetadata::empty(digest_config.cas_digest_config()),
        )?;

        extract_artifact_value(&builder, path("d1/f1"), digest_config)?.context("Not value!")?;
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use anyhow::Context as _;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::FileDigestConfig;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_core::directory::DirectoryEntry;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::file_name::FileNameBuf;
use faccess::PathExt;

use crate::directory::new_symlink;
use crate::directory::ActionDirectoryBuilder;
use crate::directory::ActionDirectoryEntry;
use crate::directory::ActionDirectoryMember;

pub fn build_entry_from_disk(
    mut path: AbsNormPathBuf,
    digest_config: FileDigestConfig,
) -> anyhow::Result<Option<ActionDirectoryEntry<ActionDirectoryBuilder>>> {
    // Get file metadata. If the file is missing, ignore it.
    let m = match std::fs::symlink_metadata(&path) {
        Ok(m) => m,
        Err(ref err) if err.kind() == std::io::ErrorKind::NotFound => return Ok(None),
        Err(err) => return Err(err.into()),
    };

    let value = if m.file_type().is_symlink() {
        DirectoryEntry::Leaf(new_symlink(fs_util::read_link(&path)?)?)
    } else if m.is_file() {
        DirectoryEntry::Leaf(ActionDirectoryMember::File(FileMetadata {
            digest: TrackedFileDigest::new(
                FileDigest::from_file(&path, digest_config)?,
                digest_config.as_cas_digest_config(),
            ),
            is_executable: path.executable(),
        }))
    } else if m.is_dir() {
        DirectoryEntry::Dir(build_dir_from_disk(&mut path, digest_config)?)
    } else {
        anyhow::bail!("Path {:?} is of an unknown file type.", path)
    };
    Ok(Some(value))
}

fn build_dir_from_disk(
    disk_path: &mut AbsNormPathBuf,
    digest_config: FileDigestConfig,
) -> anyhow::Result<ActionDirectoryBuilder> {
    let mut builder = ActionDirectoryBuilder::empty();

    for file in fs_util::read_dir(&disk_path)? {
        let file = file?;
        let filetype = file.file_type()?;
        let filename = file.file_name();

        let filename = filename
            .to_str()
            .context("Filename is not UTF-8")
            .and_then(|f| FileNameBuf::try_from(f.to_owned()))
            .with_context(|| format!("Invalid filename: {}", disk_path.display()))?;

        disk_path.push(&filename);

        if filetype.is_dir() {
            let dir = build_dir_from_disk(disk_path, digest_config)?;
            builder.insert(filename, DirectoryEntry::Dir(dir))?;
        } else if filetype.is_symlink() {
            builder.insert(
                filename,
                DirectoryEntry::Leaf(new_symlink(fs_util::read_link(&disk_path)?)?),
            )?;
        } else if filetype.is_file() {
            let metadata = FileMetadata {
                digest: TrackedFileDigest::new(
                    FileDigest::from_file(&disk_path, digest_config)?,
                    digest_config.as_cas_digest_config(),
                ),
                is_executable: file.path().executable(),
            };
            builder.insert(
                filename,
                DirectoryEntry::Leaf(ActionDirectoryMember::File(metadata)),
            )?;
        }
        disk_path.pop();
    }

    Ok(builder)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use dupe::Dupe;

/// Command-level config that can tweak how the executors work.
#[derive(Clone, Dupe, Default)]
pub struct ExecutorGlobalKnobs {
    pub enable_miniperf: bool,

    /// Whether to emit action keys to execution logs (thos are pretty verbose and omitted by
    /// default).
    pub log_action_keys: bool,
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

#![feature(never_type)]
#![feature(trait_alias)]
#![feature(try_blocks)]
#![feature(try_trait_v2)]

pub mod artifact;
pub mod artifact_utils;
pub mod artifact_value;
pub mod bxl;
pub mod digest;
pub mod digest_config;
pub mod directory;
pub mod entry;
pub mod execute;
pub mod knobs;
pub mod materialize;
pub mod output_size;
pub mod path;
pub mod re;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_core::directory::unordered_entry_walk;
use buck2_core::directory::DirectoryEntry;

use crate::artifact_value::ArtifactValue;
use crate::directory::ActionDirectory;
use crate::directory::ActionDirectoryMember;

pub struct OutputCountAndBytes {
    pub count: u64,
    pub bytes: u64,
}
pub trait OutputSize {
    fn calc_output_count_and_bytes(&self) -> OutputCountAndBytes;
}

impl OutputSize for ArtifactValue {
    fn calc_output_count_and_bytes(&self) -> OutputCountAndBytes {
        self.entry().calc_output_count_and_bytes()
    }
}

impl<D> OutputSize for DirectoryEntry<D, ActionDirectoryMember>
where
    D: ActionDirectory,
{
    fn calc_output_count_and_bytes(&self) -> OutputCountAndBytes {
        let mut bytes = 0;
        let mut count = 0;
        let mut walk = unordered_entry_walk(self.as_ref());
        while let Some((_path, entry)) = walk.next() {
            match entry {
                DirectoryEntry::Leaf(ActionDirectoryMember::File(f)) => {
                    bytes += f.digest.size();
                    count += 1;
                }
                _ => {}
            }
        }
        OutputCountAndBytes { count, bytes }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;

pub trait ArtifactDyn: Send + Sync + 'static {
    fn resolve_path(&self, fs: &ArtifactFs) -> anyhow::Result<ProjectRelativePathBuf>;
    fn is_source(&self) -> bool;
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_common::executor_config::PathSeparatorKind;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;

pub struct ExecutorFs<'a> {
    fs: &'a ArtifactFs,
    path_separator: PathSeparatorKind,
}

impl<'a> ExecutorFs<'a> {
    pub fn new(fs: &'a ArtifactFs, path_separator: PathSeparatorKind) -> Self {
        Self { fs, path_separator }
    }

    pub fn fs(&self) -> &ArtifactFs {
        self.fs
    }

    pub fn path_separator(&self) -> PathSeparatorKind {
        self.path_separator
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

pub mod artifact_dyn;
pub mod fs;
pub mod group;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_core::fs::artifact_path_resolver::ArtifactFs;

use crate::artifact::artifact_dyn::ArtifactDyn;
use crate::artifact_value::ArtifactValue;
use crate::directory::ActionDirectoryBuilder;

/// This is like `ArtifactGroupValues`, but without dependency on `Artifact`.
pub trait ArtifactGroupValuesDyn: Send + Sync + 'static {
    fn iter(&self) -> Box<dyn Iterator<Item = (&dyn ArtifactDyn, &ArtifactValue)> + '_>;

    fn add_to_directory(
        &self,
        builder: &mut ActionDirectoryBuilder,
        artifact_fs: &ArtifactFs,
    ) -> anyhow::Result<()>;
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

pub mod artifact_group_values_dyn;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_common::cas_digest::CasDigest;
use buck2_common::cas_digest::CasDigestConfig;
use buck2_common::cas_digest::TrackedCasDigest;
use buck2_common::cas_digest::TrackedCasDigestKind;

pub struct ActionDigestKind {
    _private: (),
}

impl TrackedCasDigestKind for ActionDigestKind {
    fn empty_digest(_config: CasDigestConfig) -> Option<TrackedCasDigest<Self>> {
        // No reason to optimize "empty" actions.
        None
    }
}

pub type ActionDigest = CasDigest<ActionDigestKind>;
pub type TrackedActionDigest = TrackedCasDigest<ActionDigestKind>;


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::collections::HashMap;

use buck2_common::file_ops::TrackedFileDigest;
use dupe::Dupe;
use prost::Message;

use crate::digest_config::DigestConfig;

/// Contains small blobs referenced from action messages (does not include any file contents blobs).
pub struct ActionBlobs(HashMap<TrackedFileDigest, Vec<u8>>);

impl ActionBlobs {
    pub fn new(digest_config: DigestConfig) -> Self {
        // We add empty files to the input that don't exist in disk; so add
        // the empty digest to blobs, as going to disk would fail.
        let mut blobs = HashMap::new();
        blobs.insert(
            TrackedFileDigest::empty(digest_config.cas_digest_config()),
            Vec::new(),
        );
        Self(blobs)
    }

    pub fn add_blob(&mut self, digest: TrackedFileDigest, data: Vec<u8>) {
        self.0.insert(digest, data);
    }

    pub fn add_protobuf_message(
        &mut self,
        m: &impl Message,
        digest_config: DigestConfig,
    ) -> TrackedFileDigest {
        let mut blob = Vec::new();
        // Unwrap is safe because it only fails in OOM conditions, which we pretend don't happen
        m.encode(&mut blob).unwrap();
        let digest = TrackedFileDigest::from_content(&blob, digest_config.cas_digest_config());
        self.0.insert(digest.dupe(), blob);
        digest
    }

    pub fn keys(&self) -> impl Iterator<Item = &TrackedFileDigest> {
        self.0.keys()
    }

    pub fn get(&self, digest: &TrackedFileDigest) -> Option<&Vec<u8>> {
        self.0.get(digest)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::project::ProjectRoot;
use crossbeam_channel::unbounded;
use dice::DiceComputations;
use dice::UserComputationData;
use dupe::Dupe;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use more_futures::cancellation::CancellationContext;
use tokio::sync::oneshot;
use tokio::sync::Semaphore;

#[async_trait]
pub trait BlockingExecutor: Allocative + Send + Sync + 'static {
    /// Execute a blocking I/O operation on the current thread. This should be used sparingly. It
    /// is appropriate to use in cases we are doing a minimal amount of I/O (e.g. writing to just
    /// one file), or where I/O is mixed with other blocking operations.  Those operations run with
    /// fairly high concurrency as they aren't expected to contend with each other.
    async fn execute_dyn_io_inline<'a>(
        &self,
        f: Box<dyn FnOnce() -> anyhow::Result<()> + Send + 'a>,
    ) -> anyhow::Result<()>;

    /// Execute a blocking I/O operation, possibly on a dedicated I/O pool. This should be used as
    /// the default for I/O. The operations executed here must perform _only_ I/O (since if they do
    /// something else they might contend for I/O threads with actual I/O).
    fn execute_io<'a>(
        &self,
        io: Box<dyn IoRequest>,
        cancellations: &'a CancellationContext,
    ) -> BoxFuture<'a, anyhow::Result<()>>;

    /// The size of the queue of pending I/O.
    fn queue_size(&self) -> usize;
}

impl dyn BlockingExecutor {
    pub async fn execute_io_inline<F, T>(&self, f: F) -> anyhow::Result<T>
    where
        F: FnOnce() -> anyhow::Result<T> + Send,
        T: Send,
    {
        let mut res = None;
        self.execute_dyn_io_inline(Box::new(|| {
            res = Some(f()?);
            Ok(())
        }))
        .await?;
        res.context("Inline I/O did not execute")
    }
}

pub trait IoRequest: Send + Sync + 'static {
    fn execute(self: Box<Self>, project_fs: &ProjectRoot) -> anyhow::Result<()>;
}

struct ThreadPoolIoRequest {
    io: Box<dyn IoRequest>,
    sender: oneshot::Sender<anyhow::Result<()>>,
}

#[derive(Allocative)]
pub struct BuckBlockingExecutor {
    #[allocative(skip)]
    io_data_semaphore: Semaphore,
    #[allocative(skip)]
    command_sender: crossbeam_channel::Sender<ThreadPoolIoRequest>,
}

impl BuckBlockingExecutor {
    /// We choose the default concurrency as follows:
    ///
    /// - For operations executed by the thread pool, we choose a fairly low concurrency level.
    /// This is because those operations do exclusively I/O work, and that work consists of
    /// modifying the directory structure of the FS, which scales negatively as soon as you add
    /// more than 4 threads on all systems we care about (sometimes it does so earlier, but for now
    /// 4 is the one-size-fits-all solution we have). D33922298 has benchmark details.
    ///
    /// - For operations that primarily write data, we default to the number of threads on the
    /// host. This is because those operations often have to do CPU bound work to generate the data
    /// they are trying to write, and writing to multiple files doesn't have the negative scaling
    /// issues modifying the directory structure does.
    pub fn default_concurrency(fs: ProjectRoot) -> anyhow::Result<Self> {
        static IO_THREADS: EnvHelper<usize> = EnvHelper::new("BUCK2_IO_THREADS");
        static IO_SEMAPHORE: EnvHelper<usize> = EnvHelper::new("BUCK2_IO_SEMAPHORE");

        let io_threads = IO_THREADS.get_copied()?.unwrap_or(4);
        let io_semaphore = IO_SEMAPHORE.get_copied()?.unwrap_or_else(num_cpus::get);

        let (command_sender, command_receiver) = unbounded();

        for i in 0..io_threads {
            let command_receiver = command_receiver.clone();
            let fs = fs.dupe();
            std::thread::Builder::new()
                .name(format!("buck-io-{}", i))
                .spawn(move || {
                    for ThreadPoolIoRequest { sender, io } in command_receiver.iter() {
                        let res = io.execute(&fs);
                        let _ignored = sender.send(res);
                    }
                })
                .context("Failed to spawn io worker")?;
        }

        Ok(Self {
            io_data_semaphore: Semaphore::new(io_semaphore),
            command_sender,
        })
    }
}

#[async_trait]
impl BlockingExecutor for BuckBlockingExecutor {
    async fn execute_dyn_io_inline<'a>(
        &self,
        f: Box<dyn FnOnce() -> anyhow::Result<()> + Send + 'a>,
    ) -> anyhow::Result<()> {
        let _permit = self
            .io_data_semaphore
            .acquire()
            .await
            .expect("This semaphore is never closed");

        tokio::task::block_in_place(f)
    }

    fn execute_io<'a>(
        &self,
        io: Box<dyn IoRequest>,
        cancellations: &'a CancellationContext,
    ) -> BoxFuture<'a, anyhow::Result<()>> {
        let (sender, receiver) = oneshot::channel();

        // Ignore errors sending as they'll translate to an error receiving once we drop the
        // sender.
        let _ignored = self.command_sender.send(ThreadPoolIoRequest { io, sender });

        cancellations
            .critical_section(|| async move { receiver.await.context("Pool shut down")? })
            .boxed()
    }

    fn queue_size(&self) -> usize {
        self.command_sender.len()
    }
}

pub trait SetBlockingExecutor {
    fn set_blocking_executor(&mut self, exec: Arc<dyn BlockingExecutor>);
}

pub trait HasBlockingExecutor {
    fn get_blocking_executor(&self) -> Arc<dyn BlockingExecutor>;
}

impl SetBlockingExecutor for UserComputationData {
    fn set_blocking_executor(&mut self, exec: Arc<dyn BlockingExecutor>) {
        self.data.set(exec);
    }
}

impl HasBlockingExecutor for DiceComputations {
    fn get_blocking_executor(&self) -> Arc<dyn BlockingExecutor> {
        self.per_transaction_data()
            .data
            .get::<Arc<dyn BlockingExecutor>>()
            .expect("BlockingExecutor should be set")
            .dupe()
    }
}

pub mod testing {
    use super::*;

    #[derive(Allocative)]
    pub struct DummyBlockingExecutor {
        pub fs: ProjectRoot,
    }

    #[async_trait]
    impl BlockingExecutor for DummyBlockingExecutor {
        async fn execute_dyn_io_inline<'a>(
            &self,
            f: Box<dyn FnOnce() -> anyhow::Result<()> + Send + 'a>,
        ) -> anyhow::Result<()> {
            f()
        }

        fn execute_io<'a>(
            &self,
            io: Box<dyn IoRequest>,
            _cancellations: &'a CancellationContext,
        ) -> BoxFuture<'a, anyhow::Result<()>> {
            futures::future::ready(io.execute(&self.fs)).boxed()
        }

        fn queue_size(&self) -> usize {
            0
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use async_trait::async_trait;

use crate::digest_config::DigestConfig;
use crate::execute::action_digest::ActionDigest;
use crate::execute::result::CommandExecutionResult;
use crate::execute::target::CommandExecutionTarget;

pub struct CacheUploadInfo<'a> {
    pub target: &'a dyn CommandExecutionTarget,
    pub action_digest: ActionDigest,
    pub digest_config: DigestConfig,
}

/// A single purpose trait to handle cache uploads
#[async_trait]
pub trait UploadCache: Send + Sync {
    /// Given information about the command and its result, upload the result
    /// and related items to the cache.
    /// Return a boolean to indicate whether a cache upload took place.
    async fn upload(
        &self,
        info: &CacheUploadInfo<'_>,
        execution_result: &CommandExecutionResult,
    ) -> anyhow::Result<bool>;
}

/// A no-op cache uploader for when cache uploading is disabled
pub struct NoOpCacheUploader {}

#[async_trait]
impl UploadCache for NoOpCacheUploader {
    async fn upload(
        &self,
        _info: &CacheUploadInfo<'_>,
        _execution_result: &CommandExecutionResult,
    ) -> anyhow::Result<bool> {
        Ok(false)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::fmt;
use std::sync::Arc;

use async_trait::async_trait;
use derivative::Derivative;
use derive_more::Display;
use dupe::Dupe;
use tokio::sync::Mutex;
use tokio::sync::OwnedMutexGuard;

#[async_trait]
pub trait ClaimManager: Send + Sync + 'static {
    /// Acquire a claim. the Claim mamanger is normally internally mutable and will likely lock out
    /// other executions from acquiring a claim.
    ///
    /// Actions *must* acquire a claim before writing to outputs for an action or declaring
    /// output artifacts in the materializer.
    ///
    /// The expectation is that multiple actions that attempt to write to the same output directory
    /// should not be allowed to obtain claims concurrently.
    async fn claim(self: Box<Self>) -> Box<dyn Claim>;

    /// Inform the manager that producing a result might take longer than expected.
    ///
    /// The manager may use this information to notify other parts of the system. We use this notably
    /// to unblock local execution if RE reports long queueing.
    fn on_result_delayed(&mut self);
}

pub trait Claim: Send + Sync + fmt::Debug + 'static {
    /// Release a claim. This can be done to make it available again after acquiring it. Once the
    /// claim is released, the claim holder should no longer write any output.
    fn release(self: Box<Self>) -> anyhow::Result<()>;
}

#[derive(Debug)]
pub enum ClaimStatus {
    // No claim has been attempted, or one was attempted but abandoned.
    NotClaimed,
    // A claim has been made.
    Claimed,
}

#[derive(Dupe, Clone)]
pub struct MutexClaimManager {
    mutex: Arc<Mutex<ClaimStatus>>,
}

impl MutexClaimManager {
    pub fn new() -> Self {
        Self {
            mutex: Arc::new(Mutex::new(ClaimStatus::NotClaimed)),
        }
    }
}

#[async_trait]
impl ClaimManager for MutexClaimManager {
    async fn claim(self: Box<Self>) -> Box<dyn Claim> {
        let mut guard = self.mutex.dupe().lock_owned().await;
        match *guard {
            ClaimStatus::NotClaimed => {}
            ClaimStatus::Claimed => futures::future::pending().await,
        }
        *guard = ClaimStatus::Claimed;
        Box::new(MutexClaim { guard })
    }

    fn on_result_delayed(&mut self) {}
}

#[derive(Display, Derivative)]
#[display(fmt = "MutexClaim")]
#[derivative(Debug)]
pub struct MutexClaim {
    // No point in printing this as it will *always* be Claimed.
    #[derivative(Debug = "ignore")]
    guard: OwnedMutexGuard<ClaimStatus>,
}

impl Claim for MutexClaim {
    fn release(mut self: Box<Self>) -> anyhow::Result<()> {
        *self.guard = ClaimStatus::NotClaimed;
        Ok(())
    }
}

#[cfg(test)]
mod tests {
    use std::task::Poll;

    use assert_matches::assert_matches;

    use super::*;

    #[tokio::test]
    async fn test_mutex_claim_release() {
        let claim_manager = MutexClaimManager::new();
        let claim = Box::new(claim_manager.dupe()).claim().await;

        let claim2 = Box::new(claim_manager.dupe()).claim();
        futures::pin_mut!(claim2);
        assert_matches!(futures::poll!(claim2.as_mut()), Poll::Pending);

        claim.release().expect("Can release claim");
        assert_matches!(futures::poll!(claim2.as_mut()), Poll::Ready(..));
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use anyhow::Context;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;

use crate::execute::blocking::IoRequest;

/// IoRequest we dispatch to the blocking executor to clear output paths.
pub struct CleanOutputPaths {
    pub paths: Vec<ProjectRelativePathBuf>,
}

impl CleanOutputPaths {
    pub fn clean<'a>(
        paths: impl Iterator<Item = &'a ProjectRelativePath>,
        fs: &'a ProjectRoot,
    ) -> anyhow::Result<()> {
        for path in paths {
            cleanup_path(fs, path)
                .with_context(|| format!("Error cleaning up output path `{}`", path))?;
        }
        Ok(())
    }
}

#[tracing::instrument(level = "debug", skip(fs), fields(path = %path))]
pub fn cleanup_path(fs: &ProjectRoot, path: &ProjectRelativePath) -> anyhow::Result<()> {
    // This will remove the path if it exists.
    fs.remove_path_recursive(path)?;

    let path = fs.resolve(path);
    let mut path: &AbsNormPath = &path;

    // Be aware of T85589819 - the parent directory might already exist, but as a _file_.  It might
    // be even worse, it might be 2 parents up, which will cause create_dir to fail when we try to
    // execute. So, we walk up the tree until we either find a dir we're happy with, or a file we
    // can delete. It's safe to delete this file because we know it doesn't overlap with a current
    // output, or that check would have failed, so it must be a stale file.
    loop {
        path = match path.parent() {
            Some(path) => path,
            None => {
                return Err(anyhow::anyhow!(
                    "Internal Error: reached root before finding a directory that exists!"
                ));
            }
        };

        match fs_util::symlink_metadata_if_exists(path) {
            Ok(Some(m)) => {
                if m.is_dir() {
                    // It's a dir, no need to go further, and no need to delete.
                    tracing::trace!(path = %path, "skip (is dir)");
                } else {
                    // There was a file or a symlink, so it's safe to delete and then we can exit
                    // because we'll be able to create a dir here.
                    tracing::trace!(path = %path, "remove_file");
                    fs_util::remove_file(path)?;
                }
                return Ok(());
            }
            Ok(None) if cfg!(unix) => {
                // If we get ENOENT that guarantees there is no file on the path. If there was
                // one, we would get ENOTDIR. TODO (T123279320) This probably works on Windows,
                // but it wasn't tested there.
                //
                // On non-Unix we don't have this optimization. Recursing all the way up
                // until we find the first dir (or file to delete) is fine. There will
                // eventually be *a* directory (at buck-out, then another one at the empty
                // directory, which is our cwd, and should exist by now).
                tracing::trace!(path = %path, "skip (ENOENT)");
                return Ok(());
            }
            Ok(None) => {
                tracing::trace!(path = %path, "continue (ENOENT)");
            }
            Err(e) => {
                // Continue going up. Eventually we should reach the output directory, which should
                // exist.
                tracing::trace!(path = %path, "continue (error: {:#})", e);
            }
        }
    }
}

impl IoRequest for CleanOutputPaths {
    fn execute(self: Box<Self>, project_fs: &ProjectRoot) -> anyhow::Result<()> {
        Self::clean(self.paths.iter().map(AsRef::as_ref), project_fs)
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::ops::ControlFlow;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use buck2_common::executor_config::CommandGenerationOptions;
use buck2_common::executor_config::OutputPathsBehavior;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_core::directory::FingerprintedDirectory;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use dupe::Dupe;
use more_futures::cancellation::CancellationContext;
use remote_execution as RE;
use sorted_vector_map::SortedVectorMap;

use crate::artifact::fs::ExecutorFs;
use crate::digest::CasDigestToReExt;
use crate::digest_config::DigestConfig;
use crate::execute::blobs::ActionBlobs;
use crate::execute::cache_uploader::CacheUploadInfo;
use crate::execute::cache_uploader::UploadCache;
use crate::execute::executor_stage;
use crate::execute::manager::CommandExecutionManager;
use crate::execute::prepared::PreparedAction;
use crate::execute::prepared::PreparedCommand;
use crate::execute::prepared::PreparedCommandExecutor;
use crate::execute::prepared::PreparedCommandOptionalExecutor;
use crate::execute::request::CommandExecutionInput;
use crate::execute::request::CommandExecutionRequest;
use crate::execute::request::ExecutorPreference;
use crate::execute::request::OutputType;
use crate::execute::result::CommandExecutionMetadata;
use crate::execute::result::CommandExecutionResult;

#[derive(Copy, Dupe, Clone, Debug, PartialEq, Eq)]
pub struct ActionExecutionTimingData {
    pub wall_time: Duration,
}

impl Default for ActionExecutionTimingData {
    fn default() -> Self {
        Self {
            wall_time: Duration::ZERO,
        }
    }
}

impl From<CommandExecutionMetadata> for ActionExecutionTimingData {
    fn from(command: CommandExecutionMetadata) -> Self {
        Self {
            wall_time: command.wall_time,
        }
    }
}

#[derive(Clone, Dupe)]
pub struct CommandExecutor(Arc<CommandExecutorData>);

struct CommandExecutorData {
    inner: Arc<dyn PreparedCommandExecutor>,
    cache_checker: Arc<dyn PreparedCommandOptionalExecutor>,
    artifact_fs: ArtifactFs,
    options: CommandGenerationOptions,
    re_platform: RE::Platform,
    enforce_re_timeouts: bool,
    cache_uploader: Arc<dyn UploadCache>,
}

impl CommandExecutor {
    pub fn new(
        inner: Arc<dyn PreparedCommandExecutor>,
        cache_checker: Arc<dyn PreparedCommandOptionalExecutor>,
        cache_uploader: Arc<dyn UploadCache>,
        artifact_fs: ArtifactFs,
        options: CommandGenerationOptions,
        re_platform: RE::Platform,
        enforce_re_timeouts: bool,
    ) -> Self {
        Self(Arc::new(CommandExecutorData {
            inner,
            cache_checker,
            artifact_fs,
            options,
            re_platform,
            enforce_re_timeouts,
            cache_uploader,
        }))
    }

    pub fn fs(&self) -> &ArtifactFs {
        &self.0.artifact_fs
    }

    pub fn executor_fs(&self) -> ExecutorFs {
        ExecutorFs::new(&self.0.artifact_fs, self.0.options.path_separator)
    }

    /// Check if the action can be served by the action cache.
    pub async fn action_cache(
        &self,
        manager: CommandExecutionManager,
        prepared_command: &PreparedCommand<'_, '_>,
        cancellations: &CancellationContext<'_>,
    ) -> ControlFlow<CommandExecutionResult, CommandExecutionManager> {
        self.0
            .cache_checker
            .maybe_execute(prepared_command, manager, cancellations)
            .await
    }

    pub async fn cache_upload(
        &self,
        info: &CacheUploadInfo<'_>,
        execution_result: &CommandExecutionResult,
    ) -> anyhow::Result<bool> {
        self.0.cache_uploader.upload(info, execution_result).await
    }

    /// Execute a command.
    ///
    /// This intentionally does not return a Result since we want to capture information about the
    /// execution even if there are errors. Any errors can be propagated by converting them
    /// to a result with CommandExecutionManager::error.
    pub async fn exec_cmd(
        &self,
        manager: CommandExecutionManager,
        prepared_command: &PreparedCommand<'_, '_>,
        cancellations: &CancellationContext<'_>,
    ) -> CommandExecutionResult {
        self.0
            .inner
            .exec_cmd(prepared_command, manager, cancellations)
            .await
    }

    pub fn is_local_execution_possible(&self, executor_preference: ExecutorPreference) -> bool {
        self.0
            .inner
            .is_local_execution_possible(executor_preference)
    }

    pub fn prepare_action(
        &self,
        request: &CommandExecutionRequest,
        digest_config: DigestConfig,
    ) -> anyhow::Result<PreparedAction> {
        executor_stage(buck2_data::PrepareAction {}, || {
            let input_digest = request.paths().input_directory().fingerprint();

            let action_metadata_blobs = request.inputs().iter().filter_map(|x| match x {
                CommandExecutionInput::Artifact(_) => None,
                CommandExecutionInput::ActionMetadata(metadata) => {
                    Some((metadata.data.clone(), metadata.digest.dupe()))
                }
            });
            let action = re_create_action(
                request.all_args_vec(),
                request.paths().output_paths(),
                request.working_directory().map(|p| p.as_str().to_owned()),
                request.env(),
                input_digest,
                action_metadata_blobs,
                if self.0.enforce_re_timeouts {
                    request.timeout()
                } else {
                    None
                },
                self.0.re_platform.clone(),
                false,
                digest_config,
                self.0.options.output_paths_behavior,
            )?;

            anyhow::Ok(action)
        })
    }
}

fn re_create_action(
    args: Vec<String>,
    outputs: &[(ProjectRelativePathBuf, OutputType)],
    workdir: Option<String>,
    environment: &SortedVectorMap<String, String>,
    input_digest: &TrackedFileDigest,
    blobs: impl Iterator<Item = (Vec<u8>, TrackedFileDigest)>,
    timeout: Option<Duration>,
    platform: RE::Platform,
    do_not_cache: bool,
    digest_config: DigestConfig,
    output_paths_behavior: OutputPathsBehavior,
) -> anyhow::Result<PreparedAction> {
    let mut command = RE::Command {
        arguments: args,
        platform: Some(platform),
        working_directory: workdir.unwrap_or_default(),
        environment_variables: environment
            .iter()
            .map(|(k, v)| RE::EnvironmentVariable {
                name: (*k).clone(),
                value: (*v).clone(),
            })
            .collect(),
        ..Default::default()
    };

    match output_paths_behavior {
        OutputPathsBehavior::Compatibility => {
            for (output, output_type) in outputs {
                let path = output.as_str().to_owned();

                match output_type {
                    OutputType::FileOrDirectory => {
                        command.output_files.push(path.clone());
                        command.output_directories.push(path);
                    }
                    OutputType::File => command.output_files.push(path),
                    OutputType::Directory => command.output_directories.push(path),
                }
            }
        }
        OutputPathsBehavior::Strict => {
            for (output, output_type) in outputs {
                let path = output.as_str().to_owned();

                match output_type {
                    OutputType::FileOrDirectory => {
                        command.output_files.push(path);
                    }
                    OutputType::File => command.output_files.push(path),
                    OutputType::Directory => command.output_directories.push(path),
                }
            }
        }
        OutputPathsBehavior::OutputPaths => {
            #[cfg(fbcode_build)]
            {
                return Err(anyhow::anyhow!(
                    "output_paths is not supported in fbcode_build"
                ));
            }

            #[cfg(not(fbcode_build))]
            {
                for (output, _output_type) in outputs {
                    command.output_paths.push(output.as_str().to_owned());
                }
            }
        }
    }

    let mut prepared_blobs = ActionBlobs::new(digest_config);
    for (data, digest) in blobs {
        prepared_blobs.add_blob(digest, data);
    }
    let action = RE::Action {
        input_root_digest: Some(input_digest.to_grpc()),
        command_digest: Some(
            prepared_blobs
                .add_protobuf_message(&command, digest_config)
                .to_grpc(),
        ),
        timeout: timeout
            .map(|t| t.try_into())
            .transpose()
            .context("Cannot convert timeout to GRPC")?,
        do_not_cache,
        ..Default::default()
    };

    let action = prepared_blobs.add_protobuf_message(&action, digest_config);
    Ok(PreparedAction {
        action: action.data().dupe().coerce(),
        blobs: prepared_blobs,
        platform: command
            .platform
            .expect("We did put a platform a few lines up"),
    })
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

//! Attaching command execution related data to Dice

use std::sync::Arc;

use buck2_common::executor_config::CommandExecutorConfig;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use dice::DiceComputations;
use dice::DiceData;
use dice::UserComputationData;
use dupe::Dupe;
use remote_execution as RE;

use super::prepared::PreparedCommandOptionalExecutor;
use crate::execute::cache_uploader::UploadCache;
use crate::execute::prepared::PreparedCommandExecutor;
use crate::re::manager::ManagedRemoteExecutionClient;

pub struct CommandExecutorResponse {
    pub executor: Arc<dyn PreparedCommandExecutor>,
    pub platform: RE::Platform,
    pub cache_checker: Arc<dyn PreparedCommandOptionalExecutor>,
    pub cache_uploader: Arc<dyn UploadCache>,
}

pub trait SetCommandExecutor {
    fn set_command_executor(&mut self, init: Box<dyn HasCommandExecutor + Send + Sync + 'static>);
}

pub trait HasCommandExecutor {
    fn get_command_executor(
        &self,
        artifact_fs: &ArtifactFs,
        config: &CommandExecutorConfig,
    ) -> anyhow::Result<CommandExecutorResponse>;
}

impl SetCommandExecutor for UserComputationData {
    fn set_command_executor(
        &mut self,
        delegate: Box<dyn HasCommandExecutor + Send + Sync + 'static>,
    ) {
        self.data.set(HasCommandExecutorHolder { delegate })
    }
}

impl HasCommandExecutor for DiceComputations {
    fn get_command_executor(
        &self,
        artifact_fs: &ArtifactFs,
        config: &CommandExecutorConfig,
    ) -> anyhow::Result<CommandExecutorResponse> {
        let holder = self
            .per_transaction_data()
            .data
            .get::<HasCommandExecutorHolder>()
            .expect("CommandExecutorDelegate should be set");
        holder.delegate.get_command_executor(artifact_fs, config)
    }
}

struct HasCommandExecutorHolder {
    delegate: Box<dyn HasCommandExecutor + Send + Sync + 'static>,
}

pub trait HasFallbackExecutorConfig {
    fn get_fallback_executor_config(&self) -> &Arc<CommandExecutorConfig>;
}

impl HasFallbackExecutorConfig for DiceComputations {
    fn get_fallback_executor_config(&self) -> &Arc<CommandExecutorConfig> {
        self.per_transaction_data()
            .data
            .get::<Arc<CommandExecutorConfig>>()
            .expect("CommandExecutorConfig should be set")
    }
}

pub fn set_fallback_executor_config(data: &mut DiceData, config: Arc<CommandExecutorConfig>) {
    data.set(config)
}

pub trait SetReClient {
    fn set_re_client(&mut self, re_client: ManagedRemoteExecutionClient);
}

pub trait GetReClient {
    fn get_re_client(&self) -> ManagedRemoteExecutionClient;
}

impl SetReClient for UserComputationData {
    fn set_re_client(&mut self, re_client: ManagedRemoteExecutionClient) {
        self.data.set(re_client);
    }
}

impl GetReClient for UserComputationData {
    fn get_re_client(&self) -> ManagedRemoteExecutionClient {
        self.data
            .get::<ManagedRemoteExecutionClient>()
            .expect("Materializer should be set")
            .dupe()
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::ffi::OsString;

use dupe::Dupe;
use once_cell::sync::OnceCell;

#[cfg(fbcode_build)]
const ENV_ALLOW_LIST: &[&str] = test_env_allowlist::LEGACY_TESTPILOT_ALLOW_LIST;

#[cfg(all(unix, not(fbcode_build)))]
const ENV_ALLOW_LIST: &[&str] = &["PATH", "USER", "LOGNAME", "HOME", "TMPDIR"];

// The standard (built-in) variables.
// https://ss64.com/nt/syntax-variables.html
#[cfg(all(windows, not(fbcode_build)))]
const ENV_ALLOW_LIST: &[&str] = &[
    "ALLUSERSPROFILE",
    "APPDATA",
    "COMPUTERNAME",
    "COMSPEC",
    "CommonProgramFiles",
    "CommonProgramFiles(x86)",
    "HOMEDRIVE",
    "HOMEPATH",
    "LOCALAPPDATA",
    "NUMBER_OF_PROCESSORS",
    "OS",
    "PATH",
    "PATHEXT",
    "PROCESSOR_ARCHITECTURE",
    "PROCESSOR_ARCHITEW6432",
    "PROCESSOR_IDENTIFIER",
    "PROCESSOR_LEVEL",
    "PROCESSOR_REVISION",
    "PSModulePath",
    "ProgramData",
    "ProgramFiles",
    "ProgramFiles(x86)",
    "ProgramW6432",
    "Public",
    "SYSTEMDRIVE",
    "SYSTEMROOT",
    "TEMP",
    "TMP",
    "USERDOMAIN",
    "USERNAME",
    "USERPROFILE",
    "UserDnsDomain",
    "WINDIR",
];

#[derive(Copy, Clone, Dupe, Debug)]
pub struct EnvironmentInheritance {
    clear: bool,
    values: &'static [(&'static str, OsString)],
    exclusions: &'static [&'static str],
}

impl EnvironmentInheritance {
    pub fn test_allowlist() -> Self {
        // This is made to be a list of lists in case we want to include lists from different
        // provenances, like the test_env_allowlist::ENV_LIST_HACKY.
        let allowlists = &[ENV_ALLOW_LIST];

        // We create this *once* since getenv is actually not cheap (being O(n) of the environment
        // size).
        static TEST_CELL: OnceCell<Vec<(&'static str, OsString)>> = OnceCell::new();

        let values = TEST_CELL.get_or_init(|| {
            let mut ret = Vec::new();
            for list in allowlists.iter() {
                for key in list.iter() {
                    if let Some(value) = std::env::var_os(key) {
                        ret.push((*key, value));
                    }
                }
            }
            ret
        });

        Self {
            clear: true,
            values,
            exclusions: &[],
        }
    }

    /// Exclude some vars that are known to cause issues. In an ideal world we should do a
    /// migration to lock this down everywhere.
    pub fn local_command_exclusions() -> Self {
        Self {
            clear: false,
            values: &[],
            exclusions: &[
                "PYTHONPATH",
                "PYTHONHOME",
                "PYTHONSTARTUP",
                "LD_LIBRARY_PATH",
                "LD_PRELOAD",
            ],
        }
    }

    pub fn empty() -> Self {
        Self {
            values: &[],
            exclusions: &[],
            clear: true,
        }
    }

    pub fn values(&self) -> impl Iterator<Item = (&'static str, &'static OsString)> {
        self.values.iter().map(|(k, v)| (*k, v))
    }

    pub fn exclusions(&self) -> impl Iterator<Item = &'static str> {
        self.exclusions.iter().copied()
    }

    pub fn clear(&self) -> bool {
        self.clear
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use buck2_common::file_ops::FileMetadata;
use buck2_core::directory::DirectoryEntry;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use dupe::Dupe;

use crate::directory::ActionDirectoryBuilder;
use crate::directory::ActionDirectoryMember;
use crate::execute::request::CommandExecutionInput;

pub fn inputs_directory(
    inputs: &[CommandExecutionInput],
    fs: &ArtifactFs,
) -> anyhow::Result<ActionDirectoryBuilder> {
    let mut builder = ActionDirectoryBuilder::empty();
    for input in inputs {
        match input {
            CommandExecutionInput::Artifact(group) => {
                group.add_to_directory(&mut builder, fs)?;
            }
            CommandExecutionInput::ActionMetadata(metadata) => {
                let path = fs.buck_out_path_resolver().resolve_gen(&metadata.path);
                builder.insert(
                    &path,
                    DirectoryEntry::Leaf(ActionDirectoryMember::File(FileMetadata {
                        digest: metadata.digest.dupe(),
                        is_executable: false,
                    })),
                )?;
            }
        };
    }
    Ok(builder)
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use derive_more::Display;
use sorted_vector_map::SortedVectorMap;

use crate::execute::action_digest::ActionDigest;

#[derive(Debug, Display, Clone)]
pub enum CommandExecutionKind {
    /// This action was executed locally.
    #[display(fmt = "local")]
    Local {
        // Even though this did not run on RE, we still produced this, so we might as well report
        // it.
        digest: ActionDigest,
        command: Vec<String>,
        env: SortedVectorMap<String, String>,
    },
    /// This action was executed via a remote executor.
    #[display(fmt = "remote")]
    Remote { digest: ActionDigest },
    /// This action was served by the action cache and not executed.
    #[display(fmt = "action_cache")]
    ActionCache { digest: ActionDigest },
    /// This action would have executed via a local worker but failed during worker initialization.
    #[display(fmt = "worker_init")]
    LocalWorkerInit {
        command: Vec<String>,
        env: SortedVectorMap<String, String>,
    },
    /// This action was executed via a local worker.
    #[display(fmt = "worker")]
    LocalWorker {
        digest: ActionDigest,
        command: Vec<String>,
        env: SortedVectorMap<String, String>,
        fallback_exe: Vec<String>,
    },
}

impl CommandExecutionKind {
    pub fn as_enum(&self) -> buck2_data::ActionExecutionKind {
        match self {
            Self::Local { .. } => buck2_data::ActionExecutionKind::Local,
            Self::LocalWorker { .. } | Self::LocalWorkerInit { .. } => {
                buck2_data::ActionExecutionKind::LocalWorker
            }
            Self::Remote { .. } => buck2_data::ActionExecutionKind::Remote,
            Self::ActionCache { .. } => buck2_data::ActionExecutionKind::ActionCache,
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

use std::sync::Arc;
use std::time::Duration;

use buck2_common::liveliness_observer::LivelinessObserver;
use buck2_events::dispatch::EventDispatcher;
use indexmap::IndexMap;

use crate::artifact_value::ArtifactValue;
use crate::execute::claim::Claim;
use crate::execute::claim::ClaimManager;
use crate::execute::kind::CommandExecutionKind;
use crate::execute::output::CommandStdStreams;
use crate::execute::request::CommandExecutionOutput;
use crate::execute::result::CommandExecutionMetadata;
use crate::execute::result::CommandExecutionReport;
use crate::execute::result::CommandExecutionResult;
use crate::execute::result::CommandExecutionStatus;

trait CommandExecutionManagerLike: Sized {
    /// Create a new Command execution result.
    fn result(
        self,
        status: CommandExecutionStatus,
        outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
        std_streams: CommandStdStreams,
        exit_code: Option<i32>,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult;
}

/// This tracker helps track the information that will go into the BuckCommandExecutionMetadata
pub struct CommandExecutionManager {
    pub claim_manager: Box<dyn ClaimManager>,
    pub events: EventDispatcher,
    pub liveliness_observer: Arc<dyn LivelinessObserver>,
}

impl CommandExecutionManager {
    pub fn new(
        claim_manager: Box<dyn ClaimManager>,
        events: EventDispatcher,
        liveliness_observer: Arc<dyn LivelinessObserver>,
    ) -> Self {
        Self {
            claim_manager,
            events,
            liveliness_observer,
        }
    }

    /// Acquire a claim. This might never return if the claim has been taken.
    pub async fn claim(self) -> CommandExecutionManagerWithClaim {
        let claim = self.claim_manager.claim().await;

        CommandExecutionManagerWithClaim {
            claim,
            events: self.events,
            liveliness_observer: self.liveliness_observer,
        }
    }

    pub fn on_result_delayed(&mut self) {
        self.claim_manager.on_result_delayed();
    }

    pub fn cancel(self) -> CommandExecutionResult {
        self.result(
            CommandExecutionStatus::Cancelled,
            IndexMap::new(),
            Default::default(),
            None,
            CommandExecutionMetadata::default(),
        )
    }
}

impl CommandExecutionManagerLike for CommandExecutionManager {
    fn result(
        self,
        status: CommandExecutionStatus,
        outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
        std_streams: CommandStdStreams,
        exit_code: Option<i32>,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult {
        CommandExecutionResult {
            outputs,
            report: CommandExecutionReport {
                claim: None,
                status,
                timing,
                std_streams,
                exit_code,
            },
            rejected_execution: None,
            did_cache_upload: false,
            eligible_for_full_hybrid: false,
        }
    }
}

pub struct CommandExecutionManagerWithClaim {
    pub events: EventDispatcher,
    pub liveliness_observer: Arc<dyn LivelinessObserver>,
    claim: Box<dyn Claim>,
}

/// Like CommandExecutionManager but provides access to things that are only allowed with a Claim;
impl CommandExecutionManagerWithClaim {
    /// Explicitly requires a Claim here to help implementors remember to claim things since a
    /// command can't be successful without making local changes.
    pub fn success(
        self,
        execution_kind: CommandExecutionKind,
        outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
        std_streams: CommandStdStreams,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult {
        self.result(
            CommandExecutionStatus::Success { execution_kind },
            outputs,
            std_streams,
            Some(0),
            timing,
        )
    }

    pub fn cancel_claim(self) -> CommandExecutionResult {
        self.result(
            CommandExecutionStatus::Cancelled,
            IndexMap::new(),
            Default::default(),
            None,
            CommandExecutionMetadata::default(),
        )
    }
}

impl CommandExecutionManagerLike for CommandExecutionManagerWithClaim {
    fn result(
        self,
        status: CommandExecutionStatus,
        outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
        std_streams: CommandStdStreams,
        exit_code: Option<i32>,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult {
        CommandExecutionResult {
            outputs,
            report: CommandExecutionReport {
                claim: Some(self.claim),
                status,
                timing,
                std_streams,
                exit_code,
            },
            rejected_execution: None,
            did_cache_upload: false,
            eligible_for_full_hybrid: false,
        }
    }
}

pub trait CommandExecutionManagerExt: Sized {
    fn failure(
        self,
        execution_kind: CommandExecutionKind,
        outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
        std_streams: CommandStdStreams,
        exit_code: Option<i32>,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult;

    fn timeout(
        self,
        execution_kind: CommandExecutionKind,
        duration: Duration,
        std_streams: CommandStdStreams,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult;

    fn error(self, stage: &'static str, error: impl Into<anyhow::Error>) -> CommandExecutionResult;
}

impl<T> CommandExecutionManagerExt for T
where
    T: CommandExecutionManagerLike,
{
    fn failure(
        self,
        execution_kind: CommandExecutionKind,
        outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
        std_streams: CommandStdStreams,
        exit_code: Option<i32>,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult {
        self.result(
            CommandExecutionStatus::Failure { execution_kind },
            outputs,
            std_streams,
            exit_code,
            timing,
        )
    }

    fn timeout(
        self,
        execution_kind: CommandExecutionKind,
        duration: Duration,
        std_streams: CommandStdStreams,
        timing: CommandExecutionMetadata,
    ) -> CommandExecutionResult {
        self.result(
            CommandExecutionStatus::TimedOut {
                duration,
                execution_kind,
            },
            IndexMap::new(),
            std_streams,
            None,
            timing,
        )
    }

    fn error(self, stage: &'static str, error: impl Into<anyhow::Error>) -> CommandExecutionResult {
        self.result(
            CommandExecutionStatus::Error {
                stage,
                error: error.into(),
            },
            IndexMap::new(),
            Default::default(),
            None,
            CommandExecutionMetadata::default(),
        )
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This source code is licensed under both the MIT license found in the
 * LICENSE-MIT file in the root directory of this source tree and the Apache
 * License, Version 2.0 found in the LICENSE-APACHE file in the root directory
 * of this source tree.
 */

pub mod action_digest;
pub mod blobs;
pub mod blocking;
pub mod cache_uploader;
pub mod claim;
pub mod clean_output_paths;
pub mod command_executor;
pub mod dice_data;
pub mod environment_inheritance;
pub mod inputs_directory;
pub mod kind;
pub mod manager;
pub mod output;
pub mod prepared;
pub mod request;
pub mod result;
pub mod target;
pub mod testing_dry_run;

use std::future::Future;

use buck2_events::dispatch::span;
use buck2_events::dispatch::span_async;
use futures::FutureExt;

pub fn executor_stage_async<F: Future>(
    stage: impl Into<buck2_data::executor_stage_start::Stage>,
    f: F,
) -> impl Future<Output = <F as Future>::Output> {
    // We avoid using `async fn` or `async move` here to avoid doubling the
    // future size. See https://github.com/rust-lang/rust/issues/62958
    let event = buck2_data::ExecutorStageStart {
        stage: Some(stage.into()),
    };
    span_async(event, f.map(|v| (v, buck2_data::ExecutorStageEnd {})))
}

pub fn executor_stage<F, R>(stage: impl Into<buck2_data::executor_stage_start::Stage>, f: F) -> R
where
    F: FnOnce() -> R,
{
    let event = buck2_data::ExecutorStageStart {
        stage: Some(stage.into()),
    };
    span(event, || {
        let r = f();
        (r, buck2_data::ExecutorStageEnd {})
    })
}