/*
 * 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::fmt::Debug;

use anyhow::Context;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_common::file_ops::FileDigest;
use futures::future;
use remote_execution::TDigest;

use crate::digest::CasDigestConversionResultExt;
use crate::digest::CasDigestFromReExt;
use crate::digest_config::DigestConfig;
use crate::re::manager::ManagedRemoteExecutionClient;
use crate::re::streams::RemoteCommandStdStreams;

/// A pair of streams.
#[allow(clippy::manual_non_exhaustive)]
#[derive(Debug)]
pub struct StdStreamPair<T> {
    pub stdout: T,
    pub stderr: T,
}

#[derive(Clone)]
pub enum ReStdStream {
    /// Raw bytes received inline from RE.
    Raw(Vec<u8>),

    /// Output not available inline, we have a digest.
    Digest(TDigest),

    /// This output was not available inline, but was prefetched. The prefetch might be lossy. We
    /// have a digest to access the full data if needed.
    PrefetchedLossy { data: String, digest: TDigest },

    /// There was no output made available by RE.
    None,
}

impl ReStdStream {
    pub fn new(raw: Option<Vec<u8>>, digest: Option<TDigest>) -> Self {
        match (raw, digest) {
            (Some(raw), _) if !raw.is_empty() => Self::Raw(raw),
            (_, Some(digest)) => Self::Digest(digest),
            (_, None) => Self::None,
        }
    }

    pub fn into_raw_or_digest(self) -> (Option<Vec<u8>>, Option<TDigest>) {
        match self {
            Self::Raw(raw) => (Some(raw), None),
            Self::Digest(digest) | Self::PrefetchedLossy { digest, .. } => (None, Some(digest)),
            Self::None => (None, None),
        }
    }

    pub(crate) async fn to_lossy(
        &self,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> String {
        // 4MBs seems like a reasonably large volume of output. There is no research or science behind
        // this number.
        const MAX_STREAM_DOWNLOAD_SIZE: i64 = 4 * 1024 * 1024;

        match self {
            Self::Raw(raw) => String::from_utf8_lossy(raw).into_owned(),
            Self::Digest(digest) if digest.size_in_bytes <= MAX_STREAM_DOWNLOAD_SIZE => {
                match client.download_blob(digest, use_case).await {
                    Ok(bytes) => String::from_utf8_lossy(&bytes).to_string(),
                    Err(e) => {
                        tracing::warn!("Failed to download action stderr: {:#}", e);
                        format!(
                            "Result could not be downloaded - to view type `frecli cas download-blob {}`",
                            FileDigest::from_re(digest, digest_config).as_display(),
                        )
                    }
                }
            }
            Self::PrefetchedLossy { data, .. } => data.clone(),
            Self::Digest(digest) => {
                format!(
                    "Result too large to display - to view type `frecli cas download-blob {}`",
                    FileDigest::from_re(digest, digest_config).as_display(),
                )
            }
            Self::None => String::new(),
        }
    }

    pub(crate) async fn into_bytes(
        self,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) -> anyhow::Result<Vec<u8>> {
        match self {
            Self::Raw(raw) => Ok(raw),
            Self::Digest(digest) | Self::PrefetchedLossy { digest, .. } => {
                let bytes = client
                    .download_blob(&digest, use_case)
                    .await
                    .with_context(|| {
                        format!(
                            "Error downloading from {}",
                            FileDigest::from_re(&digest, digest_config).as_display()
                        )
                    })?;
                Ok(bytes)
            }
            Self::None => Ok(Vec::new()),
        }
    }

    /// Prefetch the output, if relevant.
    pub(crate) async fn prefetch_lossy(
        &mut self,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
        digest_config: DigestConfig,
    ) {
        if let Self::Digest(digest) = &self {
            let data = self.to_lossy(client, use_case, digest_config).await;
            *self = Self::PrefetchedLossy {
                data,
                digest: digest.clone(),
            };
        }
    }
}

impl fmt::Display for ReStdStream {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        match self {
            Self::Raw(raw) => {
                write!(fmt, "raw = `{}`", String::from_utf8_lossy(raw))?;
            }
            Self::Digest(digest) | Self::PrefetchedLossy { digest, .. } => {
                write!(fmt, "digest = `{}`", digest,)?;
            }
            Self::None => {
                write!(fmt, "none")?;
            }
        }

        Ok(())
    }
}

impl fmt::Debug for ReStdStream {
    fn fmt(&self, fmt: &mut fmt::Formatter) -> fmt::Result {
        write!(fmt, "ReStdStream({})", self)
    }
}

#[derive(Debug, derive_more::From, Clone)]
pub enum CommandStdStreams {
    Local { stdout: Vec<u8>, stderr: Vec<u8> },

    Remote(RemoteCommandStdStreams),

    Empty,
}

impl Default for CommandStdStreams {
    fn default() -> Self {
        Self::Empty
    }
}

impl CommandStdStreams {
    /// Access this data as lossy stdout / stderr. This is designed for human consumption.
    pub async fn to_lossy(&self) -> StdStreamPair<String> {
        match self {
            Self::Local { stdout, stderr } => StdStreamPair {
                stdout: String::from_utf8_lossy(stdout).into_owned(),
                stderr: String::from_utf8_lossy(stderr).into_owned(),
            },
            Self::Remote(remote) => {
                let (stdout, stderr) =
                    future::join(remote.to_lossy_stdout(), remote.to_lossy_stderr()).await;
                StdStreamPair { stdout, stderr }
            }
            Self::Empty => StdStreamPair {
                stdout: String::new(),
                stderr: String::new(),
            },
        }
    }

    /// Access this data lossily, but only retain stderr.
    pub async fn to_lossy_stderr(&self) -> String {
        match self {
            Self::Local { stderr, .. } => String::from_utf8_lossy(stderr).into_owned(),
            Self::Remote(remote) => remote.to_lossy_stderr().await,
            Self::Empty => String::new(),
        }
    }

    /// Access the raw data. This is suitable for machine consumption. This will fail if we can't
    /// fetch it.
    pub async fn into_bytes(self) -> anyhow::Result<StdStreamPair<Vec<u8>>> {
        match self {
            Self::Local { stdout, stderr } => Ok(StdStreamPair { stdout, stderr }),
            Self::Remote(remote) => {
                let (stdout, stderr) = remote.into_stdout_stderr_bytes().await?;
                Ok(StdStreamPair { stdout, stderr })
            }
            Self::Empty => Ok(StdStreamPair {
                stdout: Vec::new(),
                stderr: Vec::new(),
            }),
        }
    }

    pub async fn into_re(
        self,
        client: &ManagedRemoteExecutionClient,
        use_case: RemoteExecutorUseCase,
    ) -> anyhow::Result<StdStreamPair<ReStdStream>> {
        match self {
            Self::Local { stdout, stderr } => {
                let (stdout, stderr) = future::try_join(
                    maybe_upload_to_re(client, use_case, stdout),
                    maybe_upload_to_re(client, use_case, stderr),
                )
                .await?;

                Ok(StdStreamPair { stdout, stderr })
            }
            Self::Remote(remote) => {
                // TODO (torozco): This assumes that the existing remote outputs we have have the
                // same re use case as what we passed in. Lots of things make this assumption, but
                // for the sake of being safe, check it.
                if remote.use_case() != use_case {
                    return Err(anyhow::anyhow!(
                        "Copying log outputs across RE use cases (from `{}` to `{}`) is not supported",
                        remote.use_case(),
                        use_case
                    ));
                }

                let (stdout, stderr) = remote.into_stdout_stderr();
                Ok(StdStreamPair { stdout, stderr })
            }
            Self::Empty => Ok(StdStreamPair {
                stdout: ReStdStream::None,
                stderr: ReStdStream::None,
            }),
        }
    }
}

async fn maybe_upload_to_re(
    client: &ManagedRemoteExecutionClient,
    use_case: RemoteExecutorUseCase,
    bytes: Vec<u8>,
) -> anyhow::Result<ReStdStream> {
    const MIN_STREAM_UPLOAD_SIZE: usize = 50 * 1024; // Same as RE
    if bytes.len() < MIN_STREAM_UPLOAD_SIZE {
        return Ok(ReStdStream::Raw(bytes));
    }
    let digest = client.upload_blob(bytes, use_case).await?;
    Ok(ReStdStream::Digest(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::ops::ControlFlow;

use async_trait::async_trait;
use more_futures::cancellation::CancellationContext;
use remote_execution as RE;

use crate::digest_config::DigestConfig;
use crate::execute::action_digest::ActionDigest;
use crate::execute::blobs::ActionBlobs;
use crate::execute::manager::CommandExecutionManager;
use crate::execute::request::CommandExecutionRequest;
use crate::execute::request::ExecutorPreference;
use crate::execute::result::CommandExecutionResult;
use crate::execute::target::CommandExecutionTarget;

pub struct PreparedAction {
    pub action: ActionDigest,
    // The encoded action and other messages referenced from it by digest (e.g. RE::Command).
    // Does not include the files referenced in inputs.
    pub blobs: ActionBlobs,
    pub platform: RE::Platform,
}

pub struct PreparedCommand<'a, 'b> {
    pub request: &'a CommandExecutionRequest,
    pub target: &'b dyn CommandExecutionTarget,
    pub prepared_action: &'a PreparedAction,
    pub digest_config: DigestConfig,
}

#[async_trait]
pub trait PreparedCommandExecutor: Send + Sync {
    /// 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.
    async fn exec_cmd(
        &self,
        command: &PreparedCommand<'_, '_>,
        manager: CommandExecutionManager,
        cancellations: &CancellationContext,
    ) -> CommandExecutionResult;

    /// Checks if there is any possibility for a command with a given executor preference to
    /// be executed locally.
    fn is_local_execution_possible(&self, executor_preference: ExecutorPreference) -> bool;
}

#[async_trait]
pub trait PreparedCommandOptionalExecutor: Send + Sync {
    /// Take a command and evaluate whether it needs to be actually executed (locally or remotely) or can be skipped.
    /// In the skip case, this should handle all the things that would happen in an actual execution (materialization, output declaration, etc.).
    ///
    /// Given a command, evaluate whether the execution can be skipped. (for example because it is already cached)
    /// If it can be skipped, return a CommandExecutionResult that can be used as if the action was just executed.
    /// Otherwise, return a CommandExecutionManager that can be used to execute the action.
    async fn maybe_execute(
        &self,
        command: &PreparedCommand<'_, '_>,
        manager: CommandExecutionManager,
        cancellations: &CancellationContext,
    ) -> ControlFlow<CommandExecutionResult, CommandExecutionManager>;
}

// When we don't want to check a command can be skipped, just use the NoOpCommandExecutor that always returns the continue case.
pub struct NoOpCommandExecutor {}

#[async_trait]
impl PreparedCommandOptionalExecutor for NoOpCommandExecutor {
    async fn maybe_execute(
        &self,
        _command: &PreparedCommand<'_, '_>,
        manager: CommandExecutionManager,
        _cancellations: &CancellationContext,
    ) -> ControlFlow<CommandExecutionResult, CommandExecutionManager> {
        ControlFlow::Continue(manager)
    }
}


/*
 * 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::Display;
use std::time::Duration;

use allocative::Allocative;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_common::local_resource_state::LocalResourceState;
use buck2_core::directory::DirectoryEntry;
use buck2_core::directory::DirectoryIterator;
use buck2_core::directory::FingerprintedDirectory;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::buck_out_path::BuckOutPath;
use buck2_core::fs::buck_out_path::BuckOutScratchPath;
use buck2_core::fs::buck_out_path::BuckOutTestPath;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::soft_error;
use buck2_util::collections::sorted_set::SortedSet;
use derive_more::Display;
use dupe::Dupe;
use gazebo::variants::UnpackVariants;
use host_sharing::host_sharing::HostSharingRequirements;
use indexmap::IndexSet;
use itertools::Itertools;
use sorted_vector_map::SortedVectorMap;
use thiserror::Error;

use crate::artifact::group::artifact_group_values_dyn::ArtifactGroupValuesDyn;
use crate::digest_config::DigestConfig;
use crate::directory::insert_entry;
use crate::directory::ActionDirectoryMember;
use crate::directory::ActionImmutableDirectory;
use crate::execute::environment_inheritance::EnvironmentInheritance;
use crate::execute::inputs_directory::inputs_directory;

#[derive(Clone)]
pub struct ActionMetadataBlob {
    pub data: Vec<u8>,
    pub digest: TrackedFileDigest,
    pub path: BuckOutPath,
}

pub enum CommandExecutionInput {
    Artifact(Box<dyn ArtifactGroupValuesDyn>),
    ActionMetadata(ActionMetadataBlob),
}

#[derive(PartialEq, Eq, Debug, Copy, Clone, Dupe, Hash)]
pub enum OutputCreationBehavior {
    Create,
    Parent,
}

#[derive(Debug, Error)]
#[error("Incompatible executor preferences: `{}` & `{}`", a, b)]
struct IncompatibleExecutorPreferences {
    a: ExecutorPreference,
    b: ExecutorPreference,
}

#[derive(Copy, Clone, Dupe, Display, Debug, Allocative)]
pub enum ExecutorPreference {
    Default,
    /// Fails when executed by a remote-only executor
    LocalRequired,
    /// Does not fail when executed by a remote-only executor
    LocalPreferred,
    /// Fails when executed by a local-only executor
    RemoteRequired,
    /// Does not fail when executed by a local-only executor
    RemotePreferred,
}

impl ExecutorPreference {
    pub fn and(self, other: Self) -> anyhow::Result<Self> {
        let requires_remote = self.requires_remote() || other.requires_remote();
        let requires_local = self.requires_local() || other.requires_local();

        if requires_remote && requires_local {
            return Err(IncompatibleExecutorPreferences { a: self, b: other }.into());
        }

        if requires_local {
            return Ok(Self::LocalRequired);
        }

        if requires_remote {
            return Ok(Self::RemoteRequired);
        }

        for pref in [self, other] {
            if pref.prefers_local() {
                return Ok(Self::LocalPreferred);
            }

            if pref.prefers_remote() {
                return Ok(Self::RemotePreferred);
            }
        }

        Ok(Self::Default)
    }

    pub fn requires_remote(&self) -> bool {
        match self {
            Self::LocalRequired => false,
            Self::LocalPreferred => false,
            Self::RemoteRequired => true,
            Self::RemotePreferred => false,
            Self::Default => false,
        }
    }

    pub fn requires_local(&self) -> bool {
        match self {
            Self::LocalRequired => true,
            Self::LocalPreferred => false,
            Self::RemoteRequired => false,
            Self::RemotePreferred => false,
            Self::Default => false,
        }
    }

    pub fn prefers_local(&self) -> bool {
        match self {
            Self::LocalRequired => true,
            Self::LocalPreferred => true,
            Self::RemoteRequired => false,
            Self::RemotePreferred => false,
            Self::Default => false,
        }
    }

    pub fn prefers_remote(&self) -> bool {
        match self {
            Self::LocalRequired => false,
            Self::LocalPreferred => false,
            Self::RemoteRequired => true,
            Self::RemotePreferred => true,
            Self::Default => false,
        }
    }
}

pub struct CommandExecutionPaths {
    inputs: Vec<CommandExecutionInput>,
    outputs: IndexSet<CommandExecutionOutput>,

    input_directory: ActionImmutableDirectory,
    output_paths: Vec<(ProjectRelativePathBuf, OutputType)>,

    /// Total size of input files.
    input_files_bytes: u64,
}

impl CommandExecutionPaths {
    pub fn new(
        inputs: Vec<CommandExecutionInput>,
        outputs: IndexSet<CommandExecutionOutput>,
        fs: &ArtifactFs,
        digest_config: DigestConfig,
    ) -> anyhow::Result<Self> {
        let mut builder = inputs_directory(&inputs, fs)?;

        let output_paths = outputs
            .iter()
            .map(|o| {
                let resolved = o.as_ref().resolve(fs);
                if let Some(dir) = resolved.path_to_create() {
                    builder.mkdir(dir)?;
                }
                let output_type = resolved.output_type;
                Ok((resolved.into_path(), output_type))
            })
            .collect::<anyhow::Result<Vec<_>>>()?;

        insert_entry(
            &mut builder,
            ProjectRelativePath::unchecked_new(".buckconfig"),
            DirectoryEntry::Leaf(ActionDirectoryMember::File(FileMetadata::empty(
                digest_config.cas_digest_config(),
            ))),
        )?;

        let input_directory = builder.fingerprint(digest_config.as_directory_serializer());

        let mut input_files_bytes = 0;

        for entry in input_directory
            .fingerprinted_unordered_walk()
            .without_paths()
        {
            match entry {
                DirectoryEntry::Leaf(ActionDirectoryMember::File(f)) => {
                    input_files_bytes += f.digest.size();
                }
                _ => {}
            };
        }

        Ok(Self {
            inputs,
            outputs,
            input_directory,
            output_paths,
            input_files_bytes,
        })
    }

    pub fn input_directory(&self) -> &ActionImmutableDirectory {
        &self.input_directory
    }

    pub fn output_paths(&self) -> &[(ProjectRelativePathBuf, OutputType)] {
        &self.output_paths
    }

    pub fn input_files_bytes(&self) -> u64 {
        self.input_files_bytes
    }
}

#[derive(Copy, Clone, Dupe, Debug, Display, Allocative, Hash, PartialEq, Eq)]
pub struct WorkerId(pub u64);

#[derive(Clone)]
pub struct WorkerSpec {
    pub id: WorkerId,
    pub exe: Vec<String>,
}

/// The data contains the information about the command to be executed.
pub struct CommandExecutionRequest {
    /// Optional arguments including executable prepended to `args` to get full command line.
    /// This is used by workers to separate worker arguments from executable arguments.
    exe: Vec<String>,
    args: Vec<String>,
    paths: CommandExecutionPaths,
    env: SortedVectorMap<String, String>,
    timeout: Option<Duration>,
    executor_preference: ExecutorPreference,
    // Run with a custom $TMPDIR, or just the standard system one
    custom_tmpdir: Option<BuckOutScratchPath>,
    host_sharing_requirements: HostSharingRequirements,
    // Used to disable the low pass filter for concurrent local actions. Enabled by default
    low_pass_filter: bool,
    /// Working directory, relative to the project root.
    working_directory: Option<ProjectRelativePathBuf>,
    /// Whether we should always prefetch stderr when executing. When it's needed, this lets us
    /// overlap stderr download with output downloads, which might be marginally useful to improve
    /// latency.
    prefetch_lossy_stderr: bool,
    /// Whether to cleanup outputs
    pub outputs_cleanup: bool,
    /// What environment variables to inherit from the Buck2 daemon.
    local_environment_inheritance: Option<EnvironmentInheritance>,
    /// Whether this command should override the fallback-only behavior on an hybrid executor and
    /// thus always run as if the executor was full-hybrid, assuming it is capable.
    force_full_hybrid_if_capable: bool,
    /// Whether to disable capturing performance counters for this execution.
    disable_miniperf: bool,
    required_local_resources: SortedSet<LocalResourceState>,
    /// Persistent worker to use for execution
    worker: Option<WorkerSpec>,
}

impl CommandExecutionRequest {
    pub fn new(
        exe: Vec<String>,
        args: Vec<String>,
        paths: CommandExecutionPaths,
        env: SortedVectorMap<String, String>,
    ) -> Self {
        Self {
            exe,
            args,
            paths,
            env,
            timeout: None,
            executor_preference: ExecutorPreference::Default,
            custom_tmpdir: None,
            host_sharing_requirements: HostSharingRequirements::default(),
            low_pass_filter: true,
            working_directory: None,
            prefetch_lossy_stderr: false,
            outputs_cleanup: true,
            local_environment_inheritance: None,
            force_full_hybrid_if_capable: false,
            disable_miniperf: false,
            required_local_resources: SortedSet::new(),
            worker: None,
        }
    }

    pub fn paths(&self) -> &CommandExecutionPaths {
        &self.paths
    }

    pub fn with_timeout(mut self, timeout: Duration) -> Self {
        self.timeout = Some(timeout);
        self
    }

    pub fn with_executor_preference(mut self, executor_preference: ExecutorPreference) -> Self {
        self.executor_preference = executor_preference;
        self
    }

    pub fn custom_tmpdir(&self) -> &Option<BuckOutScratchPath> {
        &self.custom_tmpdir
    }

    pub fn with_custom_tmpdir(mut self, custom_tmpdir: BuckOutScratchPath) -> Self {
        self.custom_tmpdir = Some(custom_tmpdir);
        self
    }

    pub fn with_host_sharing_requirements(
        mut self,
        host_sharing_requirements: HostSharingRequirements,
    ) -> Self {
        self.host_sharing_requirements = host_sharing_requirements;
        self
    }

    pub fn with_low_pass_filter(mut self, low_pass_filter: bool) -> Self {
        self.low_pass_filter = low_pass_filter;
        self
    }

    pub fn with_working_directory(mut self, working_directory: ProjectRelativePathBuf) -> Self {
        self.working_directory = Some(working_directory);
        self
    }

    pub fn with_prefetch_lossy_stderr(mut self, prefetch_lossy_stderr: bool) -> Self {
        self.prefetch_lossy_stderr = prefetch_lossy_stderr;
        self
    }

    pub fn with_outputs_cleanup(mut self, outputs_cleanup: bool) -> Self {
        self.outputs_cleanup = outputs_cleanup;
        self
    }

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

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

    pub fn all_args(&self) -> impl Iterator<Item = &String> {
        self.exe.iter().chain(self.args.iter())
    }

    pub fn all_args_vec(&self) -> Vec<String> {
        self.all_args().cloned().collect()
    }

    pub fn all_args_str(&self) -> String {
        self.all_args().join(" ")
    }

    pub fn exe(&self) -> &[String] {
        &self.exe
    }

    pub fn args(&self) -> &[String] {
        &self.args
    }

    pub fn worker(&self) -> &Option<WorkerSpec> {
        &self.worker
    }

    pub fn with_worker(mut self, worker: Option<WorkerSpec>) -> Self {
        self.worker = worker;
        self
    }

    pub fn inputs(&self) -> &[CommandExecutionInput] {
        &self.paths.inputs
    }

    pub fn outputs<'a>(&'a self) -> impl Iterator<Item = CommandExecutionOutputRef<'a>> + 'a {
        self.paths.outputs.iter().map(|output| output.as_ref())
    }

    pub fn env(&self) -> &SortedVectorMap<String, String> {
        &self.env
    }

    pub fn timeout(&self) -> Option<Duration> {
        self.timeout
    }

    pub fn executor_preference(&self) -> ExecutorPreference {
        self.executor_preference
    }

    pub fn host_sharing_requirements(&self) -> &HostSharingRequirements {
        &self.host_sharing_requirements
    }

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

    pub fn working_directory(&self) -> Option<&ProjectRelativePath> {
        self.working_directory.as_deref()
    }

    pub fn with_local_environment_inheritance(
        mut self,
        local_environment_inheritance: EnvironmentInheritance,
    ) -> Self {
        self.local_environment_inheritance = Some(local_environment_inheritance);
        self
    }

    pub fn local_environment_inheritance(&self) -> Option<&EnvironmentInheritance> {
        self.local_environment_inheritance.as_ref()
    }

    pub fn with_force_full_hybrid_if_capable(mut self, force_full_hybrid_if_capable: bool) -> Self {
        self.force_full_hybrid_if_capable = force_full_hybrid_if_capable;
        self
    }

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

    pub fn with_disable_miniperf(mut self, disable_miniperf: bool) -> Self {
        self.disable_miniperf = disable_miniperf;
        self
    }

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

    pub fn with_required_local_resources(
        mut self,
        required_local_resources: Vec<LocalResourceState>,
    ) -> anyhow::Result<Self> {
        let original_len = required_local_resources.len();
        self.required_local_resources = required_local_resources.into_iter().collect();
        if self.required_local_resources.len() != original_len {
            return Err(anyhow::anyhow!(
                "Each provided local resource state is supposed to come from a different target."
            ));
        }
        Ok(self)
    }

    pub fn required_local_resources(&self) -> &SortedSet<LocalResourceState> {
        &self.required_local_resources
    }
}

/// Is an output a file or a directory
#[derive(PartialEq, Eq, Hash, Debug, Copy, Clone, Dupe, Allocative)]
pub enum OutputType {
    /// We don't know - used to represent legacy code that doesn't yet declare the output type properly.
    /// We aim to mostly remove this alternative over time.
    FileOrDirectory,
    File,
    Directory,
}

#[derive(Debug, Error)]
enum OutputTypeError {
    #[error("Expected {1:?}, but `{0}` is already declared as {2:?}")]
    CheckPath(String, OutputType, OutputType),
}

impl OutputType {
    /// We are hoping for something of output type self, and got a path of an existing declared type.
    /// Raise an error if they are incompatible.
    pub fn check_path(
        self,
        path_for_error_message: impl Display,
        output_type: OutputType,
    ) -> anyhow::Result<()> {
        if self == OutputType::Directory && output_type == OutputType::FileOrDirectory {
            // If we treat paths whose declared type is FileOrDirectory like files, then that's incompatible with directory
            soft_error!(
                "declare_wrong_type",
                OutputTypeError::CheckPath(
                    path_for_error_message.to_string(),
                    self,
                    OutputType::File
                )
                .into(),
                quiet: true
            )?;
            Ok(())
        } else if self == output_type
            || self == OutputType::FileOrDirectory
            || output_type == OutputType::FileOrDirectory
        {
            Ok(())
        } else {
            Err(
                OutputTypeError::CheckPath(path_for_error_message.to_string(), self, output_type)
                    .into(),
            )
        }
    }
}

#[derive(UnpackVariants, PartialEq, Eq, Hash, Debug)]
pub enum CommandExecutionOutputRef<'a> {
    BuildArtifact {
        path: &'a BuckOutPath,
        output_type: OutputType,
    },
    TestPath {
        path: &'a BuckOutTestPath,
        create: OutputCreationBehavior,
    },
}

impl<'a> CommandExecutionOutputRef<'a> {
    /// Resolve this output to a ResolvedCommandExecutionOutput that allows access to the output
    /// path as well as any dirs to create.
    pub fn resolve(&self, fs: &ArtifactFs) -> ResolvedCommandExecutionOutput {
        match self {
            Self::BuildArtifact { path, output_type } => ResolvedCommandExecutionOutput {
                path: fs.resolve_build(path),
                create: OutputCreationBehavior::Parent,
                output_type: *output_type,
            },
            Self::TestPath { path, create } => ResolvedCommandExecutionOutput {
                path: fs.buck_out_path_resolver().resolve_test(path),
                create: *create,
                output_type: OutputType::FileOrDirectory,
            },
        }
    }

    pub fn cloned(&self) -> CommandExecutionOutput {
        match self {
            Self::BuildArtifact { path, output_type } => CommandExecutionOutput::BuildArtifact {
                path: (*path).dupe(),
                output_type: *output_type,
            },
            Self::TestPath { path, create } => CommandExecutionOutput::TestPath {
                path: (*path).clone(),
                create: *create,
            },
        }
    }
}

#[derive(UnpackVariants, PartialEq, Eq, Hash, Debug)]
pub enum CommandExecutionOutput {
    BuildArtifact {
        path: BuckOutPath,
        output_type: OutputType,
    },
    TestPath {
        path: BuckOutTestPath,
        create: OutputCreationBehavior,
    },
}

impl CommandExecutionOutput {
    pub fn as_ref<'a>(&'a self) -> CommandExecutionOutputRef<'a> {
        match self {
            Self::BuildArtifact {
                ref path,
                output_type,
            } => CommandExecutionOutputRef::BuildArtifact {
                path,
                output_type: *output_type,
            },
            Self::TestPath { ref path, create } => CommandExecutionOutputRef::TestPath {
                path,
                create: *create,
            },
        }
    }
}

/// The path for an output, and what should be created before running the action that produces this
/// output.
pub struct ResolvedCommandExecutionOutput {
    pub path: ProjectRelativePathBuf,
    pub output_type: OutputType,
    create: OutputCreationBehavior,
}

impl ResolvedCommandExecutionOutput {
    /// Return this output's path.
    pub fn path(&self) -> &ProjectRelativePath {
        &self.path
    }

    /// Return this output's path, ownedly.
    pub fn into_path(self) -> ProjectRelativePathBuf {
        self.path
    }

    /// Return the path that this output requires creating before running the action.
    pub fn path_to_create(&self) -> Option<&ProjectRelativePath> {
        match self.create {
            OutputCreationBehavior::Create => Some(&self.path),
            OutputCreationBehavior::Parent => self.path.parent(),
        }
    }
}


/*
 * 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::convert::Infallible;
use std::fmt::Display;
use std::ops::ControlFlow;
use std::ops::FromResidual;
use std::time::Duration;
use std::time::SystemTime;

use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use dupe::Dupe;
use indexmap::IndexMap;

use crate::artifact_value::ArtifactValue;
use crate::execute::claim::Claim;
use crate::execute::kind::CommandExecutionKind;
use crate::execute::output::CommandStdStreams;
use crate::execute::request::CommandExecutionOutput;
use crate::execute::request::ResolvedCommandExecutionOutput;
use crate::output_size::OutputSize;

/// "Status" of an action execution indicating how it finished. E.g. "built_remotely", "local_fallback", "action_cache".
#[derive(Debug)]
pub enum CommandExecutionStatus {
    Success {
        execution_kind: CommandExecutionKind,
    },
    Failure {
        execution_kind: CommandExecutionKind,
    },
    Error {
        stage: &'static str,
        error: anyhow::Error,
    },
    TimedOut {
        execution_kind: CommandExecutionKind,
        duration: Duration,
    },
    // TODO: We should rename this.
    Cancelled,
}

impl CommandExecutionStatus {
    pub fn execution_kind(&self) -> Option<&CommandExecutionKind> {
        match self {
            CommandExecutionStatus::Success { execution_kind, .. } => Some(execution_kind),
            CommandExecutionStatus::Failure { execution_kind } => Some(execution_kind),
            CommandExecutionStatus::Error { .. } => None,
            CommandExecutionStatus::TimedOut { execution_kind, .. } => Some(execution_kind),
            CommandExecutionStatus::Cancelled => None,
        }
    }
}

impl Display for CommandExecutionStatus {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            CommandExecutionStatus::Success { execution_kind, .. } => {
                write!(f, "success {}", execution_kind,)
            }
            CommandExecutionStatus::Failure { execution_kind } => {
                write!(f, "failure {}", execution_kind,)
            }
            CommandExecutionStatus::Error { stage, error } => {
                write!(f, "error:{}\n{:#}", stage, error)
            }
            CommandExecutionStatus::TimedOut { duration, .. } => {
                write!(f, "timed out after {:.3}s", duration.as_secs_f64())
            }
            CommandExecutionStatus::Cancelled => write!(f, "Cancelled"),
        }
    }
}

/// Unlike action where we only really have just 1 time, commands can have slightly richer timing
/// data.
#[derive(Debug, Copy, Clone, Dupe)]
pub struct CommandExecutionMetadata {
    /// How long this build actually waited for this action to complete
    pub wall_time: Duration,

    /// How long this command queued in RE. This value excludes execution time, i.e. for action cache hit,
    /// this value represents how long a request has to wait for server to handle.
    pub re_queue_time: Option<Duration>,

    /// How long this command actually took to execute. This can be different from the wall_time if
    /// this was e.g. an action cache hit, in which case this field would reflect how long the
    /// command took to actually execute but not how we had to wait for it.
    pub execution_time: Duration,

    /// When execution started.
    pub start_time: SystemTime,

    /// Additional stats.
    pub execution_stats: Option<buck2_data::CommandExecutionStats>,

    /// How long it took to materialize the action's inputs.
    pub input_materialization_duration: Duration,
}

impl Default for CommandExecutionMetadata {
    fn default() -> Self {
        Self {
            wall_time: Duration::default(),
            re_queue_time: None,
            execution_time: Duration::default(),
            start_time: SystemTime::now(),
            execution_stats: None,
            input_materialization_duration: Duration::default(),
        }
    }
}

/// CommandExecutionResult is the result of an executor executing a command.
#[derive(Debug)]
pub struct CommandExecutionResult {
    /// The outputs produced by this command
    pub outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
    /// How it executed.
    pub report: CommandExecutionReport,
    /// A previously rejected execution of this command.
    pub rejected_execution: Option<CommandExecutionReport>,
    /// Whether this was uploaded to cache, by Buck2.
    pub did_cache_upload: bool,
    /// Whether this command was eligible for hybrid execution.
    pub eligible_for_full_hybrid: bool,
}

impl CommandExecutionResult {
    /// Total size of all outputs in bytes.
    pub fn calc_output_size_bytes(&self) -> u64 {
        self.outputs
            .values()
            .map(|v| v.calc_output_count_and_bytes().bytes)
            .sum()
    }

    pub fn resolve_outputs<'a>(
        &'a self,
        fs: &'a ArtifactFs,
    ) -> impl Iterator<Item = (ResolvedCommandExecutionOutput, &ArtifactValue)> + 'a {
        self.outputs
            .iter()
            .map(|(output, value)| (output.as_ref().resolve(fs), value))
    }
}

/// Describes how a command executed.
#[derive(Debug)]
pub struct CommandExecutionReport {
    pub claim: Option<Box<dyn Claim>>,
    pub status: CommandExecutionStatus,
    pub timing: CommandExecutionMetadata,
    pub std_streams: CommandStdStreams,
    /// No exit_code means the command did not finish executing. Signals get mapped into this as
    /// 128 + SIGNUM, which is the convention shells follow.
    pub exit_code: Option<i32>,
}

/// Implement FromResidual so that it's easier to refactor functions returning a CommandExecutionResult
/// (it allows to easily factor out early returns into another function and then propagate them with `?`).
impl FromResidual<ControlFlow<Self, Infallible>> for CommandExecutionResult {
    fn from_residual(residual: ControlFlow<Self, Infallible>) -> Self {
        match residual {
            ControlFlow::Break(v) => v,
            ControlFlow::Continue(_) => unreachable!(),
        }
    }
}


/*
 * 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::Debug;

pub trait CommandExecutionTarget: Send + Sync + Debug {
    fn re_action_key(&self) -> String;

    fn re_affinity_key(&self) -> String;

    fn as_proto_action_key(&self) -> buck2_data::ActionKey;

    fn as_proto_action_name(&self) -> buck2_data::ActionName;
}


/*
 * 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::sync::Mutex;

use async_trait::async_trait;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use indexmap::IndexMap;
use more_futures::cancellation::CancellationContext;
use sorted_vector_map::SortedVectorMap;

use crate::artifact_value::ArtifactValue;
use crate::execute::action_digest::ActionDigest;
use crate::execute::kind::CommandExecutionKind;
use crate::execute::manager::CommandExecutionManager;
use crate::execute::manager::CommandExecutionManagerExt;
use crate::execute::prepared::PreparedCommand;
use crate::execute::prepared::PreparedCommandExecutor;
use crate::execute::request::CommandExecutionOutput;
use crate::execute::request::ExecutorPreference;
use crate::execute::result::CommandExecutionMetadata;
use crate::execute::result::CommandExecutionResult;

#[derive(Debug, PartialEq, Eq)]
pub struct DryRunEntry {
    pub args: Vec<String>,
    pub outputs: Vec<CommandExecutionOutput>,
    pub env: SortedVectorMap<String, String>,
}

/// Records executed commands into the provided tracker and returns a successful result for all commands.
/// If the filesystem is supplied, the dry run executor will write the executed command as contents
/// to the output file.
pub struct DryRunExecutor {
    tracker: Arc<Mutex<Vec<DryRunEntry>>>,
    fs: ArtifactFs,
}

impl DryRunExecutor {
    pub fn new(tracker: Arc<Mutex<Vec<DryRunEntry>>>, fs: ArtifactFs) -> Self {
        Self { tracker, fs }
    }
}

#[async_trait]
impl PreparedCommandExecutor for DryRunExecutor {
    async fn exec_cmd(
        &self,
        command: &PreparedCommand<'_, '_>,
        manager: CommandExecutionManager,
        _cancellations: &CancellationContext,
    ) -> CommandExecutionResult {
        let PreparedCommand {
            request,
            target: _target,
            prepared_action: _prepared_action,
            digest_config,
        } = command;

        let manager = manager.claim().await;

        let args = request.all_args_vec();
        let outputs = request.outputs().map(|o| o.cloned()).collect();
        let env = request.env().to_owned();

        self.tracker
            .lock()
            .unwrap()
            .push(DryRunEntry { args, outputs, env });

        let exec_kind = CommandExecutionKind::Local {
            digest: ActionDigest::empty(digest_config.cas_digest_config()),
            command: Default::default(),
            env: Default::default(),
        };

        match request
            .outputs()
            .map(|x| {
                let path = x.resolve(&self.fs).into_path();
                self.fs.fs().write_file(&path, "", false)?;
                Ok((x.cloned(), ArtifactValue::file(digest_config.empty_file())))
            })
            .collect::<anyhow::Result<_>>()
        {
            Ok(outputs) => manager.success(
                exec_kind,
                outputs,
                Default::default(),
                CommandExecutionMetadata::default(),
            ),
            // NOTE: This should probably be an error() but who cares.
            Err(..) => manager.failure(
                exec_kind,
                IndexMap::new(),
                Default::default(),
                Some(1),
                CommandExecutionMetadata::default(),
            ),
        }
    }

    fn is_local_execution_possible(&self, _executor_preference: ExecutorPreference) -> bool {
        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::collections::BTreeMap;
use std::fs;
use std::path::Path;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use buck2_core::directory::DirectoryEntry;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_eden::connection::EdenConnectionManager;
use buck2_util::process::background_command;
use edenfs::client::EdenService;
use edenfs::CheckoutMode;
use edenfs::EnsureMaterializedParams;
use edenfs::ObjectType;
use edenfs::RemoveRecursivelyParams;
use edenfs::SetPathObjectIdParams;
use fbinit::FacebookInit;
use more_futures::cancellation::CancellationContext;
use serde::Deserialize;
use thiserror::Error;
use tokio::sync::Semaphore;

use crate::artifact_value::ArtifactValue;
use crate::directory::ActionDirectoryMember;
use crate::re::re_get_session_id::ReGetSessionId;

pub type EdenFsClient = Arc<dyn EdenService + Sync>;

const RE_SYMLINK_PREFIX: &str = "re-symlink";

#[derive(Error, Debug)]
pub enum SetupBuckOutError {
    #[error("Could not read Eden Config")]
    EdenConfigIOError(#[source] std::io::Error),

    #[error("Error during parsing the Eden Config")]
    InvalidEdenConfig(#[source] toml::de::Error),

    #[error("Could not clean buck-out")]
    CleanBuckOutFailed(#[source] std::io::Error),

    #[error("Could not create Eden buck-out")]
    CreateEdenBuckoutOutputFailed(#[source] std::io::Error),

    #[error("Failed to run Eden clone to create buck-out: {0}")]
    FailedToRunEdenClone(String),

    #[error("Could not add {0} to the redirection list of Eden buck-out: {1}")]
    EdenRedirectAddFailed(String, #[source] std::io::Error),

    #[error("Failed to run Eden redirection add path {0} to bind a dir to buck-out: {1}")]
    FailedToRunEdenRedirectionAdd(String, String),
}

#[derive(Deserialize, Debug)]
struct Repository {
    #[serde(rename = "type")]
    repo_type: String,
}

#[derive(Deserialize, Debug)]
struct EdenConfig {
    repository: Repository,
    redirections: Option<BTreeMap<String, String>>,
}

/// Eden has a magic .eden dir for all directories and its config is in .eden/client/config.toml.
/// Read the TOML file and make sure it is recas type mount.
fn read_eden_config(config_path: &Path) -> Result<EdenConfig, SetupBuckOutError> {
    let config_contents =
        fs::read_to_string(config_path).map_err(SetupBuckOutError::EdenConfigIOError)?;
    let config: EdenConfig =
        toml::from_str(&config_contents).map_err(SetupBuckOutError::InvalidEdenConfig)?;
    Ok(config)
}

pub fn is_recas_eden_mount(buck_out: &AbsNormPathBuf) -> Result<bool, SetupBuckOutError> {
    let config_path = buck_out.as_path().join(".eden/client/config.toml");
    if fs::metadata(&config_path).is_err() {
        return Ok(false);
    }
    let config = read_eden_config(&config_path)?;
    Ok(config.repository.repo_type.as_str() == "recas")
}

fn execute_eden_clone(buck_out: &AbsNormPathBuf) -> Result<(), SetupBuckOutError> {
    // If the buck-out directory exists
    if fs::metadata(buck_out).is_ok() {
        // If it is already an eden-based buck-out, then it's good
        if is_recas_eden_mount(buck_out)? {
            return Ok(());
        }
        // If it is regular FS, remove it first
        fs::remove_dir_all(buck_out).map_err(SetupBuckOutError::CleanBuckOutFailed)?;
    }

    let eden_clone = background_command("eden")
        .arg("clone")
        .arg("")
        .arg(buck_out.as_os_str())
        .arg("--backing-store=recas")
        .arg("--allow-nested-checkout")
        .current_dir("/")
        .output()
        .map_err(SetupBuckOutError::CreateEdenBuckoutOutputFailed)?;

    if !eden_clone.status.success() {
        return Err(SetupBuckOutError::FailedToRunEdenClone(
            String::from_utf8_lossy(&eden_clone.stderr).to_string(),
        ));
    }
    Ok(())
}

fn is_path_redirect(buck_out: &AbsNormPathBuf, path: &str) -> Result<bool, SetupBuckOutError> {
    let config_path = buck_out.as_path().join(".eden/client/config.toml");
    if fs::metadata(&config_path).is_err() {
        return Ok(false);
    }
    let config = read_eden_config(&config_path)?;

    if let Some(redir_config) = config.redirections {
        if let Some(redir_type) = redir_config.get(path) {
            return Ok(redir_type.as_str() == "bind");
        }
    }

    Ok(false)
}

fn execute_eden_redirection_add(
    buck_out: &AbsNormPathBuf,
    path: &str,
) -> Result<(), SetupBuckOutError> {
    if is_path_redirect(buck_out, path)? {
        return Ok(());
    }
    let eden_direct = background_command("eden")
        .arg("redirect")
        .arg("add")
        .arg(path)
        .arg("bind")
        .current_dir(buck_out.as_os_str())
        .output()
        .map_err(|e| SetupBuckOutError::EdenRedirectAddFailed(String::from(path), e))?;

    if !eden_direct.status.success() {
        return Err(SetupBuckOutError::FailedToRunEdenRedirectionAdd(
            String::from(path),
            String::from_utf8_lossy(&eden_direct.stderr).to_string(),
        ));
    }
    Ok(())
}

fn get_symlink_object_id(target: &str) -> String {
    format!("{}:{}", RE_SYMLINK_PREFIX, target)
}

fn get_object_id_and_type(value: &ArtifactValue) -> (String, ObjectType) {
    match value.entry() {
        DirectoryEntry::Dir(dir) => (dir.fingerprint().to_string(), ObjectType::TREE),
        DirectoryEntry::Leaf(ActionDirectoryMember::File(file)) => {
            let object_type = match file.is_executable {
                true => ObjectType::EXECUTABLE_FILE,
                false => ObjectType::REGULAR_FILE,
            };
            (file.digest.to_string(), object_type)
        }
        DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(symlink)) => (
            get_symlink_object_id(symlink.target().as_str()),
            ObjectType::SYMLINK,
        ),
        DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(external_sym)) => (
            get_symlink_object_id(external_sym.to_path_buf().to_str().unwrap()),
            ObjectType::SYMLINK,
        ),
    }
}

fn setup(buck_out: &AbsNormPathBuf) -> Result<(), SetupBuckOutError> {
    // Run "eden clone" command to create a brand new buck-out
    execute_eden_clone(buck_out)?;

    // Run eden redirection to exclude write heavy dirs from using EdenFS
    ["log", "tmp", "re_logs", "cache"]
        .into_iter()
        .try_for_each(|dir| -> Result<(), SetupBuckOutError> {
            execute_eden_redirection_add(buck_out, dir)?;
            Ok(())
        })?;

    Ok(())
}

#[derive(Allocative)]
pub struct EdenBuckOut {
    /// Relative path of buck_out, i.e "buck-out/v2"
    buck_out_path: ProjectRelativePathBuf,
    connection_manager: EdenConnectionManager,
    re_get_session_id: Arc<dyn ReGetSessionId>,
}

impl EdenBuckOut {
    pub fn new(
        fb: FacebookInit,
        buck_out_path: ProjectRelativePathBuf,
        mount_point: AbsNormPathBuf,
        re_get_session_id: Arc<dyn ReGetSessionId>,
    ) -> anyhow::Result<Self> {
        setup(&mount_point)?;

        // Default to the number of CPUs. This value is very conservative
        // TODO (yipu): Benchmark and figure out optimal default concurrency
        static CONCURRENCY: EnvHelper<usize> = EnvHelper::new("BUCK2_EDEN_CONCURRENCY");
        let concurrency = CONCURRENCY.get_copied()?.unwrap_or_else(num_cpus::get);
        let connection_manager =
            EdenConnectionManager::new(fb, &mount_point, Semaphore::new(concurrency))?
                .expect("EdenFS mount does not setup correctly");

        Ok(Self {
            buck_out_path,
            connection_manager,
            re_get_session_id,
        })
    }

    // TODO(yipu): This needs to take a RE use case as input and pass it to Eden.
    pub async fn set_path_object_id(
        &self,
        path: &ProjectRelativePathBuf,
        value: &ArtifactValue,
    ) -> anyhow::Result<()> {
        // Eden's SetPathObjectId requires object_id and object_type, the object_id is a string of
        // the RE digest for a file or symlink target for a symlink or external symlink
        let (object_id, object_type) = get_object_id_and_type(value);
        let relpath_to_buck_out = path.strip_prefix(&self.buck_out_path).with_context(|| {
            format!(
                "Invalid artifact: path might not in the buck-out directory: {}",
                path
            )
        })?;

        let re_session_id = self.re_get_session_id.get_session_id().await?;

        let params = SetPathObjectIdParams {
            mountPoint: self.connection_manager.get_mount_point(),
            path: relpath_to_buck_out.as_str().as_bytes().to_vec(),
            objectId: object_id.into_bytes(),
            r#type: object_type,
            mode: CheckoutMode::FORCE,
            requestInfo: Some(BTreeMap::from([(
                String::from("session-id"),
                re_session_id,
            )])),
            ..Default::default()
        };
        self.connection_manager
            .with_eden(move |eden| eden.setPathObjectId(&params))
            .await?;
        Ok(())
    }

    async fn remove_path_recursive(
        &self,
        project_fs: &ProjectRoot,
        path: &ProjectRelativePathBuf,
        cancellations: &CancellationContext<'_>,
    ) -> anyhow::Result<()> {
        // Existence check would not trigger materialization since EdenFS will fast
        // return by not materialization the path. So the only cost is Eden will load
        // path's ancestors, which should be relatively inexpensive.
        if !project_fs.resolve(path).exists() {
            return Ok(());
        }

        let relpath_to_buck_out = path.strip_prefix(&self.buck_out_path).with_context(|| {
            format!(
                "Invalid artifact: path might not in the buck-out directory: {}",
                path
            )
        })?;

        cancellations
            .critical_section(|| async move {
                let params = RemoveRecursivelyParams {
                    mountPoint: self.connection_manager.get_mount_point(),
                    path: relpath_to_buck_out.as_str().as_bytes().to_vec(),
                    ..Default::default()
                };

                self.connection_manager
                    .with_eden(move |eden| eden.removeRecursively(&params))
                    .await?;
                Ok(())
            })
            .await
    }

    pub async fn remove_paths_recursive(
        &self,
        project_fs: &ProjectRoot,
        paths: Vec<ProjectRelativePathBuf>,
        cancellations: &CancellationContext<'_>,
    ) -> anyhow::Result<()> {
        // TODO(bobyf, torozco) does this need to be critical section

        let futs = paths.iter().map(|path| async move {
            self.remove_path_recursive(project_fs, path, cancellations)
                .await
                .with_context(|| format!("[eden] Error cleaning up path {}", path))
        });

        futures::future::try_join_all(futs).await?;

        Ok(())
    }

    pub async fn ensure_materialized(
        &self,
        paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<()> {
        let file_paths = paths
            .iter()
            .filter_map(|path| path.strip_prefix(&self.buck_out_path).ok())
            .map(|relpath| relpath.as_str().as_bytes().to_vec())
            .collect::<Vec<_>>();
        static ENSURE_MATERIALIZED_IN_BACKGROUND: EnvHelper<bool> =
            EnvHelper::new("BUCK2_EDEN_ENSURE_MATERIALIZED_IN_BACKGROUND");
        let background = ENSURE_MATERIALIZED_IN_BACKGROUND
            .get_copied()?
            .unwrap_or(true);
        let params = EnsureMaterializedParams {
            mountPoint: self.connection_manager.get_mount_point(),
            paths: file_paths,
            background,
            followSymlink: true, // Also materialize symlink targets
            ..Default::default()
        };

        self.connection_manager
            .with_eden(move |eden| eden.ensureMaterialized(&params))
            .await?;
        Ok(())
    }
}

#[cfg(all(test, not(windows)))]
mod tests {
    use std::path::PathBuf;
    use std::sync::Arc;

    use buck2_common::external_symlink::ExternalSymlink;
    use buck2_common::file_ops::FileMetadata;
    use buck2_common::file_ops::TrackedFileDigest;
    use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;
    use buck2_core::fs::paths::RelativePathBuf;
    use dupe::Dupe;

    use super::*;
    use crate::artifact_value::ArtifactValue;
    use crate::digest_config::DigestConfig;
    use crate::directory::Symlink;

    #[test]
    fn test_symlink_object_id() -> anyhow::Result<()> {
        let symlink = Symlink::new(RelativePathBuf::from("path/include"));
        assert_eq!(
            "re-symlink:path/include",
            get_symlink_object_id(symlink.target().as_str())
        );
        Ok(())
    }

    #[test]
    fn test_external_symlink_object_id() -> anyhow::Result<()> {
        let symlink = ExternalSymlink::new(
            PathBuf::from("/mnt/gvfs"),
            Some(ForwardRelativePathBuf::unchecked_new("include".to_owned())),
        )
        .unwrap();
        assert_eq!(
            "re-symlink:/mnt/gvfs/include",
            get_symlink_object_id(&Arc::new(symlink).to_path_buf().to_string_lossy())
        );
        Ok(())
    }

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

        let digest = TrackedFileDigest::from_content(b"content", digest_config.cas_digest_config());
        let metadata_executable = FileMetadata {
            digest: digest.dupe(),
            is_executable: true,
        };

        let value = ArtifactValue::from(DirectoryEntry::Leaf(ActionDirectoryMember::File(
            metadata_executable,
        )));
        let (object_id, object_type) = get_object_id_and_type(&value);
        assert_eq!(digest.to_string(), object_id);
        assert_eq!(ObjectType::EXECUTABLE_FILE, object_type);

        let metadata_regular = FileMetadata {
            digest: digest.dupe(),
            is_executable: false,
        };

        let value = ArtifactValue::from(DirectoryEntry::Leaf(ActionDirectoryMember::File(
            metadata_regular,
        )));
        let (object_id, object_type) = get_object_id_and_type(&value);
        assert_eq!(digest.to_string(), object_id);
        assert_eq!(ObjectType::REGULAR_FILE, object_type);
        Ok(())
    }

    #[test]
    fn test_get_object_id_and_type_symlink() -> anyhow::Result<()> {
        let symlink = Symlink::new(RelativePathBuf::from("path/include"));
        let value = ArtifactValue::from(DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(
            Arc::new(symlink),
        )));
        let (object_id, object_type) = get_object_id_and_type(&value);
        assert_eq!("re-symlink:path/include", object_id);
        assert_eq!(ObjectType::SYMLINK, object_type);
        Ok(())
    }

    #[test]
    fn test_get_object_id_and_type_external_symlink() -> anyhow::Result<()> {
        let symlink = Arc::new(
            ExternalSymlink::new(
                PathBuf::from("/mnt/gvfs"),
                Some(ForwardRelativePathBuf::unchecked_new("include".to_owned())),
            )
            .unwrap(),
        );

        let value = ArtifactValue::from(DirectoryEntry::Leaf(
            ActionDirectoryMember::ExternalSymlink(symlink),
        ));
        let (object_id, object_type) = get_object_id_and_type(&value);
        assert_eq!("re-symlink:/mnt/gvfs/include", object_id);
        assert_eq!(ObjectType::SYMLINK, object_type);
        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 std::io::Write;
use std::sync::Arc;
use std::time::Duration;

use allocative::Allocative;
use anyhow::Context as _;
use buck2_common::cas_digest::CasDigestConfig;
use buck2_common::cas_digest::DigestAlgorithmKind;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_common::http::counting_client::CountingHttpClient;
use buck2_common::http::retries::http_retry;
use buck2_common::http::retries::AsHttpError;
use buck2_common::http::retries::HttpError;
use buck2_common::http::HttpClient;
use buck2_core::fs::fs_util;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use bytes::Bytes;
use digest::DynDigest;
use dupe::Dupe;
use futures::stream::Stream;
use futures::StreamExt;
use hyper::Response;
use sha1::Digest;
use sha1::Sha1;
use sha2::Sha256;
use smallvec::SmallVec;
use thiserror::Error;

use crate::digest_config::DigestConfig;

#[derive(Debug, Clone, Dupe, Allocative)]
pub enum Checksum {
    Sha1(Arc<str>),
    Sha256(Arc<str>),
    Both { sha1: Arc<str>, sha256: Arc<str> },
}

impl Checksum {
    pub fn sha1(&self) -> Option<&str> {
        match self {
            Self::Sha1(sha1) => Some(sha1),
            Self::Sha256(..) => None,
            Self::Both { sha1, .. } => Some(sha1),
        }
    }

    pub fn sha256(&self) -> Option<&str> {
        match self {
            Self::Sha1(..) => None,
            Self::Sha256(sha256) => Some(sha256),
            Self::Both { sha256, .. } => Some(sha256),
        }
    }
}

#[derive(Debug, Error)]
enum HttpHeadError {
    #[error("Error performing http_head request")]
    Client(#[from] HttpError),
}

#[derive(Debug, Error)]
enum HttpDownloadError {
    #[error("Error performing http_download request")]
    Client(#[from] HttpError),

    #[error("Invalid {0} digest. Expected {1}, got {2}. URL: {3}")]
    InvalidChecksum(&'static str, String, String, String),

    #[error(transparent)]
    IoError(anyhow::Error),
}

impl AsHttpError for HttpHeadError {
    fn as_http_error(&self) -> Option<&HttpError> {
        match self {
            Self::Client(e) => Some(e),
        }
    }
}

impl AsHttpError for HttpDownloadError {
    fn as_http_error(&self) -> Option<&HttpError> {
        match self {
            Self::Client(e) => Some(e),
            Self::InvalidChecksum(..) | Self::IoError(..) => None,
        }
    }
}

pub async fn http_head(client: &dyn HttpClient, url: &str) -> anyhow::Result<Response<()>> {
    let response = http_retry(
        || async {
            client
                .head(url)
                .await
                .map_err(|e| HttpHeadError::Client(HttpError::Client(e)))
        },
        vec![2, 4, 8].into_iter().map(Duration::from_secs).collect(),
    )
    .await?;
    Ok(response)
}

pub async fn http_download(
    client: &CountingHttpClient,
    fs: &ProjectRoot,
    digest_config: DigestConfig,
    path: &ProjectRelativePath,
    url: &str,
    checksum: &Checksum,
    executable: bool,
) -> anyhow::Result<TrackedFileDigest> {
    let abs_path = fs.resolve(path);
    if let Some(dir) = abs_path.parent() {
        fs_util::create_dir_all(fs.resolve(dir))?;
    }

    Ok(http_retry(
        || async {
            let file = fs_util::create_file(&abs_path).map_err(HttpDownloadError::IoError)?;

            let stream = client
                .get(url)
                .await
                .map_err(|e| HttpDownloadError::Client(HttpError::Client(e)))?
                .into_body();
            let buf_writer = std::io::BufWriter::new(file);

            let digest = copy_and_hash(
                url,
                &abs_path,
                stream,
                buf_writer,
                digest_config.cas_digest_config(),
                checksum,
            )
            .await?;

            if executable {
                fs.set_executable(path)
                    .map_err(HttpDownloadError::IoError)?;
            }

            Result::<_, HttpDownloadError>::Ok(TrackedFileDigest::new(
                digest,
                digest_config.cas_digest_config(),
            ))
        },
        vec![2, 4, 8].into_iter().map(Duration::from_secs).collect(),
    )
    .await?)
}

/// Copy a stream into a writer while producing its digest and checksumming it.
async fn copy_and_hash(
    url: &str,
    abs_path: &(impl std::fmt::Display + ?Sized),
    mut stream: impl Stream<Item = Result<Bytes, hyper::Error>> + Unpin,
    mut writer: impl Write,
    digest_config: CasDigestConfig,
    checksum: &Checksum,
) -> Result<FileDigest, HttpDownloadError> {
    let mut digester = FileDigest::digester(digest_config);

    // For each checksum entry we have, we're going to add a validator. We might have to create
    // a new hasher, or reuse the `FileDigest::digester` if it matches.

    enum Validator {
        PrimaryDigest,
        ExtraDigest(Box<dyn DynDigest + Send>),
    }

    let mut validators = SmallVec::<[_; 2]>::new();

    if let Some(sha1) = checksum.sha1() {
        let validator = if digester.algorithm() == DigestAlgorithmKind::Sha1 {
            Validator::PrimaryDigest
        } else {
            Validator::ExtraDigest(Box::new(Sha1::new()) as _)
        };

        validators.push((validator, sha1, "sha1"));
    }

    if let Some(sha256) = checksum.sha256() {
        let validator = if digester.algorithm() == DigestAlgorithmKind::Sha256 {
            Validator::PrimaryDigest
        } else {
            Validator::ExtraDigest(Box::new(Sha256::new()) as _)
        };

        validators.push((validator, sha256, "sha256"));
    }

    while let Some(chunk) = stream.next().await {
        let chunk = chunk.map_err(|source| HttpError::Transfer {
            received: digester.bytes_read(),
            url: url.to_owned(),
            source,
        })?;
        writer
            .write(&chunk)
            .with_context(|| format!("write({})", abs_path))
            .map_err(HttpDownloadError::IoError)?;

        digester.update(&chunk);
        for (validator, _expected, _kind) in validators.iter_mut() {
            if let Validator::ExtraDigest(hasher) = validator {
                hasher.update(&chunk);
            }
        }
    }
    writer
        .flush()
        .with_context(|| format!("flush({})", abs_path))
        .map_err(HttpDownloadError::IoError)?;

    let digest = digester.finalize();

    // Validate
    for (validator, expected, kind) in validators {
        let obtained = match validator {
            Validator::PrimaryDigest => digest.raw_digest().to_string(),
            Validator::ExtraDigest(hasher) => hex::encode(hasher.finalize()),
        };

        if expected != obtained {
            return Err(HttpDownloadError::InvalidChecksum(
                kind,
                expected.to_owned(),
                obtained,
                url.to_owned(),
            ));
        }
    }

    Ok(digest)
}

#[cfg(test)]
mod test {
    use assert_matches::assert_matches;
    use buck2_common::cas_digest::testing;
    use futures::stream;

    use super::*;

    async fn do_test(
        digest_config: CasDigestConfig,
        checksum: &Checksum,
    ) -> Result<(FileDigest, Vec<u8>), HttpDownloadError> {
        let mut out = Vec::new();

        let digest = copy_and_hash(
            "test",
            "test",
            stream::iter(vec![Ok(Bytes::from("foo")), Ok(Bytes::from("bar"))]),
            &mut out,
            digest_config,
            checksum,
        )
        .await?;

        Ok((digest, out))
    }

    #[tokio::test]
    async fn test_copy_and_hash_ok() -> anyhow::Result<()> {
        let (digest, bytes) = do_test(
            testing::blake3(),
            &Checksum::Both {
                sha1: Arc::from("8843d7f92416211de9ebb963ff4ce28125932878"),
                sha256: Arc::from(
                    "c3ab8ff13720e8ad9047dd39466b3c8974e592c2fa383d4a3960714caef0c4f2",
                ),
            },
        )
        .await?;

        assert_eq!(
            digest.to_string(),
            "aa51dcd43d5c6c5203ee16906fd6b35db298b9b2e1de3fce81811d4806b76b7d:6"
        );

        assert_eq!(std::str::from_utf8(&bytes).unwrap(), "foobar");

        Ok(())
    }

    #[tokio::test]
    async fn test_copy_and_hash_invalid_primary_hash() -> anyhow::Result<()> {
        assert_matches!(
            do_test(testing::sha1(), &Checksum::Sha1(Arc::from("oops")),).await,
            Err(HttpDownloadError::InvalidChecksum(..))
        );

        assert_matches!(
            do_test(testing::sha256(), &Checksum::Sha256(Arc::from("oops")),).await,
            Err(HttpDownloadError::InvalidChecksum(..))
        );

        Ok(())
    }

    #[tokio::test]
    async fn test_copy_and_hash_invalid_secondary_hash() -> anyhow::Result<()> {
        assert_matches!(
            do_test(testing::blake3(), &Checksum::Sha1(Arc::from("oops")),).await,
            Err(HttpDownloadError::InvalidChecksum(..))
        );

        assert_matches!(
            do_test(testing::blake3(), &Checksum::Sha256(Arc::from("oops")),).await,
            Err(HttpDownloadError::InvalidChecksum(..))
        );

        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 std::fmt;
use std::sync::Arc;

use allocative::Allocative;
use async_trait::async_trait;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_common::file_ops::FileMetadata;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_core::base_deferred_key::BaseDeferredKey;
use buck2_core::directory::DirectoryEntry;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_events::dispatch::EventDispatcher;
use chrono::DateTime;
use chrono::Duration;
use chrono::Utc;
use derive_more::Display;
use dice::UserComputationData;
use dupe::Dupe;
use futures::stream::BoxStream;
use futures::stream::TryStreamExt;
use more_futures::cancellation::CancellationContext;
use thiserror::Error;

use crate::artifact_value::ArtifactValue;
use crate::directory::ActionDirectoryEntry;
use crate::directory::ActionDirectoryMember;
use crate::directory::ActionImmutableDirectory;
use crate::directory::ActionSharedDirectory;
use crate::execute::action_digest::TrackedActionDigest;
#[cfg(any(fbcode_build, cargo_internal_build))]
use crate::materialize::eden_api::EdenBuckOut;
use crate::materialize::http::Checksum;

// Add a stub EdenBuckOut for when we don't have Eden output enabled
#[cfg(not(any(fbcode_build, cargo_internal_build)))]
pub struct EdenBuckOut {
    not_implemented: !,
}

#[cfg(not(any(fbcode_build, cargo_internal_build)))]
impl EdenBuckOut {
    pub async fn remove_paths_recursive(
        &self,
        _project_fs: &buck2_core::fs::project::ProjectRoot,
        _paths: Vec<ProjectRelativePathBuf>,
        _cancellations: &CancellationContext<'_>,
    ) -> anyhow::Result<()> {
        self.not_implemented
    }

    pub async fn set_path_object_id(
        &self,
        __path: &ProjectRelativePathBuf,
        __value: &ArtifactValue,
    ) -> anyhow::Result<()> {
        self.not_implemented
    }
}

pub struct WriteRequest {
    pub path: ProjectRelativePathBuf,
    pub content: Vec<u8>,
    pub is_executable: bool,
}

#[derive(Error, Debug)]
pub enum MaterializationError {
    #[error("Error materializing artifact at path `{}`", .path)]
    Error {
        path: ProjectRelativePathBuf,

        #[source]
        source: anyhow::Error,
    },

    /// The artifact wasn't found. This typically means it expired in the CAS.
    #[error(
        "Your build requires materializing an artifact that has expired in the \
        RE CAS and Buck does not have it (path: {}, digest origin: {}). \
        This likely happened because your Buck daemon \
        has been online for a long time. This error is currently unrecoverable. \
        To proceed, you should restart Buck using `buck2 killall` (debug info: {})",
        .path,
        .info.origin.as_display_for_not_found(),
        .debug
    )]
    NotFound {
        path: ProjectRelativePathBuf,
        info: Arc<CasDownloadInfo>,
        debug: Arc<str>,
    },

    #[error("Error inserting entry into materializer state sqlite for artifact at `{}`", .path)]
    SqliteDbError {
        path: ProjectRelativePathBuf,

        #[source]
        source: anyhow::Error,
    },
}

/// A trait providing methods to asynchronously materialize artifacts.
///
/// # Invariants
///
/// 0. Before accessing an artifact on disk, `ensure_materialized` must be
///    called with its path. Alternatively, `get_materialized_file_paths` can
///    be used to *try* to find an already materialized path of a file with the
///    same contents as the desired path, if the contents are relevant but the
///    path isn't.
///
/// 1. If an artifact is declared twice followed by `ensure_materialized`, the
///    latest declaration is the one guaranteed to have been materialized.
///    However, if an external source modifies the artifact on disk after it's
///    been declared, there are no guarantees that after `ensure_materialized`
///    the declared artifact will be on disk (due to race conditions).
///
/// 2. If `declare_*` is called concurrently for two artifacts with conflicting
///    paths or subpaths, there are no guarantees on which of the declarations
///    will be used when materializing - possibly a mix of both.
///
/// 3. If `ensure_materialized` is called on a path that wasn't previously
///    declared, that path is not materialized and no errors are raised for it.
///
/// 4. Declare may delete any existing paths that conflict with the path that was
///    declared.
#[async_trait]
pub trait Materializer: Allocative + Send + Sync + 'static {
    /// The name of this materializer, for logging.
    fn name(&self) -> &str;

    /// Declare that a set of artifacts exist on disk already.
    async fn declare_existing(
        &self,
        artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
    ) -> anyhow::Result<()>;

    async fn declare_copy_impl(
        &self,
        path: ProjectRelativePathBuf,
        value: ArtifactValue,
        srcs: Vec<CopiedArtifact>,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()>;

    async fn declare_cas_many_impl<'a, 'b>(
        &self,
        info: Arc<CasDownloadInfo>,
        artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()>;

    async fn declare_http(
        &self,
        path: ProjectRelativePathBuf,
        info: HttpDownloadInfo,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()>;

    /// Write contents to paths. The output is ordered in the same order as the input. Implicitly
    /// cleans up paths that the WriteRequest declares.
    async fn declare_write<'a>(
        &self,
        gen: Box<dyn FnOnce() -> anyhow::Result<Vec<WriteRequest>> + Send + 'a>,
    ) -> anyhow::Result<Vec<ArtifactValue>>;

    /// Ask the materializer if the artifacts at the set of paths match what is on disk or
    /// declared. Returns Ok(Ok) if they do and Ok(Err) if they don't. It's a result not a boolean
    /// so you can't ignore it.
    async fn declare_match(
        &self,
        artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
    ) -> anyhow::Result<DeclareMatchOutcome>;

    /// Declare an artifact at `path` exists. This will overwrite any pre-existing materialization
    /// methods for this file and indicate that no materialization is necessary.
    async fn invalidate(&self, path: ProjectRelativePathBuf) -> anyhow::Result<()> {
        self.invalidate_many(vec![path]).await
    }

    /// Declare an artifact at `path` exists. This will overwrite any pre-existing materialization
    /// methods for this file and indicate that no materialization is necessary.
    async fn invalidate_many(&self, paths: Vec<ProjectRelativePathBuf>) -> anyhow::Result<()>;

    /// Materialize artifacts paths. Returns a Stream with each element corresponding to one of the
    /// input paths, in the order that they were passed. This method provides access to
    /// individual results, but code that doesn't care about why an individual materialization
    /// failed should probably use `ensure_materialized` instead.
    async fn materialize_many(
        &self,
        artifact_paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<BoxStream<'static, Result<(), MaterializationError>>>;

    /// Given a list of artifact paths, blocks until all previously declared
    /// artifacts on that list are materialized. An [`Err`] is returned if the
    /// materialization fails for one or more of these paths.
    ///
    /// Paths that weren't previously declared are not materialized and no
    /// errors are raised for them.
    async fn ensure_materialized(
        &self,
        artifact_paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<()> {
        Ok(self
            .materialize_many(artifact_paths)
            .await?
            .try_collect()
            .await?)
    }

    /// Similar to `ensure_materialized`, but it relaxes its most important
    /// invariant: there's no guarantee that the artifact will be materialized
    /// after calling this method. It's meant for final artifacts that are NOT
    /// required by further build steps and therefore can be skipped in some
    /// cases.
    ///
    /// The materializer returns `Ok(true)` if the materialization succeeded,
    /// `Ok(false)` if it was skipped, and [`Err`] if it was tried but failed.
    ///
    /// Calling this on an artifact that was never declared is undefined
    /// behavior.
    async fn try_materialize_final_artifact(
        &self,
        artifact_path: ProjectRelativePathBuf,
    ) -> anyhow::Result<bool>;

    /// Given a `file_path` whose contents we are interested in, *tries* to
    /// find a materialized path with the same contents. It returns [`None`] if
    /// the path leads to a file that needs to be fetched from the CAS.
    ///
    /// If none of the declared artifacts contains `file_path`, the
    /// materializer can't tell whether that's because it was materialized by
    /// an external source, or because it doesn't exist. Since this usually
    /// means the path is a source file (not an output), for simplicity the
    /// materializer returns `Some(file_path)`, not [`None`].
    ///
    /// TODO(rafaelc): analyze if we can get rid of this method without making
    /// Buck2 significantly slower, after we stop deferring local copies of
    /// artifacts that are already on disk.
    async fn get_materialized_file_paths(
        &self,
        file_paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<Vec<Result<ProjectRelativePathBuf, ArtifactNotMaterializedReason>>>;

    /// Expose Eden based buck-out if the materializer is Eden
    /// Return None if not based on Eden.
    fn eden_buck_out(&self) -> Option<&EdenBuckOut> {
        None
    }

    fn as_deferred_materializer_extension(&self) -> Option<&dyn DeferredMaterializerExtensions> {
        None
    }

    /// Currently no-op for all materializers except deferred materializer
    fn log_materializer_state(&self, _events: &EventDispatcher) {}

    /// Inject stats into a snapshot. This is also used only for the deferred materializer at this
    /// time.
    fn add_snapshot_stats(&self, _snapshot: &mut buck2_data::Snapshot) {}
}

#[derive(Copy, Clone, Dupe, Debug)]
#[must_use]
pub enum DeclareMatchOutcome {
    Match,
    NotMatch,
}

impl From<bool> for DeclareMatchOutcome {
    fn from(is_match: bool) -> Self {
        if is_match {
            Self::Match
        } else {
            Self::NotMatch
        }
    }
}

impl DeclareMatchOutcome {
    pub fn is_match(self) -> bool {
        match self {
            Self::Match => true,
            Self::NotMatch => false,
        }
    }
}

impl dyn Materializer {
    /// Declare an artifact at `path` whose files can be materialized by doing
    /// a local copy. Implicitly cleans up `path`.
    ///
    /// The dest of each copied artifact in `srcs` must either be equal to
    /// `path`, or be a subpath of `path`; otherwise, [`Err`] is returned.
    pub async fn declare_copy(
        &self,
        path: ProjectRelativePathBuf,
        value: ArtifactValue,
        srcs: Vec<CopiedArtifact>,
        cancellations: &CancellationContext<'_>,
    ) -> anyhow::Result<()> {
        self.check_declared_external_symlink(&value)?;
        self.declare_copy_impl(path, value, srcs, cancellations)
            .await
    }

    /// Declares a list of artifacts whose files can be materialized by
    /// downloading from the CAS.
    pub async fn declare_cas_many<'a, 'b>(
        &self,
        info: Arc<CasDownloadInfo>,
        artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
        cancellations: &CancellationContext<'_>,
    ) -> anyhow::Result<()> {
        for (_, value) in artifacts.iter() {
            self.check_declared_external_symlink(value)?;
        }
        self.declare_cas_many_impl(info, artifacts, cancellations)
            .await
    }

    /// External symlink is a hack used to resolve the symlink to the correct external hack.
    /// No external symlink should be declared on the materializer with a non-empty remaining
    /// path. This function runs a check on all declared artifacts and returns `Err` if they
    /// are external symlinks with an existing value on `remaining_path` and `Ok` otherwise.
    fn check_declared_external_symlink(&self, value: &ArtifactValue) -> anyhow::Result<()> {
        match value.entry() {
            DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(external_symlink)) => {
                match external_symlink.remaining_path() {
                    Some(_) => {
                        return Err(anyhow::anyhow!(
                            "Internal error: external symlink should not be declared on materializer with non-empty remaining path: '{}'",
                            external_symlink.dupe().to_path_buf().display()
                        ));
                    }
                    None => {}
                }
            }
            _ => {}
        }
        Ok(())
    }
}

/// Information to perform the copy of an artifact from `src` to `dest`.
#[derive(Debug)]
pub struct CopiedArtifact {
    /// Source path of the copied artifact.
    pub src: ProjectRelativePathBuf,
    /// Destination path.
    pub dest: ProjectRelativePathBuf,
    /// Entry of the artifact at `dest`.
    pub dest_entry: ActionDirectoryEntry<ActionImmutableDirectory>,
}

impl CopiedArtifact {
    pub fn new(
        src: ProjectRelativePathBuf,
        dest: ProjectRelativePathBuf,
        dest_entry: ActionDirectoryEntry<ActionImmutableDirectory>,
    ) -> Self {
        Self {
            src,
            dest,
            dest_entry,
        }
    }
}

#[derive(Debug)]
pub enum CasDownloadInfoOrigin {
    /// Declared by an action that executed on RE.
    Execution(ActionExecutionOrigin),

    /// Simply declared by an action.
    Declared,
}

#[derive(Debug)]
pub struct ActionExecutionOrigin {
    /// Digest of the action that led us to discover this CAS object.
    action_digest: TrackedActionDigest,

    /// When did we learn of the connection between this digest and the download it allows. This
    /// typically represents how much time has passed since we executed the action or hit in the
    /// action cache.
    ///
    /// NOTE: we do not store this as `std::time::Instant`, because `Instant::elapsed()` does
    /// not necessarily track wall time (which is what we rather care about when dealing with
    /// TTLs received from RE). In particular, `Instant::elapsed()` can be *very* far off if
    /// the system went to sleep (on MacOS, it will not increase during that time!).
    ///
    /// See: <https://github.com/rust-lang/rust/issues/79462>
    action_instant: DateTime<Utc>,

    /// The TTL we retrieved from RE for this action.
    ttl: Duration,
}

impl ActionExecutionOrigin {
    fn action_age(&self) -> Duration {
        // NOTE: This might return a negative duration if our time skewed a lot, but that's
        // actually totally fine, since we only care to know if this is < TTL.
        Utc::now() - self.action_instant
    }
}

impl fmt::Display for CasDownloadInfoOrigin {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        match self {
            Self::Execution(execution) => {
                write!(
                    f,
                    "{} retrieved {:.3} seconds ago with ttl = {:.3} seconds",
                    execution.action_digest,
                    execution.action_age().num_seconds(),
                    execution.ttl.num_seconds()
                )?;
            }
            Self::Declared => {
                write!(f, "declared")?;
            }
        }

        Ok(())
    }
}

impl CasDownloadInfoOrigin {
    pub fn as_display_for_not_found(&self) -> CasDownloadInfoOriginNotFound<'_> {
        CasDownloadInfoOriginNotFound { inner: self }
    }

    /// Does the action cache guarantee this result exist? We expect the action cache to always
    /// return a TTL that is lower than the TTL of any of the outputs referenced by the action.
    pub fn guaranteed_by_action_cache(&self) -> bool {
        match self {
            Self::Execution(execution) => execution.action_age() < execution.ttl,
            Self::Declared => {
                // If the output was just declared, then there are no action cache guarantees.
                false
            }
        }
    }
}

/// A Display wrapper for CasDownloadInfoOrigin in cases where this origin was not found (in those
/// cases we want to report potential action cache corruption).
pub struct CasDownloadInfoOriginNotFound<'a> {
    inner: &'a CasDownloadInfoOrigin,
}

impl fmt::Display for CasDownloadInfoOriginNotFound<'_> {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        self.inner.fmt(f)?;
        if self.inner.guaranteed_by_action_cache() {
            write!(f, " (not expired: action cache corruption)")?;
        }

        Ok(())
    }
}

/// Information about a CAS download we might require when an artifact is not materialized.
#[derive(Debug)]
pub struct CasDownloadInfo {
    pub origin: CasDownloadInfoOrigin,
    /// RE Use case to use when downloading this
    pub re_use_case: RemoteExecutorUseCase,
}

impl CasDownloadInfo {
    pub fn new_execution(
        action_digest: TrackedActionDigest,
        re_use_case: RemoteExecutorUseCase,
        action_instant: DateTime<Utc>,
        ttl: Duration,
    ) -> Self {
        Self {
            origin: CasDownloadInfoOrigin::Execution(ActionExecutionOrigin {
                action_digest,
                action_instant,
                ttl,
            }),
            re_use_case,
        }
    }

    pub fn new_declared(re_use_case: RemoteExecutorUseCase) -> Self {
        Self {
            origin: CasDownloadInfoOrigin::Declared,
            re_use_case,
        }
    }

    pub fn action_age(&self) -> Option<Duration> {
        match &self.origin {
            CasDownloadInfoOrigin::Execution(execution) => Some(execution.action_age()),
            CasDownloadInfoOrigin::Declared => None,
        }
    }

    pub fn action_digest(&self) -> Option<&TrackedActionDigest> {
        match &self.origin {
            CasDownloadInfoOrigin::Execution(execution) => Some(&execution.action_digest),
            CasDownloadInfoOrigin::Declared => None,
        }
    }
}

/// Information about a CAS download we might require when an artifact is not materialized.
#[derive(Debug, Display)]
#[display(fmt = "{} declared by {}", "self.url", "self.owner")]
pub struct HttpDownloadInfo {
    /// URL to download the file from.
    pub url: Arc<str>,

    /// Size, whether the file is executable. Also contains a digest, which is a bit of a shame
    /// since it's duplicative of checksum.
    pub metadata: FileMetadata,

    /// Checksum for the file, to valiate before downloading.
    pub checksum: Checksum,

    /// Target that declared the action.
    pub owner: BaseDeferredKey,
}

#[derive(Debug, Error)]
pub enum ArtifactNotMaterializedReason {
    #[error(
        "The artifact at path '{}' ({}) was produced by a RE action ({}), \
        but has not been downloaded",
        .path,
        .entry,
        .info.origin.as_display_for_not_found()
    )]
    RequiresCasDownload {
        path: ProjectRelativePathBuf,
        entry: ActionDirectoryEntry<ActionImmutableDirectory>,
        info: Arc<CasDownloadInfo>,
    },

    #[error("The artifact at path '{}' has not been downloaded yet", .path)]
    RequiresMaterialization { path: ProjectRelativePathBuf },

    #[error(
        "The artifact at path '{}' points into an entry ({}) \
        produced by a RE action ({}), but that entry does not contain \
        this exact path.",
        .path,
        .entry,
        .info.origin
    )]
    DeferredMaterializerCorruption {
        path: ProjectRelativePathBuf,
        entry: ActionDirectoryEntry<ActionSharedDirectory>,
        info: Arc<CasDownloadInfo>,
    },
}

// ==== dice ====

pub trait SetMaterializer {
    fn set_materializer(&mut self, materializer: Arc<dyn Materializer>);
}

pub trait HasMaterializer {
    fn get_materializer(&self) -> Arc<dyn Materializer>;
}

impl SetMaterializer for UserComputationData {
    fn set_materializer(&mut self, materializer: Arc<dyn Materializer>) {
        self.data.set(materializer);
    }
}

impl HasMaterializer for UserComputationData {
    fn get_materializer(&self) -> Arc<dyn Materializer> {
        self.data
            .get::<Arc<dyn Materializer>>()
            .expect("Materializer should be set")
            .dupe()
    }
}

#[derive(Clone, Copy, Debug, Dupe)]
pub enum MaterializationMethod {
    /// Materialize all immediately as they are declared
    Immediate,
    /// Materialize only when needed
    Deferred,
    /// Materialize only when needed, do not materialize final artifacts
    DeferredSkipFinalArtifacts,
    /// Let Eden delegate materialzation
    Eden,
}

#[derive(Debug, Error)]
pub enum MaterializationMethodError {
    #[error(
        "Invalid value for buckconfig `[buck2] materializations`. Got `{0}`. Expected one of `all`, `deferred`, `deferred_skip_final_artifacts` or `eden`."
    )]
    InvalidValueForConfig(String),
}

impl MaterializationMethod {
    pub fn try_new_from_config(legacy_config: Option<&LegacyBuckConfig>) -> anyhow::Result<Self> {
        Self::try_new_from_config_value(
            legacy_config.and_then(|c| c.get("buck2", "materializations")),
        )
    }

    fn try_new_from_config_value(config_value: Option<&str>) -> anyhow::Result<Self> {
        match config_value {
            None | Some("") | Some("all") => Ok(MaterializationMethod::Immediate),
            Some("deferred") => Ok(MaterializationMethod::Deferred),
            Some("deferred_skip_final_artifacts") => {
                Ok(MaterializationMethod::DeferredSkipFinalArtifacts)
            }
            Some("eden") => Ok(MaterializationMethod::Eden),
            Some(v) => Err(MaterializationMethodError::InvalidValueForConfig(v.to_owned()).into()),
        }
    }
}

/// This trait provides a level of indirection since the concrete implementation of
/// `DeferredMaterializerEntry` lives in a crate that depends on this one.
pub trait DeferredMaterializerEntry: Send + Sync + std::fmt::Display {}

/// Obtain notifications for entries as they are materialized, and request eager materialization of
/// those paths.
#[async_trait]
pub trait DeferredMaterializerSubscription: Send + Sync {
    /// Get notifications for specific paths. This also implicitly requests their eager
    /// materialization.
    fn subscribe_to_paths(&mut self, paths: Vec<ProjectRelativePathBuf>);

    /// Stop getting notifications for specific paths. In-flight notifications may still be
    /// received.
    fn unsubscribe_from_paths(&mut self, paths: Vec<ProjectRelativePathBuf>);

    /// Await the next materialization on this subscription.
    async fn next_materialization(&mut self) -> Option<ProjectRelativePathBuf>;
}

/// Extensions to the Materializer trait that are only available in the Deferred materializer.
#[async_trait]
pub trait DeferredMaterializerExtensions: Send + Sync {
    fn iterate(
        &self,
    ) -> anyhow::Result<
        BoxStream<'static, (ProjectRelativePathBuf, Box<dyn DeferredMaterializerEntry>)>,
    >;

    /// Obtain a list of files that don't match their in-memory representation. This may not catch
    /// all discrepancies.
    fn fsck(&self) -> anyhow::Result<BoxStream<'static, (ProjectRelativePathBuf, anyhow::Error)>>;

    async fn refresh_ttls(&self, min_ttl: i64) -> anyhow::Result<()>;

    async fn get_ttl_refresh_log(&self) -> anyhow::Result<String>;

    async fn clean_stale_artifacts(
        &self,
        keep_since_time: DateTime<Utc>,
        dry_run: bool,
        tracked_only: bool,
    ) -> anyhow::Result<buck2_cli_proto::CleanStaleResponse>;

    async fn test_iter(&self, count: usize) -> anyhow::Result<String>;

    /// Create a new DeferredMaterializerSubscription.
    async fn create_subscription(
        &self,
    ) -> anyhow::Result<Box<dyn DeferredMaterializerSubscription>>;
}


/*
 * 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.
 */

#[cfg(any(fbcode_build, cargo_internal_build))]
pub mod eden_api;
pub mod http;

pub mod materializer;
pub mod nodisk;


/*
 * 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 async_trait::async_trait;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use gazebo::prelude::*;
use more_futures::cancellation::CancellationContext;

use crate::artifact_value::ArtifactValue;
use crate::materialize::materializer::ArtifactNotMaterializedReason;
use crate::materialize::materializer::CasDownloadInfo;
use crate::materialize::materializer::CopiedArtifact;
use crate::materialize::materializer::DeclareMatchOutcome;
use crate::materialize::materializer::HttpDownloadInfo;
use crate::materialize::materializer::MaterializationError;
use crate::materialize::materializer::Materializer;
use crate::materialize::materializer::WriteRequest;

/// Materializer that doesn't really materialize anything, analogous to
/// materializing to /dev/null. Meant to be used in unittests that need a
/// materializer, but don't want materializations to actually happen.
#[derive(Allocative)]
pub struct NoDiskMaterializer;

#[async_trait]
impl Materializer for NoDiskMaterializer {
    fn name(&self) -> &str {
        "noop"
    }

    async fn declare_existing(
        &self,
        _artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
    ) -> anyhow::Result<()> {
        // Nothing to do, we don't keep track of state;
        Ok(())
    }

    async fn declare_copy_impl(
        &self,
        _path: ProjectRelativePathBuf,
        _value: ArtifactValue,
        _srcs: Vec<CopiedArtifact>,
        _cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        Ok(())
    }

    async fn declare_cas_many_impl<'a, 'b>(
        &self,
        _info: Arc<CasDownloadInfo>,
        _artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
        _cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        Ok(())
    }

    async fn declare_http(
        &self,
        _path: ProjectRelativePathBuf,
        _info: HttpDownloadInfo,
        _cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        Ok(())
    }

    async fn declare_match(
        &self,
        _artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
    ) -> anyhow::Result<DeclareMatchOutcome> {
        // This materializer does not keep track of state
        Ok(DeclareMatchOutcome::NotMatch)
    }

    async fn declare_write<'a>(
        &self,
        _gen: Box<dyn FnOnce() -> anyhow::Result<Vec<WriteRequest>> + Send + 'a>,
    ) -> anyhow::Result<Vec<ArtifactValue>> {
        Err(anyhow::anyhow!("NoDiskMaterializer cannot write"))
    }

    async fn invalidate_many(&self, _paths: Vec<ProjectRelativePathBuf>) -> anyhow::Result<()> {
        Ok(())
    }

    async fn materialize_many(
        &self,
        artifact_paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<BoxStream<'static, Result<(), MaterializationError>>> {
        Ok(stream::iter(artifact_paths.into_iter().map(|_| Ok(()))).boxed())
    }

    async fn try_materialize_final_artifact(
        &self,
        _artifact_path: ProjectRelativePathBuf,
    ) -> anyhow::Result<bool> {
        Ok(false)
    }

    async fn get_materialized_file_paths(
        &self,
        paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<Vec<Result<ProjectRelativePathBuf, ArtifactNotMaterializedReason>>> {
        Ok(paths.into_map(Ok))
    }
}