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

use std::collections::HashMap;
use std::sync::Arc;

use anyhow::Context;
use buck2_common::file_ops::FileType;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_core::soft_error;
use buck2_events::dispatch::EventDispatcher;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::execute::clean_output_paths::CleanOutputPaths;
use chrono::DateTime;
use chrono::Utc;
use derivative::Derivative;
use dupe::Dupe;
use futures::future::BoxFuture;
use futures::FutureExt;
use more_futures::cancellation::CancellationContext;
use thiserror::Error;
use tokio::sync::oneshot::Sender;
use tracing::error;

use crate::materializers::deferred::extension::ExtensionCommand;
use crate::materializers::deferred::join_all_existing_futs;
use crate::materializers::deferred::ArtifactMaterializationData;
use crate::materializers::deferred::ArtifactMaterializationStage;
use crate::materializers::deferred::ArtifactTree;
use crate::materializers::deferred::DefaultIoHandler;
use crate::materializers::deferred::DeferredMaterializerCommandProcessor;
use crate::materializers::sqlite::MaterializerStateSqliteDb;

#[derive(Derivative)]
#[derivative(Debug)]
pub struct CleanStaleArtifacts {
    pub keep_since_time: DateTime<Utc>,
    pub dry_run: bool,
    pub tracked_only: bool,
    #[derivative(Debug = "ignore")]
    pub sender: Sender<BoxFuture<'static, anyhow::Result<buck2_cli_proto::CleanStaleResponse>>>,
    pub dispatcher: EventDispatcher,
}

fn skip_clean_response_with_message(
    message: &str,
) -> anyhow::Result<(
    BoxFuture<'static, anyhow::Result<()>>,
    buck2_cli_proto::CleanStaleResponse,
)> {
    Ok((
        futures::future::ready(Ok(())).boxed(),
        buck2_cli_proto::CleanStaleResponse {
            message: Some(message.to_owned()),
            stats: None,
        },
    ))
}

impl ExtensionCommand<DefaultIoHandler> for CleanStaleArtifacts {
    fn execute(
        self: Box<Self>,
        processor: &mut DeferredMaterializerCommandProcessor<DefaultIoHandler>,
    ) {
        let res = if let Some(sqlite_db) = processor.sqlite_db.as_mut() {
            if !processor.defer_write_actions {
                skip_clean_response_with_message(
                    "Skipping clean, set buck2.defer_write_actions to use clean --stale",
                )
            } else {
                gather_clean_futures_for_stale_artifacts(
                    &mut processor.tree,
                    self.keep_since_time,
                    self.dry_run,
                    self.tracked_only,
                    sqlite_db,
                    &processor.io,
                    processor.digest_config,
                    processor.cancellations,
                    &self.dispatcher,
                )
            }
        } else {
            skip_clean_response_with_message(
                "Skipping clean, set buck2.sqlite_materializer_state to use clean --stale",
            )
        };
        let fut = async move {
            let (fut, response) = res?;
            fut.await?;
            tracing::trace!("finished cleaning stale artifacts");
            Ok(response)
        }
        .boxed();
        let _ignored = self.sender.send(fut);
    }
}

#[derive(Debug, Clone, Error)]
#[error("Internal error: materializer state exists (num db entries: {}) but no artifacts were found by clean ({:?}). Not cleaning untracked artifacts.", .db_size, .stats)]
pub(crate) struct CleanStaleError {
    db_size: usize,
    stats: buck2_data::CleanStaleStats,
}

fn gather_clean_futures_for_stale_artifacts(
    tree: &mut ArtifactTree,
    keep_since_time: DateTime<Utc>,
    dry_run: bool,
    tracked_only: bool,
    sqlite_db: &mut MaterializerStateSqliteDb,
    io: &Arc<DefaultIoHandler>,
    digest: DigestConfig,
    cancellations: &'static CancellationContext,
    dispatcher: &EventDispatcher,
) -> anyhow::Result<(
    BoxFuture<'static, anyhow::Result<()>>,
    buck2_cli_proto::CleanStaleResponse,
)> {
    let gen_path = io
        .buck_out_path
        .join(ProjectRelativePathBuf::unchecked_new("gen".to_owned()));
    let gen_dir = io.fs.resolve(&gen_path);
    if !fs_util::try_exists(&gen_dir)? {
        return skip_clean_response_with_message("Nothing to clean");
    }
    tracing::trace!(gen_dir = %gen_dir, "Scanning");

    let mut stats = buck2_data::CleanStaleStats::default();
    let mut paths_to_remove = Vec::new();
    let mut paths_to_invalidate = Vec::new();

    if tracked_only {
        find_stale_tracked_only(tree, keep_since_time, &mut stats, &mut paths_to_invalidate)?
    } else {
        let gen_subtree = tree
            .get_subtree(&mut gen_path.iter())
            .context("Found a file where gen dir expected")?;

        let empty;

        let gen_subtree = match gen_subtree {
            Some(t) => t,
            None => {
                empty = HashMap::new();
                &empty
            }
        };

        StaleFinder {
            fs: &io.fs,
            dispatcher,
            keep_since_time,
            stats: &mut stats,
            paths_to_remove: &mut paths_to_remove,
            paths_to_invalidate: &mut paths_to_invalidate,
        }
        .visit_recursively(gen_path, gen_subtree)?;
    };

    // If no stale or retained artifact founds, the db should be empty.
    if stats.stale_artifact_count + stats.retained_artifact_count == 0 {
        // Just need to know if any entries exist, could be a simpler query.
        // Checking the db directly in case tree is somehow not in sync.
        let materializer_state = sqlite_db.materializer_state_table().read_all(digest)?;

        // Entries in the db should have been found in buck-out, return error and skip cleaning untracked artifacts.
        if !materializer_state.is_empty() {
            let error = CleanStaleError {
                db_size: materializer_state.len(),
                stats,
            };
            // quiet just because it's also returned, soft_error to log to scribe
            return Err(soft_error!("clean_stale_error", error.into(), quiet: true)?);
        }
    }

    let fut = if dry_run {
        futures::future::ready(Ok(())).boxed()
    } else {
        let io = io.dupe();

        stats.cleaned_path_count = paths_to_remove.len() as u64;
        stats.cleaned_artifact_count = stats.stale_artifact_count + stats.untracked_artifact_count;
        stats.cleaned_bytes = stats.untracked_bytes + stats.stale_bytes;

        let existing_futs =
            tree.invalidate_paths_and_collect_futures(paths_to_invalidate, Some(sqlite_db))?;

        async move {
            // Wait for all in-progress operations to finish on the paths we are about to
            // remove from disk.
            join_all_existing_futs(existing_futs).await?;

            // Then actually delete them. Note that we kick off one CleanOutputPaths per path. We
            // do this to get parallelism.
            futures::future::try_join_all(paths_to_remove.into_iter().map(|path| {
                io.io_executor.execute_io(
                    Box::new(CleanOutputPaths { paths: vec![path] }),
                    cancellations,
                )
            }))
            .await?;

            anyhow::Ok(())
        }
        .boxed()
    };

    Ok((
        fut,
        buck2_cli_proto::CleanStaleResponse {
            message: None,
            stats: Some(stats),
        },
    ))
}

/// Get file size or directory size, without following symlinks
pub fn get_size(path: &AbsNormPath) -> anyhow::Result<u64> {
    let mut result = 0;
    if path.is_dir() {
        for entry in fs_util::read_dir(path)? {
            result += get_size(&entry?.path())?;
        }
    } else {
        result = path.symlink_metadata()?.len();
    }
    Ok(result)
}

struct StaleFinder<'a> {
    fs: &'a ProjectRoot,
    dispatcher: &'a EventDispatcher,
    keep_since_time: DateTime<Utc>,
    stats: &'a mut buck2_data::CleanStaleStats,
    /// Those paths will be deleted on disk.
    paths_to_remove: &'a mut Vec<ProjectRelativePathBuf>,
    /// Those paths will be invalidated in the materiaizer.
    paths_to_invalidate: &'a mut Vec<ProjectRelativePathBuf>,
}

impl<'a> StaleFinder<'a> {
    /// Start from `path` and `subtree` and visit everything below.
    fn visit_recursively<'t>(
        &mut self,
        path: ProjectRelativePathBuf,
        subtree: &'t HashMap<FileNameBuf, ArtifactTree>,
    ) -> anyhow::Result<()> {
        let mut queue = vec![(path, subtree)];

        while let Some((path, tree)) = queue.pop() {
            self.visit(&path, tree, &mut queue)?;
        }

        Ok(())
    }

    /// Visit one directory.
    fn visit<'t>(
        &mut self,
        path: &ProjectRelativePath,
        subtree: &'t HashMap<FileNameBuf, ArtifactTree>,
        queue: &mut Vec<(
            ProjectRelativePathBuf,
            &'t HashMap<FileNameBuf, ArtifactTree>,
        )>,
    ) -> anyhow::Result<()> {
        let abs_path = self.fs.resolve(path);

        for child in fs_util::read_dir(&abs_path)? {
            let child = child?;

            let file_name = child.file_name();
            let file_name = file_name.to_str().and_then(|f| FileName::new(f).ok());

            let file_name = match file_name {
                Some(file_name) => file_name,
                None => {
                    // If the file name is invalid, then it can't be tracked by the materializer.
                    // We should ideally delete this, but currently we don't support doing that.
                    continue;
                }
            };

            let path = path.join(file_name);

            let file_type = FileType::from(child.file_type()?);

            let subtree = match subtree.get(file_name) {
                Some(subtree) => subtree,
                None => {
                    // This path is not tracked by the materializer, we can delete it.
                    tracing::trace!(path = %path, file_type = ?file_type, "marking as untracked");
                    self.stats.untracked_artifact_count += 1;
                    self.stats.untracked_bytes += get_size(&child.path())?;
                    if self.stats.untracked_artifact_count <= 2000 {
                        self.dispatcher.instant_event(buck2_data::UntrackedFile {
                            path: path.to_string(),
                            file_type: format!("{:?}", file_type),
                        });
                    }
                    self.paths_to_remove.push(path);
                    continue;
                }
            };

            match subtree {
                ArtifactTree::Tree(subtree) if file_type.is_dir() => {
                    queue.push((path, subtree));
                }
                ArtifactTree::Data(box ArtifactMaterializationData {
                    stage:
                        ArtifactMaterializationStage::Materialized {
                            active: false,
                            last_access_time,
                            metadata,
                        },
                    ..
                }) if *last_access_time < self.keep_since_time => {
                    // This is something we can invalidate.
                    tracing::trace!(path = %path, file_type = ?file_type, "marking as stale");
                    self.stats.stale_artifact_count += 1;
                    self.stats.stale_bytes += metadata.size();
                    self.paths_to_invalidate.push(path.clone());
                    self.paths_to_remove.push(path);
                }
                ArtifactTree::Data(box ArtifactMaterializationData {
                    stage: ArtifactMaterializationStage::Materialized { metadata, .. },
                    ..
                }) => {
                    tracing::trace!(path = %path, file_type = ?file_type, "marking as retained");
                    self.stats.retained_artifact_count += 1;
                    self.stats.retained_bytes += metadata.size();
                }
                _ => {
                    // What we have on disk does not match what we have in the materializer (which is
                    // not clean stale's problem to fix).
                    tracing::trace!(path = %path, file_type = ?file_type, "skipping");
                }
            }
        }

        Ok(())
    }
}

fn find_stale_tracked_only(
    tree: &ArtifactTree,
    keep_since_time: DateTime<Utc>,
    stats: &mut buck2_data::CleanStaleStats,
    paths_to_invalidate: &mut Vec<ProjectRelativePathBuf>,
) -> anyhow::Result<()> {
    for (f_path, v) in tree.iter_with_paths() {
        if let ArtifactMaterializationStage::Materialized {
            last_access_time,
            active,
            ..
        } = &v.stage
        {
            let path = ProjectRelativePathBuf::from(f_path);
            if *last_access_time < keep_since_time && !active {
                tracing::trace!(path = %path, "stale artifact");
                stats.stale_artifact_count += 1;
                paths_to_invalidate.push(path);
            } else {
                tracing::trace!(path = %path, "retaining artifact");
                stats.retained_artifact_count += 1;
            }
        }
    }
    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::fmt::Debug;
use std::fmt::Display;
use std::fmt::Write;
use std::sync::Arc;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_core::directory::DirectoryEntry;
use buck2_core::fs::fs_util;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_events::dispatch::get_dispatcher;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::materialize::materializer::DeferredMaterializerEntry;
use buck2_execute::materialize::materializer::DeferredMaterializerExtensions;
use buck2_execute::materialize::materializer::DeferredMaterializerSubscription;
use chrono::DateTime;
use chrono::Duration;
use chrono::TimeZone;
use chrono::Utc;
use derivative::Derivative;
use dupe::Dupe;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use tokio::sync::mpsc;
use tokio::sync::mpsc::UnboundedSender;
use tokio::sync::oneshot;
use tokio::sync::oneshot::Sender;
use tokio::task::JoinHandle;
use tokio_stream::wrappers::UnboundedReceiverStream;

use crate::materializers::deferred::clean_stale::CleanStaleArtifacts;
use crate::materializers::deferred::io_handler::create_ttl_refresh;
use crate::materializers::deferred::subscriptions::MaterializerSubscriptionOperation;
use crate::materializers::deferred::ArtifactMaterializationMethod;
use crate::materializers::deferred::ArtifactMaterializationStage;
use crate::materializers::deferred::DefaultIoHandler;
use crate::materializers::deferred::DeferredMaterializer;
use crate::materializers::deferred::DeferredMaterializerCommandProcessor;
use crate::materializers::deferred::MaterializerCommand;

pub(super) trait ExtensionCommand<T>: Debug + Sync + Send + 'static {
    fn execute(self: Box<Self>, processor: &mut DeferredMaterializerCommandProcessor<T>);
}

#[derive(Debug)]
enum PathData {
    Materialized {
        ts: DateTime<Utc>,
        size: Option<u64>,
    },
    Declared(Arc<ArtifactMaterializationMethod>),
}

impl Display for PathData {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        match self {
            PathData::Materialized { ts, size } => {
                if let Some(size) = size {
                    write!(f, "materialized (ts={:?}, size={})", ts, size)
                } else {
                    write!(f, "materialized (ts={:?})", ts)
                }
            }
            PathData::Declared(method) => write!(f, "declared: {}", method),
        }
    }
}

impl DeferredMaterializerEntry for PathData {}

#[derive(Derivative)]
#[derivative(Debug)]
struct Iterate {
    /// This is for debug commands so we use an unbounded channel to avoid locking up the
    /// materializer command thread.
    #[derivative(Debug = "ignore")]
    sender: UnboundedSender<(ProjectRelativePathBuf, Box<dyn DeferredMaterializerEntry>)>,
}

impl ExtensionCommand<DefaultIoHandler> for Iterate {
    fn execute(
        self: Box<Self>,
        processor: &mut DeferredMaterializerCommandProcessor<DefaultIoHandler>,
    ) {
        for (path, data) in processor.tree.iter_with_paths() {
            let path_data = match &data.stage {
                ArtifactMaterializationStage::Declared { method, .. } => {
                    PathData::Declared(method.dupe())
                }
                ArtifactMaterializationStage::Materialized {
                    last_access_time,
                    metadata,
                    ..
                } => {
                    let size = match &metadata.0 {
                        DirectoryEntry::Dir(meta) => meta.total_size,
                        DirectoryEntry::Leaf(ActionDirectoryMember::File(file_metadata)) => {
                            file_metadata.digest.size()
                        }
                        DirectoryEntry::Leaf(_) => 0,
                    };
                    // drop nano-seconds
                    let ts = Utc
                        .timestamp_opt(last_access_time.timestamp(), 0)
                        .single()
                        .unwrap();
                    PathData::Materialized {
                        ts,
                        size: Some(size),
                    }
                }
            };

            let path = ProjectRelativePathBuf::from(path);

            match self.sender.send((path, Box::new(path_data) as _)) {
                Ok(..) => {}
                Err(..) => break, // No use sending more if the client disconnected.
            }
        }
    }
}

#[derive(Derivative)]
#[derivative(Debug)]
struct Fsck {
    /// This is for debug commands so we use an unbounded channel to avoid locking up the
    /// materializer command thread.
    #[derivative(Debug = "ignore")]
    sender: UnboundedSender<(ProjectRelativePathBuf, anyhow::Error)>,
}

impl ExtensionCommand<DefaultIoHandler> for Fsck {
    fn execute(
        self: Box<Self>,
        processor: &mut DeferredMaterializerCommandProcessor<DefaultIoHandler>,
    ) {
        for (path, data) in processor.tree.iter_with_paths() {
            match &data.stage {
                ArtifactMaterializationStage::Declared { .. } => {
                    continue;
                }
                ArtifactMaterializationStage::Materialized { .. } => {}
            };

            // We actually block the thread here. This is to ensure we don't try to delete things
            // when we check them. This is primarily a debug command. We don't run this while
            // actual things are in flight.

            let path = ProjectRelativePathBuf::from(path);
            let res = fs_util::symlink_metadata(processor.io.fs.resolve(&path));
            match res {
                Ok(..) => {}
                Err(e) => {
                    let _ignored = self.sender.send((path, e));
                }
            }
        }
    }
}

#[derive(Derivative)]
#[derivative(Debug)]
struct RefreshTtls {
    sender: Sender<Option<JoinHandle<anyhow::Result<()>>>>,
    min_ttl: i64,
}

impl ExtensionCommand<DefaultIoHandler> for RefreshTtls {
    fn execute(
        self: Box<Self>,
        processor: &mut DeferredMaterializerCommandProcessor<DefaultIoHandler>,
    ) {
        let task = create_ttl_refresh(
            &processor.tree,
            &processor.io.re_client_manager,
            Duration::seconds(self.min_ttl),
            processor.digest_config,
        )
        .map(|f| processor.rt.spawn(f));
        let _ignored = self.sender.send(task);
    }
}

#[derive(Derivative)]
#[derivative(Debug)]
struct GetTtlRefreshLog {
    sender: Sender<String>,
}

impl ExtensionCommand<DefaultIoHandler> for GetTtlRefreshLog {
    fn execute(
        self: Box<Self>,
        processor: &mut DeferredMaterializerCommandProcessor<DefaultIoHandler>,
    ) {
        // We normally poll this very lazily, so actually force it to happen here.
        processor.poll_current_ttl_refresh();

        let mut out = String::new();

        for entry in &processor.ttl_refresh_history {
            write!(&mut out, "{:?}\t", entry.at).unwrap();
            match &entry.outcome {
                None => {
                    writeln!(&mut out, "SKIP").unwrap();
                }
                Some(Ok(())) => {
                    writeln!(&mut out, "OK").unwrap();
                }
                Some(Err(e)) => {
                    writeln!(&mut out, "ERR\t{:#}", e).unwrap();
                }
            }
        }

        let _ignored = self.sender.send(out);
    }
}

#[derive(Derivative)]
#[derivative(Debug)]
struct TestIter {
    sender: Sender<String>,
    count: usize,
}

impl ExtensionCommand<DefaultIoHandler> for TestIter {
    fn execute(
        self: Box<Self>,
        processor: &mut DeferredMaterializerCommandProcessor<DefaultIoHandler>,
    ) {
        let mut out = String::new();

        let now = std::time::Instant::now();

        for _i in 0..self.count {
            let it = processor.tree.iter_without_paths();

            for e in it {
                let _e = e;
            }
        }

        writeln!(
            &mut out,
            "Elapsed for iter() ({} times): {:?}",
            self.count,
            now.elapsed()
        )
        .unwrap();

        let now = std::time::Instant::now();

        for _i in 0..self.count {
            let it = processor.tree.iter_with_paths();

            for e in it {
                let _e = e;
            }
        }

        writeln!(
            &mut out,
            "Elapsed for iter().with_paths() ({} times): {:?}",
            self.count,
            now.elapsed()
        )
        .unwrap();

        let _ignored = self.sender.send(out);
    }
}

#[async_trait]
impl DeferredMaterializerExtensions for DeferredMaterializer {
    fn iterate(
        &self,
    ) -> anyhow::Result<
        BoxStream<'static, (ProjectRelativePathBuf, Box<dyn DeferredMaterializerEntry>)>,
    > {
        let (sender, receiver) = mpsc::unbounded_channel();
        self.command_sender.send(MaterializerCommand::Extension(
            Box::new(Iterate { sender }) as _
        ))?;
        Ok(UnboundedReceiverStream::new(receiver).boxed())
    }

    fn fsck(&self) -> anyhow::Result<BoxStream<'static, (ProjectRelativePathBuf, anyhow::Error)>> {
        let (sender, receiver) = mpsc::unbounded_channel();
        self.command_sender.send(MaterializerCommand::Extension(
            Box::new(Fsck { sender }) as _
        ))?;
        Ok(UnboundedReceiverStream::new(receiver).boxed())
    }

    async fn refresh_ttls(&self, min_ttl: i64) -> anyhow::Result<()> {
        let (sender, receiver) = oneshot::channel();
        self.command_sender
            .send(MaterializerCommand::Extension(
                Box::new(RefreshTtls { sender, min_ttl }) as _,
            ))?;
        match receiver.await.context("No response from materializer")? {
            Some(task) => task
                .await
                .context("Refresh task aborted")?
                .context("Refresh failed")?,
            None => {}
        };
        Ok(())
    }

    async fn get_ttl_refresh_log(&self) -> anyhow::Result<String> {
        let (sender, receiver) = oneshot::channel();
        self.command_sender
            .send(MaterializerCommand::Extension(
                Box::new(GetTtlRefreshLog { sender }) as _,
            ))?;
        receiver.await.context("No response from materializer")
    }

    async fn clean_stale_artifacts(
        &self,
        keep_since_time: DateTime<Utc>,
        dry_run: bool,
        tracked_only: bool,
    ) -> anyhow::Result<buck2_cli_proto::CleanStaleResponse> {
        let dispatcher = get_dispatcher();
        let (sender, recv) = oneshot::channel();
        self.command_sender
            .send(MaterializerCommand::Extension(Box::new(
                CleanStaleArtifacts {
                    keep_since_time,
                    dry_run,
                    tracked_only,
                    sender,
                    dispatcher,
                },
            )))?;
        recv.await?.await
    }

    async fn test_iter(&self, count: usize) -> anyhow::Result<String> {
        let (sender, receiver) = oneshot::channel();
        self.command_sender
            .send(MaterializerCommand::Extension(
                Box::new(TestIter { sender, count }) as _,
            ))?;
        receiver.await.context("No response from materializer")
    }

    async fn create_subscription(
        &self,
    ) -> anyhow::Result<Box<dyn DeferredMaterializerSubscription>> {
        let (sender, receiver) = oneshot::channel();
        self.command_sender.send(MaterializerCommand::Subscription(
            MaterializerSubscriptionOperation::Create { sender },
        ))?;
        Ok(Box::new(receiver.await.context("No response from materializer")?) as _)
    }
}


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

//! Data structure akin to a map, but where the key is a sequence.
//!
//! Restrictions:
//! - Storing both a key A and a key B which is a prefix of A is not possible.
//!
//! Special operations:
//! - Using key A to search for a value at a key B when B is a prefix of A.
//!
//! This is useful when artifacts are directories and we need to query the map
//! to figure out which artifact a path belongs to. E.g. we have an artifact at
//! "foo/bar", and we need to find out which artifact "foo/bar/c" belongs to.

use std::borrow::Borrow;
use std::collections::hash_map::Entry;
use std::collections::hash_map::IntoIter;
use std::collections::hash_map::Iter;
use std::collections::HashMap;
use std::hash::Hash;
use std::marker::PhantomData;

use buck2_core::fs::paths::file_name::FileNameBuf;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePathBuf;

pub type FileTree<V> = DataTree<FileNameBuf, V>;

/// Tree that stores data in the leaves. Think of the key as the path to the
/// leaf containing the value. The data/value is of type `V`, and each edge
/// is of type `K` (making the key to a value a sequence of `K`).
/// TODO(scottcao): This trie is not implemented properly. It should be merged into the directory trie
/// we have at buck2_core/src/directory.
#[derive(Debug)]
pub enum DataTree<K, V> {
    /// Stores data of type `V` with key of type `Iterator<Item = K>`.
    Tree(HashMap<K, DataTree<K, V>>),
    Data(V),
}

impl<K: 'static + Eq + Hash + Clone, V: 'static> DataTree<K, V> {
    pub fn new() -> Self {
        Self::Tree(HashMap::new())
    }

    /// Gets the value at `key` or one of its prefixes, and returns it.
    ///
    /// When a value is found and [`Some`] is returned, it's guaranteed that
    /// only enough of `key` to find the returned value was consumed.
    /// E.g. if `key` is (A, B, C, D) and there's a value present at (A, B),
    /// then after this method returns (C, D) can still be consumed from `key`.
    ///
    /// There are no guarantees on how much is consumed from `key` when
    /// [`None`] is returned.
    pub fn prefix_get<'a, I, Q>(&self, key: &mut I) -> Option<&V>
    where
        K: 'a + Borrow<Q>,
        Q: 'a + Hash + Eq + ?Sized,
        I: Iterator<Item = &'a Q>,
    {
        if let Self::Data(data) = self {
            // return early so we don't consume from key_iter unnecessarily
            return Some(data);
        }
        let mut node = self;
        for k in key {
            node = match node.children().unwrap().get(k) {
                None => return None,
                Some(node) => match node {
                    Self::Tree(_) => node,
                    Self::Data(data) => return Some(data),
                },
            };
        }
        None
    }

    /// Similar to `prefix_get`, but takes and returns `&mut`.
    pub fn prefix_get_mut<'a, I, Q>(&mut self, key: &mut I) -> Option<&mut V>
    where
        K: 'a + Borrow<Q>,
        Q: 'a + Hash + Eq + ?Sized,
        I: Iterator<Item = &'a Q>,
    {
        if let Self::Data(data) = self {
            // return early so we don't consume from key_iter unnecessarily
            return Some(data);
        }
        let mut node = self;
        for k in key {
            node = match node.children_mut().unwrap().get_mut(k) {
                None => return None,
                Some(node) => match node {
                    Self::Tree(_) => node,
                    Self::Data(data) => return Some(data),
                },
            };
        }
        None
    }

    /// Get the subtree at `key`, if the entry exists and is a tree.
    /// Return an error if there is an entry but it is not a tree.
    pub fn get_subtree<'a, I, Q>(&self, key: &mut I) -> anyhow::Result<Option<&HashMap<K, Self>>>
    where
        K: 'a + Borrow<Q>,
        Q: 'a + Hash + Eq + ?Sized,
        I: Iterator<Item = &'a Q>,
    {
        let mut entries = match self {
            Self::Tree(ref t) => t,
            Self::Data(..) => {
                return Err(anyhow::anyhow!("Data found where tree expected"));
            }
        };

        for k in key {
            let node = match entries.get(k) {
                None => return Ok(None),
                Some(v) => v,
            };

            entries = match node {
                Self::Tree(ref t) => t,
                Self::Data(..) => {
                    return Err(anyhow::anyhow!("Data found where tree expected"));
                }
            };
        }

        Ok(Some(entries))
    }

    /// Inserts a key-value pair into the tree.
    ///
    /// If there is already a key in the map that is a prefix of the inserted
    /// key, that key is removed.
    ///
    /// If the inserted key is a prefix of one or more keys in the map, all
    /// those keys are removed.
    pub fn insert<I: Iterator<Item = K>>(&mut self, mut key: I, value: V) {
        if let Some(k) = key.next() {
            if matches!(self, Self::Data(_)) {
                *self = Self::new();
            }
            let child = match self.children_mut().unwrap().entry(k) {
                Entry::Occupied(e) => e.into_mut(),
                Entry::Vacant(e) => e.insert(Self::new()),
            };
            child.insert(key, value);
        } else {
            *self = Self::Data(value);
        }
    }

    /// Removes a key from the tree, returning the subtree at the key if the key
    /// was previously in the tree.
    /// If the prefix of `key` exists as a leaf on the tree, that leaf is removed
    /// and the value of that leaf is returned. We guarantee that only enough of
    /// `key` to find the returned value was consumed.
    pub fn remove<'a, I, Q>(&mut self, mut key: I) -> Option<DataTree<K, V>>
    where
        K: 'a + Borrow<Q>,
        Q: 'a + Hash + Eq + ?Sized,
        I: Iterator<Item = &'a Q>,
    {
        if matches!(self, Self::Data(_)) {
            return Some(std::mem::replace(self, Self::new()));
        }
        if let Some(k) = key.next() {
            if let Some(node) = self.children_mut().unwrap().get_mut(k) {
                let data = node.remove(key);
                let remove_node = match node {
                    Self::Tree(children) => children.is_empty(),
                    Self::Data(_) => true,
                };
                if remove_node {
                    self.children_mut().unwrap().remove(k);
                }
                data
            } else {
                None
            }
        } else {
            Some(std::mem::replace(self, Self::new()))
        }
    }

    pub fn children(&self) -> Option<&HashMap<K, DataTree<K, V>>> {
        match self {
            Self::Tree(children) => Some(children),
            Self::Data(_) => None,
        }
    }

    fn children_mut(&mut self) -> Option<&mut HashMap<K, DataTree<K, V>>> {
        match self {
            Self::Tree(children) => Some(children),
            Self::Data(_) => None,
        }
    }

    /// Returns an iterator over DataTree<K, V>.
    pub fn iter<T>(&self) -> DataTreeIterator<'_, K, V, T> {
        match self {
            Self::Tree(t) => DataTreeIterator::Stack(vec![(None, t.iter())], PhantomData),
            Self::Data(v) => DataTreeIterator::Entry(Some(v)),
        }
    }

    /// Take ownership of the values in DataTree<K, V> and iterate.
    pub fn into_iter<T>(self) -> DataTreeIntoIterator<K, V, T> {
        match self {
            Self::Tree(t) => DataTreeIntoIterator::Stack(vec![(None, t.into_iter())], PhantomData),
            Self::Data(v) => DataTreeIntoIterator::Entry(Some(v)),
        }
    }
}

struct NoopCollector;

impl<'a> FromIterator<&'a FileNameBuf> for NoopCollector {
    fn from_iter<I>(_iter: I) -> Self
    where
        I: IntoIterator<Item = &'a FileNameBuf>,
    {
        NoopCollector
    }
}

impl<V: 'static> FileTree<V> {
    pub fn iter_with_paths(&self) -> impl Iterator<Item = (ForwardRelativePathBuf, &V)> {
        self.iter::<Option<ForwardRelativePathBuf>>()
            .map(|(k, v)| (k.unwrap_or_else(ForwardRelativePathBuf::empty), v))
    }

    pub fn iter_without_paths(&self) -> impl Iterator<Item = &V> {
        self.iter::<NoopCollector>().map(|(NoopCollector, v)| v)
    }

    pub fn into_iter_with_paths(self) -> impl Iterator<Item = (ForwardRelativePathBuf, V)> {
        self.into_iter::<Option<ForwardRelativePathBuf>>()
            .map(|(k, v)| (k.unwrap_or_else(ForwardRelativePathBuf::empty), v))
    }

    #[allow(unused)]
    pub fn into_iter_without_paths(self) -> impl Iterator<Item = V> {
        self.into_iter::<NoopCollector>()
            .map(|(NoopCollector, v)| v)
    }
}

pub enum DataTreeIterator<'a, K, V, T> {
    Stack(
        Vec<(Option<&'a K>, Iter<'a, K, DataTree<K, V>>)>,
        PhantomData<T>,
    ),
    Entry(Option<&'a V>),
}

impl<'a, K, V, T> Iterator for DataTreeIterator<'a, K, V, T>
where
    T: for<'x> FromIterator<&'x K>,
{
    type Item = (T, &'a V);

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Stack(ref mut stack, _) => loop {
                let (_, ref mut last) = stack.last_mut()?;

                match last.next() {
                    Some((k, DataTree::Tree(t))) => {
                        stack.push((Some(k), t.iter()));
                    }
                    Some((k, DataTree::Data(v))) => {
                        let it = stack
                            .iter()
                            .filter_map(|(k, _)| k.as_deref())
                            .chain(std::iter::once(k));
                        return Some((it.collect(), v));
                    }
                    None => {
                        stack.pop();
                    }
                }
            },
            Self::Entry(v) => v.take().map(|v| (std::iter::empty().collect(), v)),
        }
    }
}

pub enum DataTreeIntoIterator<K, V, T> {
    Stack(
        Vec<(Option<K>, IntoIter<K, DataTree<K, V>>)>,
        PhantomData<T>,
    ),
    Entry(Option<V>),
}

impl<K, V, T> Iterator for DataTreeIntoIterator<K, V, T>
where
    T: for<'x> FromIterator<&'x K>,
{
    type Item = (T, V);

    fn next(&mut self) -> Option<Self::Item> {
        match self {
            Self::Stack(ref mut stack, _) => loop {
                let (_, ref mut last) = stack.last_mut()?;

                match last.next() {
                    Some((k, DataTree::Tree(t))) => {
                        stack.push((Some(k), t.into_iter()));
                    }
                    Some((k, DataTree::Data(v))) => {
                        let it = stack
                            .iter()
                            .filter_map(|(k, _)| k.as_ref())
                            .chain(std::iter::once(&k));
                        return Some((it.collect(), v));
                    }
                    None => {
                        stack.pop();
                    }
                }
            },
            Self::Entry(v) => v.take().map(|v| (std::iter::empty().collect(), v)),
        }
    }
}

#[cfg(test)]
mod tests {
    use std::collections::BTreeMap;

    use assert_matches::assert_matches;

    use super::*;

    #[derive(Debug)]
    struct CopyCollector<T>(Vec<T>);

    impl<'a, T> FromIterator<&'a T> for CopyCollector<T>
    where
        T: Copy + 'static,
    {
        fn from_iter<I>(iter: I) -> Self
        where
            I: IntoIterator<Item = &'a T>,
        {
            Self(iter.into_iter().copied().collect())
        }
    }

    #[test]
    fn test_iter() {
        let expected: BTreeMap<Vec<i32>, String> = [
            (vec![1, 2, 3], "123".to_owned()),
            (vec![1, 2, 4], "124".to_owned()),
            (vec![1, 2, 5, 6], "1256".to_owned()),
            (vec![1, 3, 4], "134".to_owned()),
            (vec![1, 3, 5], "135".to_owned()),
        ]
        .into_iter()
        .collect();

        let mut tree = DataTree::<i32, String>::new();
        for (k, v) in expected.iter() {
            tree.insert(k.clone().into_iter(), v.clone());
        }
        let actual = tree
            .iter::<CopyCollector<_>>()
            .map(|(k, v)| (k.0.into_iter().collect::<Vec<_>>(), v.to_owned()))
            .collect::<BTreeMap<_, _>>();

        assert_eq!(expected, actual);
    }

    #[test]
    fn test_into_iter() {
        let expected: BTreeMap<Vec<i32>, String> = [
            (vec![1, 2, 3], "123".to_owned()),
            (vec![1, 2, 4], "124".to_owned()),
            (vec![1, 2, 5, 6], "1256".to_owned()),
            (vec![1, 3, 4], "134".to_owned()),
            (vec![1, 3, 5], "135".to_owned()),
        ]
        .into_iter()
        .collect();

        let mut tree = DataTree::<i32, String>::new();
        for (k, v) in expected.iter() {
            tree.insert(k.clone().into_iter(), v.clone());
        }
        let actual = tree
            .into_iter::<CopyCollector<_>>()
            .map(|(k, v)| (k.0.into_iter().collect::<Vec<_>>(), v))
            .collect::<BTreeMap<_, _>>();

        assert_eq!(expected, actual);
    }

    #[test]
    fn test_get() {
        let path = ForwardRelativePathBuf::unchecked_new("foo/bar".to_owned());
        let contents = "contents_of_foobar".to_owned();

        let mut file_path_to_contents: DataTree<FileNameBuf, String> = DataTree::new();
        file_path_to_contents.insert(path.iter().map(|f| f.to_owned()), contents.clone());

        assert_eq!(
            file_path_to_contents
                .prefix_get_mut(&mut path.iter())
                .as_deref(),
            Some(&contents)
        );
    }

    #[test]
    fn test_prefix_remove() {
        // We need remove any artifacts at leaves when we invalidate
        // and declare artifacts at a subdirectory of that leaf path.
        let mut tree = DataTree::<i32, String>::new();
        tree.insert(vec![1, 2, 3].into_iter(), "123".to_owned());
        let key = vec![1, 2, 3, 4];
        let mut key_iter = key.iter();
        let removed = tree.remove(&mut key_iter);
        assert_matches!(removed, Some(DataTree::Data(val)) if val == *"123");
        // Test that only enough of key_iter to find the returned value was consumed.
        assert_eq!(key_iter.next(), Some(&4));
        assert_eq!(key_iter.next(), None);

        // Check tree is empty
        assert_matches!(tree.iter::<CopyCollector<_>>().next(), None);
    }

    #[test]
    fn test_suffix_remove() {
        let mut tree = DataTree::<i32, String>::new();
        tree.insert(vec![1, 2, 3].into_iter(), "123".to_owned());
        tree.remove(vec![1, 2].iter());

        assert_matches!(tree.iter::<CopyCollector<_>>().next(), None);
    }

    #[test]
    fn test_remove_empty_dirs() {
        let mut tree = DataTree::<i32, String>::new();
        tree.insert(vec![1, 2, 3].into_iter(), "123".to_owned());
        tree.remove(vec![1, 2, 3].iter());
        assert_matches!(tree, DataTree::Tree(m) if m.is_empty());
    }
}


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

use std::collections::HashMap;
use std::collections::HashSet;
use std::str::FromStr;
use std::sync::Arc;

use anyhow::Context;
use async_trait::async_trait;
use buck2_common::file_ops::FileDigest;
use buck2_common::http::counting_client::CountingHttpClient;
use buck2_common::result::SharedError;
use buck2_common::result::ToSharedResultExt;
use buck2_core::directory::unordered_entry_walk;
use buck2_core::directory::DirectoryEntry;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_events::dispatch::EventDispatcher;
use buck2_execute::digest::CasDigestFromReExt;
use buck2_execute::digest::CasDigestToReExt;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::ActionDirectoryEntry;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::directory::ActionSharedDirectory;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::execute::blocking::IoRequest;
use buck2_execute::execute::clean_output_paths::cleanup_path;
use buck2_execute::materialize::http::http_download;
use buck2_execute::output_size::OutputSize;
use buck2_execute::re::manager::ReConnectionManager;
use chrono::Duration;
use chrono::Utc;
use dupe::Dupe;
use futures::future::BoxFuture;
use futures::future::Future;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use gazebo::prelude::VecExt;
use more_futures::cancellation::CancellationContext;
use once_cell::sync::Lazy;
use remote_execution::NamedDigest;
use remote_execution::NamedDigestWithPermissions;
use remote_execution::REClientError;
use remote_execution::TCode;
use remote_execution::TDigest;
use tracing::instrument;

use crate::materializers::deferred::ArtifactMaterializationMethod;
use crate::materializers::deferred::ArtifactMaterializationStage;
use crate::materializers::deferred::ArtifactTree;
use crate::materializers::deferred::LowPriorityMaterializerCommand;
use crate::materializers::deferred::MaterializationMethodToProto;
use crate::materializers::deferred::MaterializeEntryError;
use crate::materializers::deferred::MaterializerSender;
use crate::materializers::deferred::SharedMaterializingError;
use crate::materializers::deferred::Version;
use crate::materializers::deferred::WriteFile;
use crate::materializers::io::materialize_files;
use crate::materializers::io::MaterializeTreeStructure;

pub(super) struct DefaultIoHandler {
    pub(super) fs: ProjectRoot,
    pub(super) digest_config: DigestConfig,
    pub(super) buck_out_path: ProjectRelativePathBuf,
    pub(super) re_client_manager: Arc<ReConnectionManager>,
    /// Executor for blocking IO operations
    pub(super) io_executor: Arc<dyn BlockingExecutor>,
    pub(super) http_client: CountingHttpClient,
}

struct MaterializationStat {
    file_count: u64,
    total_bytes: u64,
}

#[async_trait]
pub(super) trait IoHandler: Sized + Sync + Send + 'static {
    fn write<'a>(
        self: &Arc<Self>,
        path: ProjectRelativePathBuf,
        write: Arc<WriteFile>,
        version: Version,
        command_sender: MaterializerSender<Self>,
        cancellations: &'a CancellationContext<'a>,
    ) -> BoxFuture<'a, Result<(), SharedMaterializingError>>;

    fn clean_path<'a>(
        self: &Arc<Self>,
        path: ProjectRelativePathBuf,
        version: Version,
        command_sender: MaterializerSender<Self>,
        cancellations: &'a CancellationContext,
    ) -> BoxFuture<'a, Result<(), SharedError>>;

    async fn materialize_entry(
        self: &Arc<Self>,
        path: ProjectRelativePathBuf,
        method: Arc<ArtifactMaterializationMethod>,
        entry: ActionDirectoryEntry<ActionSharedDirectory>,
        event_dispatcher: EventDispatcher,
        cancellations: &CancellationContext,
    ) -> Result<(), MaterializeEntryError>;

    fn create_ttl_refresh(
        self: &Arc<Self>,
        tree: &ArtifactTree,
        min_ttl: Duration,
        digest_config: DigestConfig,
    ) -> Option<BoxFuture<'static, anyhow::Result<()>>>;
}

impl DefaultIoHandler {
    /// Materializes an `entry` at `path`, using the materialization `method`
    #[instrument(level = "debug", skip(self, stat, cancellations), fields(path = %path, method = %method, entry = %entry))]
    async fn materialize_entry_span(
        &self,
        path: ProjectRelativePathBuf,
        method: Arc<ArtifactMaterializationMethod>,
        entry: ActionDirectoryEntry<ActionSharedDirectory>,
        stat: &mut MaterializationStat,
        cancellations: &CancellationContext<'_>,
    ) -> Result<(), MaterializeEntryError> {
        // Materialize the dir structure, and symlinks
        self.io_executor
            .execute_io(
                Box::new(MaterializeTreeStructure {
                    path: path.clone(),
                    entry: entry.dupe(),
                }),
                cancellations,
            )
            .await?;

        // Materialize files
        match method.as_ref() {
            ArtifactMaterializationMethod::CasDownload { info } => {
                let mut files = Vec::new();

                {
                    let mut walk = unordered_entry_walk(entry.as_ref());

                    while let Some((entry_path, entry)) = walk.next() {
                        if let DirectoryEntry::Leaf(ActionDirectoryMember::File(f)) = entry {
                            let name = path.join_normalized(entry_path.get())?;
                            let digest = maybe_tombstone_digest(f.digest.data())?.to_re();

                            tracing::trace!(name = %name, digest = %digest, "push download");
                            let name = self
                                .fs
                                .resolve(&name)
                                .as_maybe_relativized_str()?
                                .to_owned();

                            files.push(NamedDigestWithPermissions {
                                named_digest: NamedDigest {
                                    name,
                                    digest,
                                    ..Default::default()
                                },
                                is_executable: f.is_executable,
                                ..Default::default()
                            });
                        }
                    }
                }
                stat.file_count = files.len().try_into().unwrap_or_default();
                stat.total_bytes = files
                    .iter()
                    .map(|x| u64::try_from(x.named_digest.digest.size_in_bytes).unwrap_or_default())
                    .sum();

                let connection = self.re_client_manager.get_re_connection();
                let re_client = connection.get_client();

                re_client
                    .materialize_files(files, info.re_use_case)
                    .await
                    .map_err(|e| match e.downcast_ref::<REClientError>() {
                        Some(e) if e.code == TCode::NOT_FOUND => MaterializeEntryError::NotFound {
                            info: info.dupe(),
                            debug: Arc::from(e.message.as_str()),
                        },
                        _ => MaterializeEntryError::Error(e.context({
                            format!(
                                "Error materializing files declared by action: {}",
                                info.origin
                            )
                        })),
                    })?;
            }
            ArtifactMaterializationMethod::HttpDownload { info } => {
                async {
                    let downloaded = http_download(
                        &self.http_client,
                        &self.fs,
                        self.digest_config,
                        &path,
                        &info.url,
                        &info.checksum,
                        info.metadata.is_executable,
                    )
                    .await?;

                    // Check that the size we got was the one that we expected. This isn't strictly
                    // speaking necessary here, but since an invalid size would break actions
                    // running on RE, it's a good idea to catch it here when materializing so that
                    // our test suite can surface bugs when downloading things locally.
                    if downloaded.size() != info.metadata.digest.size() {
                        return Err(anyhow::anyhow!(
                            "Downloaded size ({}) does not match expected size ({})",
                            downloaded.size(),
                            info.metadata.digest.size(),
                        ));
                    }
                    stat.file_count = 1;
                    stat.total_bytes = info.metadata.digest.size();
                    Ok(())
                }
                .boxed()
                .await
                .with_context(|| {
                    format!(
                        "Error materializing HTTP resource declared by target `{}`",
                        info.owner
                    )
                })?;
            }
            ArtifactMaterializationMethod::LocalCopy(_, copied_artifacts) => {
                self.io_executor
                    .execute_io_inline(|| {
                        for a in copied_artifacts {
                            let count_and_bytes = a.dest_entry.calc_output_count_and_bytes();
                            stat.file_count += count_and_bytes.count;
                            stat.total_bytes += count_and_bytes.bytes;

                            materialize_files(
                                a.dest_entry.as_ref(),
                                &self.fs.root().join(&a.src),
                                &self.fs.root().join(&a.dest),
                            )?;
                        }
                        Ok(())
                    })
                    .await?;
            }
            ArtifactMaterializationMethod::Write(write) => {
                stat.file_count = 1;
                self.io_executor
                    .execute_io_inline(|| {
                        let data =
                            zstd::bulk::decompress(&write.compressed_data, write.decompressed_size)
                                .context("Error decompressing data")?;
                        stat.total_bytes = write.decompressed_size as u64;
                        self.fs.write_file(&path, data, write.is_executable)
                    })
                    .await?;
            }
            #[cfg(test)]
            ArtifactMaterializationMethod::Test => unimplemented!(),
        };
        Ok(())
    }
}

#[async_trait]
impl IoHandler for DefaultIoHandler {
    fn write<'a>(
        self: &Arc<Self>,
        path: ProjectRelativePathBuf,
        write: Arc<WriteFile>,
        version: Version,
        command_sender: MaterializerSender<Self>,
        cancellations: &'a CancellationContext,
    ) -> BoxFuture<'a, Result<(), SharedMaterializingError>> {
        self.io_executor
            .execute_io(
                Box::new(WriteIoRequest {
                    path,
                    write,
                    version,
                    command_sender,
                }),
                cancellations,
            )
            .map_err(|e| SharedMaterializingError::Error(e.into()))
            .boxed()
    }

    fn clean_path<'a>(
        self: &Arc<Self>,
        path: ProjectRelativePathBuf,
        version: Version,
        command_sender: MaterializerSender<Self>,
        cancellations: &'a CancellationContext,
    ) -> BoxFuture<'a, Result<(), SharedError>> {
        self.io_executor
            .execute_io(
                Box::new(CleanIoRequest {
                    path,
                    version,
                    command_sender,
                }),
                cancellations,
            )
            .map(|r| r.shared_error())
            .boxed()
    }

    /// Materializes an `entry` at `path`, using the materialization `method`
    #[instrument(level = "debug", skip(self, cancellations), fields(path = %path, method = %method, entry = %entry))]
    async fn materialize_entry(
        self: &Arc<Self>,
        path: ProjectRelativePathBuf,
        method: Arc<ArtifactMaterializationMethod>,
        entry: ActionDirectoryEntry<ActionSharedDirectory>,
        event_dispatcher: EventDispatcher,
        cancellations: &CancellationContext,
    ) -> Result<(), MaterializeEntryError> {
        let materialization_start = buck2_data::MaterializationStart {
            action_digest: match method.as_ref() {
                ArtifactMaterializationMethod::CasDownload { info } => {
                    info.action_digest().map(|digest| digest.to_string())
                }
                _ => None,
            },
        };
        event_dispatcher
            .span_async(materialization_start, async move {
                let path_string = path.as_str().to_owned();
                let mut stat = MaterializationStat {
                    file_count: 0,
                    total_bytes: 0,
                };
                let res = self
                    .materialize_entry_span(path, method.dupe(), entry, &mut stat, cancellations)
                    .await;
                let error = res.as_ref().err().map(|e| format!("{:#}", e));

                (
                    res,
                    buck2_data::MaterializationEnd {
                        action_digest: None,
                        file_count: stat.file_count,
                        total_bytes: stat.total_bytes,
                        path: path_string,
                        success: error.is_none(),
                        error,
                        method: Some(method.to_proto() as i32),
                    },
                )
            })
            .await?;
        Ok(())
    }

    fn create_ttl_refresh(
        self: &Arc<Self>,
        tree: &ArtifactTree,
        min_ttl: Duration,
        digest_config: DigestConfig,
    ) -> Option<BoxFuture<'static, anyhow::Result<()>>> {
        create_ttl_refresh(tree, &self.re_client_manager, min_ttl, digest_config).map(|f| f.boxed())
    }
}

/// This is used for testing to ingest digests (via BUCK2_TEST_TOMBSTONED_DIGESTS).
fn maybe_tombstone_digest(digest: &FileDigest) -> anyhow::Result<&FileDigest> {
    // This has to be of size 1 since size 0 will result in the RE client just producing an empty
    // instead of a not-found error.
    static TOMBSTONE_DIGEST: Lazy<FileDigest> = Lazy::new(|| FileDigest::new_sha1([0; 20], 1));

    fn convert_digests(val: &str) -> anyhow::Result<HashSet<FileDigest>> {
        val.split(' ')
            .map(|digest| {
                let digest = TDigest::from_str(digest)
                    .with_context(|| format!("Invalid digest: `{}`", digest))?;
                // This code is only used by E2E tests, so while it's not *a test*, testing_default
                // is an OK choice here.
                let digest = FileDigest::from_re(&digest, DigestConfig::testing_default())?;
                anyhow::Ok(digest)
            })
            .collect()
    }

    static TOMBSTONED_DIGESTS: EnvHelper<HashSet<FileDigest>> =
        EnvHelper::with_converter("BUCK2_TEST_TOMBSTONED_DIGESTS", convert_digests);

    if let Some(digests) = TOMBSTONED_DIGESTS.get()? {
        if digests.contains(digest) {
            return Ok(&*TOMBSTONE_DIGEST);
        }
    }

    Ok(digest)
}

/// Spawn a task to refresh TTLs.
pub(super) fn create_ttl_refresh(
    tree: &ArtifactTree,
    re_manager: &Arc<ReConnectionManager>,
    min_ttl: Duration,
    digest_config: DigestConfig,
) -> Option<impl Future<Output = anyhow::Result<()>>> {
    let mut digests_to_refresh = HashMap::<_, HashSet<_>>::new();

    let ttl_deadline = Utc::now() + min_ttl;

    for data in tree.iter_without_paths() {
        match &data.stage {
            ArtifactMaterializationStage::Declared { entry, method } => match method.as_ref() {
                ArtifactMaterializationMethod::CasDownload { info } => {
                    let mut walk = unordered_entry_walk(entry.as_ref());
                    while let Some((_entry_path, entry)) = walk.next() {
                        if let DirectoryEntry::Leaf(ActionDirectoryMember::File(file)) = entry {
                            let needs_refresh = file.digest.expires() < ttl_deadline;
                            tracing::trace!("{} needs_refresh: {}", file, needs_refresh);
                            if needs_refresh {
                                digests_to_refresh
                                    .entry(info.re_use_case)
                                    .or_default()
                                    .insert(file.digest.dupe());
                            }
                        }
                    }
                }
                _ => {}
            },
            _ => {}
        }
    }

    if digests_to_refresh.is_empty() {
        return None;
    }

    let re_manager = re_manager.dupe();

    let fut = async move {
        // We need to pick *a number* to not send an unbounded amount of digests here. 500 seems
        // broadly reasonable.
        const REFRESH_CHUNK_SIZE: usize = 500;

        let re_connection = re_manager.get_re_connection();
        let re_client = re_connection.get_client();

        for (use_case, digests_to_refresh) in digests_to_refresh {
            let mut digests_to_refresh = digests_to_refresh.into_iter().collect::<Vec<_>>();
            digests_to_refresh.sort();

            for chunk in digests_to_refresh.as_slice().chunks(REFRESH_CHUNK_SIZE) {
                tracing::debug!("Update {} TTLs", chunk.len());

                let digests_expires = re_client
                    .get_digest_expirations(chunk.iter().map(|d| d.to_re()).collect(), use_case)
                    .await?;

                let mut digests_expires = digests_expires.into_try_map(|(digest, expires)| {
                    anyhow::Ok((FileDigest::from_re(&digest, digest_config)?, expires))
                })?;
                digests_expires.sort();

                if chunk.len() != digests_expires.len() {
                    return Err(anyhow::anyhow!(
                        "Invalid response from get_digests_ttl: expected {}, got {} digests",
                        chunk.len(),
                        digests_expires.len()
                    ));
                }

                for (digest, (matching_digest, expires)) in chunk.iter().zip(&digests_expires) {
                    if digest.data() != matching_digest {
                        return Err(anyhow::anyhow!("Invalid response from get_digests_ttl"));
                    }

                    digest.update_expires(*expires);
                }
            }
        }

        anyhow::Ok(())
    }
    .map(|res| {
        if let Err(e) = &res {
            tracing::info!("TTL Refresh failed: {:#}", e);
        }

        res
    });

    Some(fut)
}

struct WriteIoRequest {
    path: ProjectRelativePathBuf,
    write: Arc<WriteFile>,
    version: Version,
    command_sender: MaterializerSender<DefaultIoHandler>,
}

impl WriteIoRequest {
    fn execute_inner(&self, project_fs: &ProjectRoot) -> anyhow::Result<()> {
        cleanup_path(project_fs, &self.path)?;
        let data =
            zstd::bulk::decompress(&self.write.compressed_data, self.write.decompressed_size)
                .context("Error decompressing data")?;
        project_fs.write_file(&self.path, data, self.write.is_executable)?;
        Ok(())
    }
}

impl IoRequest for WriteIoRequest {
    fn execute(self: Box<Self>, project_fs: &ProjectRoot) -> anyhow::Result<()> {
        // NOTE: No spans here! We should perhaps add one, but this needs to be considered
        // carefully as it's a lot of spans, and we haven't historically emitted those for writes.
        let res = self.execute_inner(project_fs).shared_error();

        // If the materializer has shut down, we ignore this.
        let _ignored = self.command_sender.send_low_priority(
            LowPriorityMaterializerCommand::MaterializationFinished {
                path: self.path,
                timestamp: Utc::now(),
                version: self.version,
                result: res.dupe().map_err(SharedMaterializingError::Error),
            },
        );

        Ok(res?)
    }
}

struct CleanIoRequest {
    path: ProjectRelativePathBuf,
    version: Version,
    command_sender: MaterializerSender<DefaultIoHandler>,
}

impl IoRequest for CleanIoRequest {
    fn execute(self: Box<Self>, project_fs: &ProjectRoot) -> anyhow::Result<()> {
        // NOTE: No spans here! We should perhaps add one, but this needs to be considered
        // carefully as it's a lot of spans, and we haven't historically emitted those for writes.
        let res = cleanup_path(project_fs, &self.path).shared_error();

        // If the materializer has shut down, we ignore this.
        let _ignored = self.command_sender.send_low_priority(
            LowPriorityMaterializerCommand::CleanupFinished {
                path: self.path,
                version: self.version,
                result: res.dupe().map_err(SharedMaterializingError::Error),
            },
        );

        Ok(res?)
    }
}