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

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

use async_trait::async_trait;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_core::fs::artifact_path_resolver::ArtifactFs;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::soft_error;
use buck2_events::dispatch::span_async;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::execute::action_digest::ActionDigest;
use buck2_execute::execute::blobs::ActionBlobs;
use buck2_execute::execute::kind::CommandExecutionKind;
use buck2_execute::execute::manager::CommandExecutionManager;
use buck2_execute::execute::manager::CommandExecutionManagerExt;
use buck2_execute::execute::output::CommandStdStreams;
use buck2_execute::execute::prepared::PreparedAction;
use buck2_execute::execute::prepared::PreparedCommand;
use buck2_execute::execute::prepared::PreparedCommandExecutor;
use buck2_execute::execute::request::CommandExecutionPaths;
use buck2_execute::execute::request::CommandExecutionRequest;
use buck2_execute::execute::request::ExecutorPreference;
use buck2_execute::execute::result::CommandExecutionResult;
use buck2_execute::execute::target::CommandExecutionTarget;
use buck2_execute::knobs::ExecutorGlobalKnobs;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::re::action_identity::ReActionIdentity;
use buck2_execute::re::client::ExecuteResponseOrCancelled;
use buck2_execute::re::manager::ManagedRemoteExecutionClient;
use buck2_execute::re::remote_action_result::RemoteActionResult;
use dupe::Dupe;
use futures::FutureExt;
use indexmap::IndexMap;
use more_futures::cancellation::CancellationContext;
use remote_execution as RE;
use remote_execution::ExecuteResponse;
use remote_execution::TCode;
use thiserror::Error;
use tracing::info;

use crate::re::download::download_action_results;
use crate::re::download::DownloadResult;

#[derive(Debug, Error)]
pub enum RemoteExecutorError {
    #[error("Trying to execute a `local_only = True` action on remote executor")]
    LocalOnlyAction,
}

pub struct ReExecutor {
    pub artifact_fs: ArtifactFs,
    pub project_fs: ProjectRoot,
    pub materializer: Arc<dyn Materializer>,
    pub re_client: ManagedRemoteExecutionClient,
    pub re_use_case: RemoteExecutorUseCase,
    pub re_action_key: Option<String>,
    pub re_max_input_files_bytes: u64,
    pub knobs: ExecutorGlobalKnobs,
    pub skip_cache_read: bool,
    pub skip_cache_write: bool,
    pub re_max_queue_time_ms: Option<u64>,
}

impl ReExecutor {
    /// Indicate whether an action is too big to run on RE.
    pub fn is_action_too_large(&self, paths: &CommandExecutionPaths) -> bool {
        paths.input_files_bytes() > self.re_max_input_files_bytes
    }

    async fn upload(
        &self,
        manager: CommandExecutionManager,
        blobs: &ActionBlobs,
        paths: &CommandExecutionPaths,
        digest_config: DigestConfig,
    ) -> ControlFlow<CommandExecutionResult, CommandExecutionManager> {
        let re_client = &self.re_client;

        let upload_response = span_async(buck2_data::ReUploadStart {}, async move {
            let res = re_client
                .upload(
                    &self.project_fs,
                    &self.materializer,
                    blobs,
                    ProjectRelativePath::empty(),
                    paths.input_directory(),
                    self.re_use_case,
                    digest_config,
                )
                .await;
            match res {
                Ok(stats) => (
                    Ok(()),
                    buck2_data::ReUploadEnd {
                        digests_uploaded: Some(stats.digests_uploaded),
                        bytes_uploaded: Some(stats.bytes_uploaded),
                    },
                ),
                Err(e) => (Err(e), buck2_data::ReUploadEnd::default()),
            }
        })
        .await;

        match upload_response {
            Ok(()) => {}
            Err(e) => return ControlFlow::Break(manager.error("remote_upload_error", e)),
        };

        ControlFlow::Continue(manager)
    }

    async fn re_execute(
        &self,
        mut manager: CommandExecutionManager,
        action: &dyn CommandExecutionTarget,
        request: &CommandExecutionRequest,
        action_digest: &ActionDigest,
        digest_config: DigestConfig,
        platform: &RE::Platform,
    ) -> ControlFlow<CommandExecutionResult, (CommandExecutionManager, ExecuteResponse)> {
        info!(
            "RE command line:\n```\n$ {}\n```\n for action `{}`",
            request.all_args_str(),
            action_digest,
        );

        let identity =
            ReActionIdentity::new(action, self.re_action_key.as_deref(), request.paths());

        let execute_response = self
            .re_client
            .execute(
                action_digest.dupe(),
                platform,
                self.re_use_case,
                &identity,
                &mut manager,
                self.skip_cache_read,
                self.skip_cache_write,
                self.re_max_queue_time_ms.map(Duration::from_millis),
                &self.knobs,
            )
            .await;

        let response = match execute_response {
            Ok(ExecuteResponseOrCancelled::Response(result)) => result,
            Ok(ExecuteResponseOrCancelled::Cancelled) => {
                return ControlFlow::Break(manager.cancel());
            }
            Err(e) => return ControlFlow::Break(manager.error("remote_call_error", e)),
        };

        let action_result = &response.action_result;

        if response.error.code != TCode::OK {
            let res = if let Some(out) = as_missing_outputs_error(&response.error) {
                // TODO: Add a dedicated report variant for this.
                // NOTE: We don't get stdout / stderr from RE when this happens, so the best we can
                // do here is just pass on the error.
                manager.failure(
                    CommandExecutionKind::Remote {
                        digest: action_digest.dupe(),
                    },
                    IndexMap::new(),
                    CommandStdStreams::Local {
                        stdout: Vec::new(),
                        stderr: out.to_owned().into(),
                    },
                    // We also don't get this output so don't put trash in here.
                    None,
                    Default::default(),
                )
            } else if is_timeout_error(&response.error) && request.timeout().is_some() {
                manager.timeout(
                    CommandExecutionKind::Remote {
                        digest: action_digest.dupe(),
                    },
                    // Checked above: we fallthrough to the error path if we didn't set a timeout
                    // and yet received one.
                    request.timeout().unwrap(),
                    CommandStdStreams::Remote(response.std_streams(
                        &self.re_client,
                        self.re_use_case,
                        digest_config,
                    )),
                    response.timing(),
                )
            } else {
                manager.error(
                    "remote_exec_error",
                    ReErrorWrapper {
                        action_digest: action_digest.dupe(),
                        inner: response.error,
                    },
                )
            };

            return ControlFlow::Break(res);
        }

        if let Some(timeout) = request.timeout() {
            let execution_time = response.timing().execution_time;

            if execution_time > timeout {
                let res = soft_error!(
                    "re_timeout_exceeded",
                    anyhow::anyhow!(
                        "Command {} exceeded its timeout (ran for {}s, timeout was {}s)",
                        action.re_action_key(),
                        execution_time.as_secs(),
                        timeout.as_secs(),
                    )
                );

                if let Err(e) = res {
                    return ControlFlow::Break(manager.error("re_timeout_exceeded", e));
                }
            }
        }

        if action_result.exit_code != 0 {
            return ControlFlow::Break(manager.failure(
                CommandExecutionKind::Remote {
                    digest: action_digest.dupe(),
                },
                // TODO: we want to expose RE outputs even when actions fail,
                //   this will allow tpx to correctly retrieve the output of
                //   failing tests running on RE. See D34344489 for context.
                IndexMap::new(),
                CommandStdStreams::Remote(response.std_streams(
                    &self.re_client,
                    self.re_use_case,
                    digest_config,
                )),
                Some(action_result.exit_code),
                response.timing(),
            ));
        }

        ControlFlow::Continue((manager, response))
    }
}

#[async_trait]
impl PreparedCommandExecutor for ReExecutor {
    async fn exec_cmd(
        &self,
        command: &PreparedCommand<'_, '_>,
        manager: CommandExecutionManager,
        cancellations: &CancellationContext,
    ) -> CommandExecutionResult {
        let PreparedCommand {
            request,
            target,
            prepared_action:
                PreparedAction {
                    action: action_digest,
                    blobs,
                    platform,
                },
            digest_config,
        } = command;

        if command.request.executor_preference().requires_local() {
            return ControlFlow::Break(
                manager.error("remote_prepare", RemoteExecutorError::LocalOnlyAction),
            )?;
        }

        // TODO(bobyf, torozco): remote execution probably needs to explicitly handle cancellations
        let manager = self
            .upload(manager, blobs, request.paths(), *digest_config)
            .await?;

        let (manager, response) = self
            .re_execute(
                manager,
                *target,
                request,
                action_digest,
                *digest_config,
                platform,
            )
            .await?;

        let res = download_action_results(
            request,
            &*self.materializer,
            &self.re_client,
            self.re_use_case,
            *digest_config,
            manager,
            buck2_data::ReStage {
                stage: Some(buck2_data::ReDownload {}.into()),
            }
            .into(),
            request.paths(),
            request.outputs(),
            action_digest,
            &response,
            cancellations,
        )
        .boxed()
        .await;

        let DownloadResult::Result(res) = res;

        res
    }

    fn is_local_execution_possible(&self, _executor_preference: ExecutorPreference) -> bool {
        false
    }
}

#[derive(Error, Debug)]
#[error(
    "action_digest={}, re_code={}, re_location={}, re_message={}",
    .action_digest,
    .inner.code,
    .inner.error_location,
    .inner.message
)]
pub struct ReErrorWrapper {
    action_digest: ActionDigest,
    inner: remote_execution::REError,
}

fn as_missing_outputs_error(err: &remote_execution::REError) -> Option<&str> {
    // A dedicated error code would be better for this :(
    if err.message.contains("OUTMISS") {
        Some(&err.message)
    } else {
        None
    }
}

fn is_timeout_error(err: &remote_execution::REError) -> bool {
    #[cfg(fbcode_build)]
    {
        // Not ideal, but DEADLINE_EXCEEDED will show up if you e.g. timeout connecting to RE, so we
        // need to actually match on the message :(
        err.code == TCode::DEADLINE_EXCEEDED && err.message.contains("Execution timed out")
    }

    #[cfg(not(fbcode_build))]
    {
        // Not obvious what this looks like in the GRPC API.
        let _ignored = err;
        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::HashMap;
use std::ffi::OsString;
use std::sync::Arc;
use std::time::Duration;

use buck2_common::client_utils::get_channel_uds;
use buck2_common::client_utils::retrying;
use buck2_common::liveliness_observer::LivelinessGuard;
use buck2_common::liveliness_observer::LivelinessObserver;
use buck2_common::result::SharedError;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::file_name::FileName;
use buck2_events::dispatch::EventDispatcher;
use buck2_execute::execute::kind::CommandExecutionKind;
use buck2_execute::execute::manager::CommandExecutionManagerExt;
use buck2_execute::execute::manager::CommandExecutionManagerWithClaim;
use buck2_execute::execute::output::CommandStdStreams;
use buck2_execute::execute::request::CommandExecutionRequest;
use buck2_execute::execute::request::WorkerId;
use buck2_execute::execute::request::WorkerSpec;
use buck2_execute::execute::result::CommandExecutionMetadata;
use buck2_execute::execute::result::CommandExecutionResult;
use buck2_forkserver::client::ForkserverClient;
use buck2_forkserver::run::GatherOutputStatus;
use buck2_worker_proto::execute_command::EnvironmentEntry;
use buck2_worker_proto::worker_client::WorkerClient;
use buck2_worker_proto::ExecuteCommand;
use buck2_worker_proto::ExecuteResponse;
use futures::future::BoxFuture;
use futures::future::Shared;
use futures::FutureExt;
use indexmap::IndexMap;
use tokio::task::JoinHandle;
use tonic::transport::Channel;

#[derive(thiserror::Error, Debug)]
pub enum WorkerInitError {
    #[error("Worker failed to spawn: {0}")]
    SpawnFailed(String),
    #[error("Worker exited before connecting")]
    EarlyExit {
        exit_code: Option<i32>,
        stdout: String,
        stderr: String,
    },
    #[error("Worker failed to connect within `{0:.2}` seconds: {1}")]
    ConnectionTimeout(f64, String),
    /// Any error not related to worker behavior
    #[error("Error initializing worker `{0}`")]
    InternalError(SharedError),
}

impl WorkerInitError {
    pub(crate) fn to_command_execution_result(
        &self,
        request: &CommandExecutionRequest,
        manager: CommandExecutionManagerWithClaim,
    ) -> CommandExecutionResult {
        let worker_spec = request.worker().as_ref().unwrap();
        let execution_kind = CommandExecutionKind::LocalWorkerInit {
            command: worker_spec.exe.clone(),
            env: request.env().clone(),
        };

        match self {
            WorkerInitError::EarlyExit {
                exit_code,
                stdout,
                stderr,
            } => {
                let std_streams = CommandStdStreams::Local {
                    stdout: stdout.to_owned().into(),
                    stderr: stderr.to_owned().into(),
                };
                // TODO(ctolliday) this should be a new failure type (worker_init_failure), not conflated with a "command failure" which
                // implies that it is the primary command and that exit code != 0
                manager.failure(
                    execution_kind,
                    IndexMap::default(),
                    std_streams,
                    *exit_code,
                    CommandExecutionMetadata::default(),
                )
            }
            // TODO(ctolliday) as above, use a new failure type (worker_init_failure) that indicates this is a worker initialization error.
            WorkerInitError::ConnectionTimeout(..) | WorkerInitError::SpawnFailed(..) => manager
                .failure(
                    execution_kind,
                    IndexMap::default(),
                    CommandStdStreams::Local {
                        stdout: Default::default(),
                        stderr: format!("Error initializing worker: {}", self).into_bytes(),
                    },
                    None,
                    CommandExecutionMetadata::default(),
                ),
            WorkerInitError::InternalError(error) => {
                manager.error("get_worker_failed", error.clone())
            }
        }
    }
}

#[cfg(unix)]
fn spawn_via_forkserver(
    forkserver: ForkserverClient,
    exe: OsString,
    args: Vec<OsString>,
    env: Vec<(OsString, OsString)>,
    working_directory: AbsNormPathBuf,
    liveliness_observer: impl LivelinessObserver + 'static,
    stdout_path: &AbsNormPathBuf,
    stderr_path: &AbsNormPathBuf,
    socket_path: &AbsNormPathBuf,
) -> JoinHandle<anyhow::Result<GatherOutputStatus>> {
    use std::os::unix::ffi::OsStrExt;

    use crate::executors::local::apply_local_execution_environment;

    let stdout_path = stdout_path.clone();
    let stderr_path = stderr_path.clone();

    let socket_path = socket_path.clone();
    tokio::spawn(async move {
        let mut req = buck2_forkserver_proto::CommandRequest {
            exe: exe.as_bytes().into(),
            argv: args.into_iter().map(|s| s.as_bytes().into()).collect(),
            cwd: Some(buck2_forkserver_proto::WorkingDirectory {
                path: working_directory.as_path().as_os_str().as_bytes().into(),
            }),
            env: vec![],
            timeout: None,
            enable_miniperf: false,
            std_redirects: Some(buck2_forkserver_proto::command_request::StdRedirectPaths {
                stdout: stdout_path.as_os_str().as_bytes().into(),
                stderr: stderr_path.as_os_str().as_bytes().into(),
            }),
        };
        apply_local_execution_environment(&mut req, &working_directory, env, None);
        let res = forkserver
            .execute(req, async move { liveliness_observer.while_alive().await })
            .await
            .map(|(status, _, _)| status);

        // Socket is created by worker so won't exist if initialization fails.
        if fs_util::try_exists(&socket_path)? {
            // TODO(ctolliday) delete directory (after logs are moved to buck-out)
            fs_util::remove_file(&socket_path)?;
        }
        res
    })
}

#[cfg(not(unix))]
fn spawn_via_forkserver(
    _forkserver: ForkserverClient,
    _exe: OsString,
    _args: Vec<OsString>,
    _env: Vec<(OsString, OsString)>,
    _working_directory: AbsNormPathBuf,
    _liveliness_observer: impl LivelinessObserver + 'static,
    _stdout_path: &AbsNormPathBuf,
    _stderr_path: &AbsNormPathBuf,
    _socket_path: &AbsNormPathBuf,
) -> JoinHandle<anyhow::Result<GatherOutputStatus>> {
    unreachable!("workers should not be initialized off unix")
}

async fn spawn_worker(
    worker_spec: &WorkerSpec,
    env: impl IntoIterator<Item = (OsString, OsString)>,
    root: &AbsNormPathBuf,
    forkserver: ForkserverClient,
    dispatcher: EventDispatcher,
) -> Result<WorkerHandle, WorkerInitError> {
    // Use fixed length path at /tmp to avoid 108 character limit for unix domain sockets
    let dir_name = format!("{}-{}", dispatcher.trace_id(), worker_spec.id);
    let worker_dir = AbsNormPathBuf::from("/tmp/buck2_worker".to_owned())
        .map_err(|e| WorkerInitError::InternalError(e.into()))?
        .join(FileName::unchecked_new(&dir_name));
    let socket_path = worker_dir.join(FileName::unchecked_new("socket"));
    if fs_util::try_exists(&worker_dir).map_err(|e| WorkerInitError::InternalError(e.into()))? {
        return Err(WorkerInitError::InternalError(
            anyhow::anyhow!("Directory for worker already exists: {:?}", worker_dir).into(),
        ));
    }
    // TODO(ctolliday) put these in buck-out/<iso>/workers and only use /tmp dir for sockets
    let stdout_path = worker_dir.join(FileName::unchecked_new("stdout"));
    let stderr_path = worker_dir.join(FileName::unchecked_new("stderr"));
    fs_util::create_dir_all(&worker_dir).map_err(|e| WorkerInitError::InternalError(e.into()))?;

    let args = worker_spec.exe.to_vec();
    tracing::info!(
        "Starting worker with logs at {}:\n$ {}\n",
        worker_dir,
        args.join(" ")
    );

    let worker_env = vec![("WORKER_SOCKET", socket_path.as_os_str())]
        .into_iter()
        .map(|(k, v)| (OsString::from(k), OsString::from(v)));
    let env: Vec<(OsString, OsString)> = env.into_iter().chain(worker_env).collect();

    let (liveliness_observer, liveliness_guard) = LivelinessGuard::create();

    let spawn_fut = spawn_via_forkserver(
        forkserver,
        OsString::from(args[0].clone()),
        args[1..].iter().map(OsString::from).collect(),
        env.clone(),
        root.clone(),
        liveliness_observer,
        &stdout_path,
        &stderr_path,
        &socket_path,
    );

    let initial_delay = Duration::from_millis(50);
    let max_delay = Duration::from_millis(500);
    // Might want to make this configurable, and/or measure impact of worker initialization on critical path
    let timeout = Duration::from_secs(60);
    let channel = {
        let stdout_path = &stdout_path;
        let stderr_path = &stderr_path;
        let socket_path = &socket_path;

        let connect = retrying(initial_delay, max_delay, timeout, move || {
            // TODO(ctolliday) T153604304
            // add handshake over grpc before returning a handle, to make sure the worker is responding
            get_channel_uds(socket_path, false)
        });

        let check_exit = async move {
            spawn_fut
                .await
                .map_err(|e| WorkerInitError::InternalError(e.into()))?
        };
        futures::pin_mut!(connect);
        futures::pin_mut!(check_exit);

        match futures::future::select(connect, check_exit).await {
            futures::future::Either::Left((connection_result, _)) => match connection_result {
                Ok(channel) => Ok(channel),
                Err(e) => Err(WorkerInitError::ConnectionTimeout(
                    timeout.as_secs_f64(),
                    e.to_string(),
                )),
            },
            futures::future::Either::Right((command_result, _)) => Err(match command_result {
                Ok(GatherOutputStatus::SpawnFailed(e)) => WorkerInitError::SpawnFailed(e),
                Ok(GatherOutputStatus::Finished { exit_code, .. }) => {
                    let stdout = fs_util::read_to_string(stdout_path)
                        .map_err(|e| WorkerInitError::InternalError(e.into()))?;
                    let stderr = fs_util::read_to_string(stderr_path)
                        .map_err(|e| WorkerInitError::InternalError(e.into()))?;
                    WorkerInitError::EarlyExit {
                        exit_code: Some(exit_code),
                        stdout,
                        stderr,
                    }
                }
                Ok(GatherOutputStatus::Cancelled | GatherOutputStatus::TimedOut(_)) => {
                    WorkerInitError::InternalError(
                        anyhow::anyhow!("Worker cancelled by buck").into(),
                    )
                }
                Err(e) => WorkerInitError::InternalError(e.into()),
            }),
        }?
    };

    tracing::info!("Connected to socket for spawned worker: {}", socket_path);
    let client = WorkerClient::new(channel);
    Ok(WorkerHandle {
        client,
        stdout_path,
        stderr_path,
        _liveliness_guard: liveliness_guard,
    })
}

type WorkerFuture = Shared<BoxFuture<'static, Result<Arc<WorkerHandle>, Arc<WorkerInitError>>>>;

pub struct WorkerPool {
    workers: Arc<parking_lot::Mutex<HashMap<WorkerId, WorkerFuture>>>,
}

impl WorkerPool {
    pub fn new() -> WorkerPool {
        tracing::info!("Creating new WorkerPool");
        WorkerPool {
            workers: Arc::new(parking_lot::Mutex::new(HashMap::default())),
        }
    }

    pub fn get_or_create_worker(
        &self,
        worker_spec: &WorkerSpec,
        env: impl IntoIterator<Item = (OsString, OsString)>,
        root: &AbsNormPathBuf,
        forkserver: ForkserverClient,
        dispatcher: EventDispatcher,
    ) -> WorkerFuture {
        let mut workers = self.workers.lock();
        if let Some(worker_fut) = workers.get(&worker_spec.id) {
            worker_fut.clone()
        } else {
            let worker_id = worker_spec.id;
            let worker_spec = worker_spec.clone();
            let root = root.clone();
            let env: Vec<(OsString, OsString)> = env.into_iter().collect();
            let fut = async move {
                match spawn_worker(&worker_spec, env, &root, forkserver, dispatcher).await {
                    Ok(worker) => Ok(Arc::new(worker)),
                    Err(e) => Err(Arc::new(e)),
                }
            }
            .boxed()
            .shared();

            workers.insert(worker_id, fut.clone());
            fut
        }
    }
}

pub struct WorkerHandle {
    client: WorkerClient<Channel>,
    stdout_path: AbsNormPathBuf,
    stderr_path: AbsNormPathBuf,
    _liveliness_guard: LivelinessGuard,
}

#[cfg(unix)]
fn env_entries(env: &[(OsString, OsString)]) -> Vec<EnvironmentEntry> {
    use std::os::unix::ffi::OsStrExt;
    env.iter()
        .map(|(k, v)| EnvironmentEntry {
            key: k.as_bytes().into(),
            value: v.as_bytes().into(),
        })
        .collect()
}

#[cfg(not(unix))]
fn env_entries(_env: &[(OsString, OsString)]) -> Vec<EnvironmentEntry> {
    unreachable!("worker should not exist off unix")
}

impl WorkerHandle {
    pub async fn exec_cmd(
        &self,
        args: &[String],
        env: Vec<(OsString, OsString)>,
    ) -> (GatherOutputStatus, Vec<u8>, Vec<u8>) {
        tracing::info!(
            "Sending worker command:\nExecuteCommand {{ argv: {:?}, env: {:?} }}\n",
            args,
            env,
        );
        let argv: Vec<Vec<u8>> = args.iter().map(|s| s.as_str().into()).collect();
        let env: Vec<EnvironmentEntry> = env_entries(&env);

        let request = ExecuteCommand { argv, env };
        let response = self.client.clone().execute(request).await;

        match response {
            Ok(response) => {
                let exec_response: ExecuteResponse = response.into_inner();
                tracing::info!("Worker response:\n{:?}\n", exec_response);
                (
                    GatherOutputStatus::Finished {
                        exit_code: exec_response.exit_code,
                        execution_stats: None,
                    },
                    vec![],
                    exec_response.stderr.into(),
                )
            }
            Err(err) => {
                (
                    GatherOutputStatus::SpawnFailed(format!(
                        "Error sending ExecuteCommand to worker: {:?}, see worker logs:\n{}\n{}",
                        err, self.stdout_path, self.stderr_path,
                    )),
                    // stdout/stderr logs for worker are for multiple commands, probably do not want to dump contents here
                    vec![],
                    vec![],
                )
            }
        }
    }
}


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

use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_common::http::counting_client::CountingHttpClient;
use buck2_core::directory::DirectoryEntry;
use buck2_core::directory::FingerprintedDirectory;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_execute::artifact_value::ArtifactValue;
use buck2_execute::digest::CasDigestToReExt;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::insert_artifact;
use buck2_execute::directory::ActionDirectoryBuilder;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::execute::blobs::ActionBlobs;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::materialize::eden_api::EdenBuckOut;
use buck2_execute::materialize::materializer::ArtifactNotMaterializedReason;
use buck2_execute::materialize::materializer::CasDownloadInfo;
use buck2_execute::materialize::materializer::CopiedArtifact;
use buck2_execute::materialize::materializer::DeclareMatchOutcome;
use buck2_execute::materialize::materializer::HttpDownloadInfo;
use buck2_execute::materialize::materializer::MaterializationError;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::materialize::materializer::WriteRequest;
use buck2_execute::re::manager::ReConnectionManager;
use dupe::Dupe;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use gazebo::prelude::*;
use more_futures::cancellation::CancellationContext;
use remote_execution::InlinedBlobWithDigest;
use remote_execution::NamedDigest;

use crate::materializers::immediate::ImmediateMaterializer;

#[derive(Allocative)]
pub struct EdenMaterializer {
    re_client_manager: Arc<ReConnectionManager>,
    delegator: Arc<dyn Materializer>,
    eden_buck_out: EdenBuckOut,
    fs: ProjectRoot,
    re_use_case: RemoteExecutorUseCase,
    digest_config: DigestConfig,
}

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

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

    /// For Eden, copying could be an expensive operation when large amount of
    /// file system materialization is required. Instead, uploading to CAS then
    /// declare on Eden would be faster so not all tree nodes would be actually materialized.
    async fn declare_copy_impl(
        &self,
        path: ProjectRelativePathBuf,
        value: ArtifactValue,
        srcs: Vec<CopiedArtifact>,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        // Use eden's remove_paths_recursive because it's faster.
        self.eden_buck_out
            .remove_paths_recursive(&self.fs, vec![path.to_owned()], cancellations)
            .await?;

        // First upload the src to CAS if missing
        let mut files: Vec<remote_execution::NamedDigest> = Vec::new();
        let mut directories: Vec<remote_execution::Path> = Vec::new();
        for copied_artifact in srcs {
            match copied_artifact.dest_entry {
                DirectoryEntry::Leaf(ActionDirectoryMember::File(file)) => {
                    files.push(NamedDigest {
                        name: self
                            .fs
                            .resolve(&copied_artifact.src)
                            .as_maybe_relativized_str()?
                            .to_owned(),
                        digest: file.digest.to_re(),
                        ..Default::default()
                    });
                }
                DirectoryEntry::Dir(dir) => {
                    directories.push(remote_execution::Path {
                        path: self
                            .fs
                            .resolve(&copied_artifact.src)
                            .as_maybe_relativized_str()?
                            .to_owned(),
                        follow_symlinks: false,
                        digest: Some(dir.fingerprint().to_re()),
                        ..Default::default()
                    });
                }
                DirectoryEntry::Leaf(..) => continue,
            };
        }

        self.re_client_manager
            .get_re_connection()
            .get_client()
            .upload_files_and_directories(files, directories, Vec::new(), self.re_use_case)
            .await?;

        // Second upload the tree structure that contains directories/file/symlink metadata
        // TODO(yipu) We don't need to upload CAS, and we should pass ArtifactValue to eden directly
        let mut builder = ActionDirectoryBuilder::empty();
        insert_artifact(&mut builder, path.as_ref(), &value)?;
        let input_dir = builder.fingerprint(self.digest_config.as_directory_serializer());

        self.re_client_manager
            .get_re_connection()
            .get_client()
            .upload(
                &self.fs,
                &self.delegator,
                &ActionBlobs::new(self.digest_config),
                ProjectRelativePath::empty(),
                &input_dir,
                self.re_use_case,
                self.digest_config,
            )
            .await?;

        self.eden_buck_out
            .set_path_object_id(&path, &value)
            .await
            .with_context(|| {
                format!(
                    "[eden] Error declaring artifact {:?} at path {}",
                    value, path
                )
            })?;

        Ok(())
    }

    // This method will call Eden's setPathObjectId method, which is to placehold a
    // tree or a blob to a path of an Eden mount.
    async fn declare_cas_many_impl<'a, 'b>(
        &self,
        _info: Arc<CasDownloadInfo>,
        artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        // Use eden's remove_paths_recursive because it's faster.
        self.eden_buck_out
            .remove_paths_recursive(
                &self.fs,
                artifacts.map(|(p, _)| p.to_owned()),
                cancellations,
            )
            .await?;

        let futs = artifacts.iter().map(|(path, value)| async move {
            self.eden_buck_out
                .set_path_object_id(path, value)
                .await
                .with_context(|| {
                    format!(
                        "[eden] Error declaring artifact {:?} at path {}",
                        value, path
                    )
                })
        });

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

        Ok(())
    }

    async fn declare_http(
        &self,
        path: ProjectRelativePathBuf,
        info: HttpDownloadInfo,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        // Use eden's remove_paths_recursive because it's faster.
        self.eden_buck_out
            .remove_paths_recursive(&self.fs, vec![path.to_owned()], cancellations)
            .await?;

        self.delegator.declare_http(path, info, cancellations).await
    }

    async fn declare_match(
        &self,
        _artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
    ) -> anyhow::Result<DeclareMatchOutcome> {
        // This would require querying Eden at a minimum.
        Ok(DeclareMatchOutcome::NotMatch)
    }

    async fn declare_write<'a>(
        &self,
        gen: Box<dyn FnOnce() -> anyhow::Result<Vec<WriteRequest>> + Send + 'a>,
    ) -> anyhow::Result<Vec<ArtifactValue>> {
        let (paths, values) = write_to_cas(
            self.re_client_manager.as_ref(),
            self.re_use_case,
            self.digest_config,
            gen,
        )
        .await?;

        futures::future::try_join_all(
            std::iter::zip(paths.iter(), values.iter())
                .map(|(path, value)| self.eden_buck_out.set_path_object_id(path, value)),
        )
        .await?;

        Ok(values)
    }

    async fn materialize_many(
        &self,
        artifact_paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<BoxStream<'static, Result<(), MaterializationError>>> {
        // EdenFS' thrift method ensureMaterialized will force materializing a list of provided paths
        self.eden_buck_out
            .ensure_materialized(artifact_paths.clone())
            .await?;
        Ok(stream::iter(artifact_paths.into_iter().map(|_| Ok(()))).boxed())
    }

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

    async fn try_materialize_final_artifact(
        &self,
        artifact_path: ProjectRelativePathBuf,
    ) -> anyhow::Result<bool> {
        self.ensure_materialized(vec![artifact_path]).await?;
        Ok(true)
    }

    async fn invalidate_many(&self, paths: Vec<ProjectRelativePathBuf>) -> anyhow::Result<()> {
        self.delegator.invalidate_many(paths).await
    }

    fn eden_buck_out(&self) -> Option<&EdenBuckOut> {
        Some(&self.eden_buck_out)
    }
}

impl EdenMaterializer {
    pub fn new(
        fs: ProjectRoot,
        digest_config: DigestConfig,
        re_client_manager: Arc<ReConnectionManager>,
        blocking_executor: Arc<dyn BlockingExecutor>,
        eden_buck_out: EdenBuckOut,
        http_client: CountingHttpClient,
    ) -> anyhow::Result<Self> {
        Ok(Self {
            re_client_manager: re_client_manager.dupe(),
            delegator: Arc::new(ImmediateMaterializer::new(
                fs.dupe(),
                digest_config,
                re_client_manager,
                blocking_executor,
                http_client,
            )),
            eden_buck_out,
            fs,
            digest_config,
            re_use_case: RemoteExecutorUseCase::buck2_default(), // TODO (yipu): Should this be configurable?
        })
    }
}

async fn write_to_cas<'a>(
    re: &ReConnectionManager,
    re_use_case: RemoteExecutorUseCase,
    digest_config: DigestConfig,
    gen: Box<dyn FnOnce() -> anyhow::Result<Vec<WriteRequest>> + Send + 'a>,
) -> anyhow::Result<(Vec<ProjectRelativePathBuf>, Vec<ArtifactValue>)> {
    let contents = gen()?;

    let mut uploads = Vec::with_capacity(contents.len());
    let mut paths = Vec::with_capacity(contents.len());
    let mut values = Vec::with_capacity(contents.len());

    for WriteRequest {
        path,
        content,
        is_executable,
    } in contents
    {
        let digest = TrackedFileDigest::from_content(&content, digest_config.cas_digest_config());

        let meta = FileMetadata {
            digest,
            is_executable,
        };

        uploads.push(InlinedBlobWithDigest {
            blob: content,
            digest: meta.digest.to_re(),
            ..Default::default()
        });
        paths.push(path);
        values.push(ArtifactValue::file(meta));
    }

    re.get_re_connection()
        .get_client()
        .upload_files_and_directories(Vec::new(), Vec::new(), uploads, re_use_case)
        .await?;

    Ok((paths, values))
}


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

use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use async_trait::async_trait;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_common::http::counting_client::CountingHttpClient;
use buck2_core::directory::unordered_entry_walk;
use buck2_core::directory::DirectoryEntry;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_execute::artifact_value::ArtifactValue;
use buck2_execute::digest::CasDigestToReExt;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::execute::blocking::BlockingExecutor;
use buck2_execute::execute::clean_output_paths::cleanup_path;
use buck2_execute::execute::clean_output_paths::CleanOutputPaths;
use buck2_execute::materialize::http::http_download;
use buck2_execute::materialize::materializer::ArtifactNotMaterializedReason;
use buck2_execute::materialize::materializer::CasDownloadInfo;
use buck2_execute::materialize::materializer::CopiedArtifact;
use buck2_execute::materialize::materializer::DeclareMatchOutcome;
use buck2_execute::materialize::materializer::HttpDownloadInfo;
use buck2_execute::materialize::materializer::MaterializationError;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::materialize::materializer::WriteRequest;
use buck2_execute::re::manager::ReConnectionManager;
use dupe::Dupe;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use gazebo::prelude::*;
use more_futures::cancellation::CancellationContext;
use remote_execution::NamedDigest;
use remote_execution::NamedDigestWithPermissions;

use crate::materializers::io::materialize_files;
use crate::materializers::io::MaterializeTreeStructure;

/// Materializer that materializes everything immediately on declare.
#[derive(Allocative)]
pub struct ImmediateMaterializer {
    fs: ProjectRoot,
    digest_config: DigestConfig,
    re_client_manager: Arc<ReConnectionManager>,
    io_executor: Arc<dyn BlockingExecutor>,
    http_client: CountingHttpClient,
}

impl ImmediateMaterializer {
    pub fn new(
        fs: ProjectRoot,
        digest_config: DigestConfig,
        re_client_manager: Arc<ReConnectionManager>,
        io_executor: Arc<dyn BlockingExecutor>,
        http_client: CountingHttpClient,
    ) -> Self {
        Self {
            fs,
            digest_config,
            re_client_manager,
            io_executor,
            http_client,
        }
    }
}

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

    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<()> {
        self.io_executor
            .execute_io(
                Box::new(CleanOutputPaths {
                    paths: vec![path.to_owned()],
                }),
                cancellations,
            )
            .await?;

        // TODO: display [materializing] in superconsole
        self.io_executor
            .execute_io(
                Box::new(MaterializeTreeStructure {
                    path: path.clone(),
                    entry: value.entry().dupe(),
                }),
                cancellations,
            )
            .await?;

        self.io_executor
            .execute_io_inline(|| {
                for copied_artifact in srcs {
                    // Make sure `path` is a prefix of `dest`, so we don't
                    // materialize anything outside `path`.
                    copied_artifact.dest.strip_prefix(&path)
                        .with_context(|| format!(
                            "declare_copy: artifact at `{}` copies into `{}`. This is a bug in Buck, not a user error.",
                            path,
                            &copied_artifact.dest,
                        ))?;
                    materialize_files(
                        copied_artifact.dest_entry.as_ref(),
                        &self.fs.root().join(&copied_artifact.src),
                        &self.fs.root().join(&copied_artifact.dest),
                    )?;
                }
                Ok(())
            })
            .await
    }

    async fn declare_cas_many_impl<'a, 'b>(
        &self,
        info: Arc<CasDownloadInfo>,
        artifacts: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        self.io_executor
            .execute_io(
                Box::new(CleanOutputPaths {
                    paths: artifacts.map(|(p, _)| p.to_owned()),
                }),
                cancellations,
            )
            .await?;

        for (path, value) in artifacts.iter() {
            self.io_executor
                .execute_io(
                    Box::new(MaterializeTreeStructure {
                        path: path.to_owned(),
                        entry: value.entry().dupe(),
                    }),
                    cancellations,
                )
                .await?;
        }

        let mut files = Vec::new();
        for (path, value) in artifacts.iter() {
            let mut walk = unordered_entry_walk(value.entry().as_ref());
            while let Some((entry_path, entry)) = walk.next() {
                if let DirectoryEntry::Leaf(ActionDirectoryMember::File(m)) = entry {
                    files.push(NamedDigestWithPermissions {
                        named_digest: NamedDigest {
                            digest: m.digest.to_re(),
                            name: self
                                .fs
                                .resolve(&path.join_normalized(entry_path.get())?)
                                .as_maybe_relativized_str()?
                                .to_owned(),
                            ..Default::default()
                        },
                        is_executable: m.is_executable,
                        ..Default::default()
                    });
                }
            }
        }

        let re_conn = self.re_client_manager.get_re_connection();
        let re_client = re_conn.get_client();
        cancellations
            .critical_section(|| re_client.materialize_files(files, info.re_use_case))
            .await?;
        Ok(())
    }

    async fn declare_http(
        &self,
        path: ProjectRelativePathBuf,
        info: HttpDownloadInfo,
        cancellations: &CancellationContext,
    ) -> anyhow::Result<()> {
        self.io_executor
            .execute_io(
                Box::new(CleanOutputPaths {
                    paths: vec![path.to_owned()],
                }),
                cancellations,
            )
            .await?;

        http_download(
            &self.http_client,
            &self.fs,
            self.digest_config,
            &path,
            &info.url,
            &info.checksum,
            info.metadata.is_executable,
        )
        .await?;

        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>> {
        write_to_disk(&self.fs, self.io_executor.as_ref(), self.digest_config, gen).await
    }

    async fn invalidate_many(&self, _paths: Vec<ProjectRelativePathBuf>) -> anyhow::Result<()> {
        // Nothing to do, we don't keep track of anything.
        Ok(())
    }

    async fn materialize_many(
        &self,
        artifact_paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<BoxStream<'static, Result<(), MaterializationError>>> {
        // We materialize on `declare`, so at this point all declared artifacts
        // have already been materialized.
        Ok(stream::iter(artifact_paths.into_iter().map(|_| Ok(()))).boxed())
    }

    async fn try_materialize_final_artifact(
        &self,
        _artifact_path: ProjectRelativePathBuf,
    ) -> anyhow::Result<bool> {
        // As long as it was declared, it was already materialized
        Ok(true)
    }

    async fn get_materialized_file_paths(
        &self,
        paths: Vec<ProjectRelativePathBuf>,
    ) -> anyhow::Result<Vec<Result<ProjectRelativePathBuf, ArtifactNotMaterializedReason>>> {
        // We materialize on `declare`, so all declared paths are already
        // materialized. We can simply return them as is.
        Ok(paths.into_map(Ok))
    }
}

pub async fn write_to_disk<'a>(
    fs: &ProjectRoot,
    io_executor: &dyn BlockingExecutor,
    digest_config: DigestConfig,
    gen: Box<dyn FnOnce() -> anyhow::Result<Vec<WriteRequest>> + Send + 'a>,
) -> anyhow::Result<Vec<ArtifactValue>> {
    io_executor
        .execute_io_inline({
            move || {
                let requests = gen()?;
                let mut values = Vec::with_capacity(requests.len());

                for WriteRequest {
                    path,
                    content,
                    is_executable,
                } in requests
                {
                    let digest = TrackedFileDigest::from_content(
                        &content,
                        digest_config.cas_digest_config(),
                    );
                    cleanup_path(fs, &path)?;
                    fs.write_file(&path, &content, is_executable)?;

                    values.push(ArtifactValue::file(FileMetadata {
                        digest,
                        is_executable,
                    }));
                }

                Ok(values)
            }
        })
        .await
}


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

use std::collections::HashMap;

use buck2_core::directory::DirectoryEntry;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
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_execute::directory::ActionDirectory;
use buck2_execute::directory::ActionDirectoryEntry;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::directory::ActionSharedDirectory;
use buck2_execute::execute::blocking::IoRequest;

pub struct MaterializeTreeStructure {
    pub path: ProjectRelativePathBuf,
    pub entry: ActionDirectoryEntry<ActionSharedDirectory>,
}

impl IoRequest for MaterializeTreeStructure {
    fn execute(self: Box<Self>, project_fs: &ProjectRoot) -> anyhow::Result<()> {
        materialize_dirs_and_syms(self.entry.as_ref(), project_fs.root().join(&self.path))?;

        Ok(())
    }
}

/// Materializes the entry at `dest`.
///
/// - `materialize_dirs_and_syms`: if `true`, materializes directories and
///   symlinks.
/// - `file_src`: takes the destination path of a file, and returns its
///   source path (where it should be copied from). If it returns [`None`],
///   the file is not materialized.
fn materialize<F, D>(
    entry: DirectoryEntry<&D, &ActionDirectoryMember>,
    dest: &AbsNormPath,
    materialize_dirs_and_syms: bool,
    mut file_src: F,
) -> anyhow::Result<()>
where
    F: FnMut(&AbsNormPath) -> Option<AbsNormPathBuf>,
    D: ActionDirectory,
{
    let mut dest = dest.to_owned();
    if materialize_dirs_and_syms {
        // create the directory where we'll materialize the entry
        if let Some(parent) = dest.parent() {
            fs_util::create_dir_all(parent)?;
        }
    }
    materialize_recursively(entry, &mut dest, materialize_dirs_and_syms, &mut file_src)
}

/// Materializes the directories and symlinks of an entry at `dest`. Files
/// are not materialized.
pub(crate) fn materialize_dirs_and_syms<P, D>(
    entry: DirectoryEntry<&D, &ActionDirectoryMember>,
    dest: P,
) -> anyhow::Result<()>
where
    P: AsRef<AbsNormPath>,
    D: ActionDirectory,
{
    materialize(entry, dest.as_ref(), true, |_: &AbsNormPath| None)
}

/// Materializes the files of an the entry rooted at `dest`.
///
/// Files are copied from `src`. In other words, if a file would be
/// materialized at `dest/p`, then it's copied from `src/p`.
pub(crate) fn materialize_files<P, D>(
    entry: DirectoryEntry<&D, &ActionDirectoryMember>,
    src: P,
    dest: P,
) -> anyhow::Result<()>
where
    P: AsRef<AbsNormPath>,
    D: ActionDirectory,
{
    let src = src.as_ref();
    let dest = dest.as_ref();
    let file_src = |d: &AbsNormPath| {
        // It's safe to unwrap because `materialize_impl` always gives us a
        // path inside `dest`.
        let subpath = d.strip_prefix(dest).unwrap();
        if subpath.as_str().is_empty() {
            // `dest` itself is a file
            Some(src.to_buf())
        } else {
            Some(src.join(subpath))
        }
    };
    materialize(entry, dest, false, file_src)
}

/// Materializes the files of an entry rooted at `dest`.
///
/// For a file at path `file_dest` in the entry, if `file_dest` exists in
/// `srcs` with value `file_src`, the file is copied from `file_src` to
/// `file_dest`. It's then removed from `srcs`.
fn _materialize_files_from_map<P, D>(
    entry: DirectoryEntry<&D, &ActionDirectoryMember>,
    srcs: &mut HashMap<AbsNormPathBuf, AbsNormPathBuf>,
    dest: P,
) -> anyhow::Result<()>
where
    P: AsRef<AbsNormPath>,
    D: ActionDirectory,
{
    let file_src = |d: &AbsNormPath| srcs.remove(d);
    materialize(entry, dest.as_ref(), false, file_src)
}

fn materialize_recursively<F, D>(
    entry: DirectoryEntry<&D, &ActionDirectoryMember>,
    dest: &mut AbsNormPathBuf,
    materialize_dirs_and_syms: bool,
    file_src: &mut F,
) -> anyhow::Result<()>
where
    F: FnMut(&AbsNormPath) -> Option<AbsNormPathBuf>,
    D: ActionDirectory + ?Sized,
{
    match entry {
        DirectoryEntry::Dir(d) => {
            if materialize_dirs_and_syms {
                fs_util::create_dir_all(&dest)?;
            }
            for (name, entry) in d.entries() {
                dest.push(name);
                materialize_recursively(entry, dest, materialize_dirs_and_syms, file_src)?;
                dest.pop();
            }
            Ok(())
        }
        DirectoryEntry::Leaf(ActionDirectoryMember::File(_)) => {
            if let Some(src) = file_src(dest) {
                fs_util::copy(src, dest)?;
            }
            Ok(())
        }
        DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(s)) => {
            if materialize_dirs_and_syms && fs_util::symlink_metadata(&dest).is_err() {
                fs_util::symlink(s.target().as_str(), dest)?;
            }
            Ok(())
        }
        DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(s)) => {
            if materialize_dirs_and_syms && fs_util::symlink_metadata(&dest).is_err() {
                fs_util::symlink(s.target(), dest)?;
            }
            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.
 */

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

pub mod deferred;
pub mod immediate;
pub mod io;
pub mod sqlite;


/*
 * 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 allocative::Allocative;
use anyhow::Context;
use buck2_common::external_symlink::ExternalSymlink;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_common::sqlite::KeyValueSqliteTable;
use buck2_core::directory::DirectoryEntry;
use buck2_core::fs::fs_util;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::paths::abs_norm_path::AbsNormPathBuf;
use buck2_core::fs::paths::file_name::FileName;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::directory::Symlink;
use buck2_execute::execute::blocking::BlockingExecutor;
use chrono::DateTime;
use chrono::TimeZone;
use chrono::Utc;
use derive_more::Display;
use derive_more::From;
use dupe::Dupe;
use gazebo::prelude::*;
use itertools::Itertools;
use once_cell::sync::Lazy;
use parking_lot::Mutex;
use rusqlite::Connection;
use thiserror::Error;

use crate::materializers::deferred::ArtifactMetadata;
use crate::materializers::deferred::DirectoryMetadata;

#[derive(Display, Allocative, Clone, From, PartialEq, Eq, Debug)]
pub struct MaterializerStateIdentity(String);

/// Hand-maintained schema version for the materializer state sqlite db.
/// PLEASE bump this version if you are making a breaking change to the
/// materializer state sqlite db schema! If you forget to bump this version,
/// then you can fix forward by bumping the `buck2.sqlite_materializer_state_version`
/// buckconfig in the project root's .buckconfig.
pub const DB_SCHEMA_VERSION: u64 = 6;

const STATE_TABLE_NAME: &str = "materializer_state";
const IDENTITY_KEY: &str = "timestamp_on_initialization";

pub type MaterializerState = Vec<(ProjectRelativePathBuf, (ArtifactMetadata, DateTime<Utc>))>;

#[derive(Error, Debug, PartialEq, Eq)]
pub(crate) enum ArtifactMetadataSqliteConversionError {
    #[error("Internal error: expected field `{}` to be not null for artifact type '{}'", .field, .artifact_type)]
    ExpectedNotNull {
        field: String,
        artifact_type: String,
    },

    #[error("Internal error: found unknown value '{}' for enum `artifact_type`", .0)]
    UnknownArtifactType(String),
}

/// Sqlite representation of sha1. Can be converted directly into BLOB type
/// in rusqlite. Note we use a Vec and not a fixed length array here because
/// rusqlite can only convert to Vec.
pub(crate) type SqliteDigest = Vec<u8>;

/// Sqlite representation of `ArtifactMetadata`. All datatypes used implement
/// rusqlite's `FromSql` trait.
#[derive(Debug)]
pub(crate) struct ArtifactMetadataSqliteEntry {
    pub artifact_type: String,
    pub entry_size: Option<u64>,
    pub entry_hash: Option<SqliteDigest>,
    pub entry_hash_kind: Option<u8>,
    pub file_is_executable: Option<bool>,
    pub symlink_target: Option<String>,
    pub directory_size: Option<u64>,
}

impl ArtifactMetadataSqliteEntry {
    pub(crate) fn new(
        artifact_type: String,
        entry_size: Option<u64>,
        entry_hash: Option<Vec<u8>>,
        entry_hash_kind: Option<u8>,
        file_is_executable: Option<bool>,
        symlink_target: Option<String>,
        directory_size: Option<u64>,
    ) -> Self {
        Self {
            artifact_type,
            entry_size,
            entry_hash,
            entry_hash_kind,
            file_is_executable,
            symlink_target,
            directory_size,
        }
    }
}

impl From<&ArtifactMetadata> for ArtifactMetadataSqliteEntry {
    fn from(metadata: &ArtifactMetadata) -> Self {
        fn digest_parts(digest: &TrackedFileDigest) -> (u64, Vec<u8>, u8) {
            (
                digest.size(),
                digest.raw_digest().as_bytes().to_vec(),
                digest.raw_digest().algorithm() as _,
            )
        }

        let (
            artifact_type,
            entry_size,
            entry_hash,
            entry_hash_kind,
            file_is_executable,
            symlink_target,
            directory_size,
        ) = match &metadata.0 {
            DirectoryEntry::Dir(meta) => {
                let (entry_size, entry_hash, entry_hash_kind) = digest_parts(&meta.fingerprint);
                (
                    "directory",
                    Some(entry_size),
                    Some(entry_hash),
                    Some(entry_hash_kind),
                    None,
                    None,
                    Some(meta.total_size),
                )
            }
            DirectoryEntry::Leaf(ActionDirectoryMember::File(file_metadata)) => {
                let (entry_size, entry_hash, entry_hash_kind) = digest_parts(&file_metadata.digest);
                (
                    "file",
                    Some(entry_size),
                    Some(entry_hash),
                    Some(entry_hash_kind),
                    Some(file_metadata.is_executable),
                    None,
                    None,
                )
            }
            DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(symlink)) => (
                "symlink",
                None,
                None,
                None,
                None,
                Some(symlink.target().as_str().to_owned()),
                None,
            ),
            DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(external_symlink)) => (
                "external_symlink",
                None,
                None,
                None,
                None,
                Some(external_symlink.target_str().to_owned()),
                None,
            ),
        };

        Self {
            artifact_type: artifact_type.to_owned(),
            entry_size,
            entry_hash,
            entry_hash_kind,
            file_is_executable,
            symlink_target,
            directory_size,
        }
    }
}

fn convert_artifact_metadata(
    sqlite_entry: ArtifactMetadataSqliteEntry,
    digest_config: DigestConfig,
) -> anyhow::Result<ArtifactMetadata> {
    fn digest(
        size: Option<u64>,
        entry_hash: Option<Vec<u8>>,
        entry_hash_kind: Option<u8>,
        artifact_type: &str,
        digest_config: DigestConfig,
    ) -> anyhow::Result<TrackedFileDigest> {
        let size = size.ok_or_else(|| {
            anyhow::anyhow!(ArtifactMetadataSqliteConversionError::ExpectedNotNull {
                field: "size".to_owned(),
                artifact_type: artifact_type.to_owned()
            })
        })?;
        let entry_hash = entry_hash.ok_or_else(|| {
            anyhow::anyhow!(ArtifactMetadataSqliteConversionError::ExpectedNotNull {
                field: "entry_hash".to_owned(),
                artifact_type: artifact_type.to_owned()
            })
        })?;
        let entry_hash_kind = entry_hash_kind.ok_or_else(|| {
            anyhow::anyhow!(ArtifactMetadataSqliteConversionError::ExpectedNotNull {
                field: "entry_hash_kind".to_owned(),
                artifact_type: artifact_type.to_owned()
            })
        })?;
        let entry_hash_kind = entry_hash_kind
            .try_into()
            .with_context(|| format!("Invalid entry_hash_kind: `{}`", entry_hash_kind))?;

        let file_digest = FileDigest::from_digest_bytes(entry_hash_kind, &entry_hash, size)?;
        Ok(TrackedFileDigest::new(
            file_digest,
            digest_config.cas_digest_config(),
        ))
    }

    let metadata = match sqlite_entry.artifact_type.as_str() {
        "directory" => DirectoryEntry::Dir(DirectoryMetadata {
            fingerprint: digest(
                sqlite_entry.entry_size,
                sqlite_entry.entry_hash,
                sqlite_entry.entry_hash_kind,
                sqlite_entry.artifact_type.as_str(),
                digest_config,
            )?,
            total_size: sqlite_entry
                .directory_size
                .context("Missing directory size")?,
        }),
        "file" => DirectoryEntry::Leaf(ActionDirectoryMember::File(FileMetadata {
            digest: digest(
                sqlite_entry.entry_size,
                sqlite_entry.entry_hash,
                sqlite_entry.entry_hash_kind,
                sqlite_entry.artifact_type.as_str(),
                digest_config,
            )?,
            is_executable: sqlite_entry.file_is_executable.ok_or_else(|| {
                anyhow::anyhow!(ArtifactMetadataSqliteConversionError::ExpectedNotNull {
                    field: "file_is_executable".to_owned(),
                    artifact_type: sqlite_entry.artifact_type.clone()
                })
            })?,
        })),
        "symlink" => {
            let symlink = Symlink::new(
                sqlite_entry
                    .symlink_target
                    .ok_or_else(|| {
                        anyhow::anyhow!(ArtifactMetadataSqliteConversionError::ExpectedNotNull {
                            field: "symlink_target".to_owned(),
                            artifact_type: sqlite_entry.artifact_type.clone()
                        })
                    })?
                    .into(),
            );
            DirectoryEntry::Leaf(ActionDirectoryMember::Symlink(Arc::new(symlink)))
        }
        "external_symlink" => {
            let external_symlink = ExternalSymlink::new(
                sqlite_entry
                    .symlink_target
                    .ok_or_else(|| {
                        anyhow::anyhow!(ArtifactMetadataSqliteConversionError::ExpectedNotNull {
                            field: "symlink_target".to_owned(),
                            artifact_type: sqlite_entry.artifact_type.clone()
                        })
                    })?
                    .into(),
                None,
            )?;
            DirectoryEntry::Leaf(ActionDirectoryMember::ExternalSymlink(Arc::new(
                external_symlink,
            )))
        }
        artifact_type => {
            return Err(anyhow::anyhow!(
                ArtifactMetadataSqliteConversionError::UnknownArtifactType(
                    artifact_type.to_owned()
                )
            ));
        }
    };

    Ok(ArtifactMetadata(metadata))
}

pub(crate) struct MaterializerStateSqliteTable {
    connection: Arc<Mutex<Connection>>,
}

impl MaterializerStateSqliteTable {
    pub fn new(connection: Arc<Mutex<Connection>>) -> Self {
        Self { connection }
    }

    pub(crate) fn create_table(&self) -> anyhow::Result<()> {
        let sql = format!(
            "CREATE TABLE {} (
                path                    TEXT NOT NULL PRIMARY KEY,
                artifact_type           TEXT CHECK(artifact_type IN ('directory','file','symlink','external_symlink')) NOT NULL,
                digest_size             INTEGER NULL DEFAULT NULL,
                entry_hash              BLOB NULL DEFAULT NULL,
                entry_hash_kind         INTEGER NULL DEFAULT NULL,
                file_is_executable      INTEGER NULL DEFAULT NULL,
                symlink_target          TEXT NULL DEFAULT NULL,
                last_access_time        INTEGER NOT NULL,
                directory_size          INTEGER NULL DEFAULT NULL
            )",
            STATE_TABLE_NAME,
        );
        tracing::trace!(sql = %*sql, "creating table");
        self.connection
            .lock()
            .execute(&sql, [])
            .with_context(|| format!("creating sqlite table {}", STATE_TABLE_NAME))?;
        Ok(())
    }

    pub(crate) fn insert(
        &self,
        path: &ProjectRelativePath,
        metadata: &ArtifactMetadata,
        timestamp: DateTime<Utc>,
    ) -> anyhow::Result<()> {
        let entry: ArtifactMetadataSqliteEntry = metadata.into();
        static SQL: Lazy<String> = Lazy::new(|| {
            format!(
                "INSERT INTO {} (path, artifact_type, digest_size, entry_hash, entry_hash_kind, file_is_executable, symlink_target, directory_size, last_access_time) VALUES (?1, ?2, ?3, ?4, ?5, ?6, ?7, ?8, ?9)",
                STATE_TABLE_NAME
            )
        });
        tracing::trace!(sql = %*SQL, entry = ?entry, "inserting into table");
        self.connection
            .lock()
            .execute(
                &SQL,
                rusqlite::params![
                    path.as_str(),
                    entry.artifact_type,
                    entry.entry_size,
                    entry.entry_hash,
                    entry.entry_hash_kind,
                    entry.file_is_executable,
                    entry.symlink_target,
                    entry.directory_size,
                    timestamp.timestamp(),
                ],
            )
            .with_context(|| {
                format!(
                    "inserting `{}` into sqlite table {}",
                    path, STATE_TABLE_NAME
                )
            })?;
        Ok(())
    }

    pub(crate) fn update_access_time(
        &self,
        path: &ProjectRelativePath,
        timestamp: DateTime<Utc>,
    ) -> anyhow::Result<()> {
        static SQL: Lazy<String> = Lazy::new(|| {
            format!(
                "UPDATE {} SET last_access_time = (?1) WHERE path = (?2)",
                STATE_TABLE_NAME,
            )
        });
        tracing::trace!(sql = %*SQL, now = %timestamp, "updating last_access_time");
        self.connection
            .lock()
            .execute(
                &SQL,
                rusqlite::params![timestamp.timestamp(), path.as_str()],
            )
            .with_context(|| format!("updating sqlite table {}", STATE_TABLE_NAME))?;
        Ok(())
    }

    pub(crate) fn read_all(
        &self,
        digest_config: DigestConfig,
    ) -> anyhow::Result<MaterializerState> {
        static SQL: Lazy<String> = Lazy::new(|| {
            format!(
                "SELECT path, artifact_type, digest_size, entry_hash, entry_hash_kind, file_is_executable, symlink_target, directory_size, last_access_time FROM {}",
                STATE_TABLE_NAME,
            )
        });
        tracing::trace!(sql = %*SQL, "reading all from table");
        let connection = self.connection.lock();
        let mut stmt = connection.prepare(&SQL)?;
        let result = stmt
            .query_map(
                [],
                |row| -> rusqlite::Result<(String, ArtifactMetadataSqliteEntry, i64)> {
                    Ok((
                        row.get(0)?,
                        ArtifactMetadataSqliteEntry::new(
                            row.get(1)?,
                            row.get(2)?,
                            row.get(3)?,
                            row.get(4)?,
                            row.get(5)?,
                            row.get(6)?,
                            row.get(7)?,
                        ),
                        row.get(8)?,
                    ))
                },
            )?
            .collect::<Result<Vec<_>, _>>()
            .with_context(|| format!("reading from sqlite table {}", STATE_TABLE_NAME))?;

        result
            .into_try_map(
                |(path, entry, last_access_time)| -> anyhow::Result<(
                    ProjectRelativePathBuf,
                    (ArtifactMetadata, DateTime<Utc>),
                )> {
                    let path = ProjectRelativePathBuf::unchecked_new(path);
                    let metadata = convert_artifact_metadata(entry, digest_config)?;
                    let timestamp = Utc
                        .timestamp_opt(last_access_time, 0)
                        .single()
                        .with_context(|| "invalid timestamp")?;
                    Ok((path, (metadata, timestamp)))
                },
            )
            .with_context(|| format!("error reading row of sqlite table {}", STATE_TABLE_NAME))
    }

    pub(crate) fn delete(&self, paths: Vec<ProjectRelativePathBuf>) -> anyhow::Result<usize> {
        if paths.is_empty() {
            return Ok(0);
        }

        let mut rows_deleted = 0;

        for chunk in paths.chunks(100) {
            let sql = format!(
                "DELETE FROM {} WHERE path IN ({})",
                STATE_TABLE_NAME,
                // According to rusqlite docs this is the best way to generate the right
                // number of query placeholders.
                itertools::repeat_n("?", chunk.len()).join(","),
            );

            tracing::trace!(sql = %sql, chunk = ?chunk, "deleting from table");
            rows_deleted += self
                .connection
                .lock()
                .execute(
                    &sql,
                    rusqlite::params_from_iter(chunk.iter().map(|p| p.as_str())),
                )
                .with_context(|| format!("deleting from sqlite table {}", STATE_TABLE_NAME))?;
        }

        Ok(rows_deleted)
    }
}

#[derive(Error, Debug, PartialEq, Eq)]
enum MaterializerStateSqliteDbError {
    #[error("Path {} does not exist", .0)]
    PathDoesNotExist(AbsNormPathBuf),

    #[error("Expected versions {:?}. Found versions {:?} in sqlite db at {}", .expected, .found, .path)]
    VersionMismatch {
        expected: HashMap<String, String>,
        found: HashMap<String, String>,
        path: AbsNormPathBuf,
    },

    #[error("Materializer identity was rejected: {}", .identity)]
    RejectedIdentity { identity: MaterializerStateIdentity },
}

/// DB that opens the sqlite connection to the materializer state db on disk and
/// holds all the sqlite tables we need for storing/querying materializer state
pub struct MaterializerStateSqliteDb {
    tables: MaterializerStateTables,
    /// A unique ID identifying this particular instance of the database. This will reset when we
    /// recreate it.
    identity: MaterializerStateIdentity,
}

impl MaterializerStateSqliteDb {
    const DB_FILENAME: &'static str = "db.sqlite";

    fn new(tables: MaterializerStateTables) -> anyhow::Result<Self> {
        let identity = tables
            .created_by_table
            .get(IDENTITY_KEY)
            .context("Error reading creation metadata")?
            .map(MaterializerStateIdentity)
            .with_context(|| format!("Identity key is missing in db: `{}`", IDENTITY_KEY))?;

        Ok(Self { tables, identity })
    }

    /// Given path to the sqlite DB, attempts to read `MaterializerState` from the DB. If we encounter
    /// any failure along the way, such as if the DB path does not exist, the sqlite read fails,
    /// or the DB has a different set of versions than the versions this buck2 expects, we
    /// throw away the existing DB and initialize a new DB. Returns (1) the connected sqlite DB and
    /// (2) the `MaterializerState` if loading was successful or the load error.
    /// The `Result<MaterializerState>` captures any failure encountered when attempting to load
    /// from the existing DB. These failures are expected if db doesn't exist or versions don't match.
    /// The outer `Result` captures any failure encountered when trying to delete the existing DB and
    /// create a new one.
    /// TODO(scottcao): pull this method into a shared trait once we add a another sqlite DB
    pub async fn initialize(
        materializer_state_dir: AbsNormPathBuf,
        versions: HashMap<String, String>,
        current_instance_metadata: HashMap<String, String>,
        // Using `BlockingExecutor` out of convenience. This function should be called during startup
        // when there's not a lot of I/O so it shouldn't matter.
        io_executor: Arc<dyn BlockingExecutor>,
        digest_config: DigestConfig,
        reject_identity: Option<&MaterializerStateIdentity>,
    ) -> anyhow::Result<(Self, anyhow::Result<MaterializerState>)> {
        io_executor
            .execute_io_inline(|| {
                Self::initialize_impl(
                    materializer_state_dir,
                    versions,
                    current_instance_metadata,
                    digest_config,
                    reject_identity,
                )
            })
            .await
    }

    fn initialize_impl(
        materializer_state_dir: AbsNormPathBuf,
        versions: HashMap<String, String>,
        mut current_instance_metadata: HashMap<String, String>,
        digest_config: DigestConfig,
        reject_identity: Option<&MaterializerStateIdentity>,
    ) -> anyhow::Result<(Self, anyhow::Result<MaterializerState>)> {
        let timestamp_on_initialization = Utc::now().to_rfc3339();
        current_instance_metadata.insert(IDENTITY_KEY.to_owned(), timestamp_on_initialization);

        let db_path = materializer_state_dir.join(FileName::unchecked_new(Self::DB_FILENAME));

        let result: anyhow::Result<(Self, MaterializerState)> = try {
            // try reading the existing db, if it exists.
            if !db_path.exists() {
                Err(anyhow::anyhow!(
                    MaterializerStateSqliteDbError::PathDoesNotExist(db_path.clone())
                ))?
            }

            let tables = MaterializerStateTables::open(&db_path)?;

            // First check that versions match
            let read_versions = tables.versions_table.read_all()?;
            if read_versions != versions {
                Err(MaterializerStateSqliteDbError::VersionMismatch {
                    expected: versions.clone(),
                    found: read_versions,
                    path: db_path.clone(),
                })?;
            }

            // Update "last_read_by" inside of the try block so that
            // just in case it fails, we can create a new db and start over
            tables
                .last_read_by_table
                .insert_all(current_instance_metadata.clone())?;

            let mut db = Self::new(tables)?;

            if let Some(reject_identity) = reject_identity {
                if db.identity == *reject_identity {
                    Err(MaterializerStateSqliteDbError::RejectedIdentity {
                        identity: db.identity.clone(),
                    })?;
                }
            }

            let state = db.materializer_state_table().read_all(digest_config)?;

            (db, state)
        };

        match result {
            Ok((db, state)) => Ok((db, Ok(state))),
            Err(e) => {
                // Loading failed. Initialize a new db from scratch.

                // Delete the existing materializer_state directory and create a new one.
                // We delete the entire directory and not just the db file because sqlite
                // can leave behind other files.
                if materializer_state_dir.exists() {
                    fs_util::remove_dir_all(&materializer_state_dir)?;
                }
                fs_util::create_dir_all(&materializer_state_dir)?;

                // Initialize a new db
                let tables = MaterializerStateTables::open(&db_path)?;
                tables.create_all_tables()?;
                tables.versions_table.insert_all(versions)?;
                // Update both "last_read_by" and "created_by"
                tables
                    .created_by_table
                    .insert_all(current_instance_metadata.clone())?;
                tables
                    .last_read_by_table
                    .insert_all(current_instance_metadata)?;

                Ok((Self::new(tables)?, Err(e)))
            }
        }
    }

    pub(crate) fn materializer_state_table(&mut self) -> &MaterializerStateSqliteTable {
        &self.tables.materializer_state_table
    }

    pub fn identity(&self) -> &MaterializerStateIdentity {
        &self.identity
    }
}

struct MaterializerStateTables {
    /// Table storing actual materializer state
    materializer_state_table: MaterializerStateSqliteTable,
    /// Table for holding any metadata used to check version match. When loading
    /// from an existing db, we check if the versions from this table match the
    /// versions this buck2 binary expects. If the versions don't match, we throw
    /// away the entire db and initialize a new one. If versions do match, then
    /// we try to read all state from `materializer_state_table`.
    versions_table: KeyValueSqliteTable,
    /// Table for logging metadata associated with the buck2 that created the db.
    created_by_table: KeyValueSqliteTable,
    /// Table for logging metadata associated with the buck2 that last updated the db.
    last_read_by_table: KeyValueSqliteTable,
}

impl MaterializerStateTables {
    /// Given path to sqlite DB, opens and returns a new connection to the DB.
    fn open(path: &AbsNormPath) -> anyhow::Result<Self> {
        let connection = Connection::open(path)?;
        // TODO: make this work on Windows too
        if cfg!(unix) {
            connection.pragma_update(None, "journal_mode", "WAL")?;
        }

        // Setting synchronous to anything but OFF prevents data corruption in case of power loss,
        // but for the deferred materializer state, we are rather happy to run the risk of data
        // corruption (which we recover from by just dropping the state and pretending we have
        // none), rather than running a `fsync` at any point during a build, which tends to be
        // *very* slow, because if we do a `fsync` from the deferred materializer, that will tend
        // to occur after a lot of writes have been done, which means a lot of data needs syncing!
        //
        // Note that upon power loss there isn't really a guarantee of ordering of things written
        // across different files anyway, so we always run some risk of having our state be
        // incorrect if that happens, unless we fsync after every single write, but, that's
        // definitely not an option.
        //
        // This problem notably manifests itself when routing writes through the deferred
        // materializer on a benchmark build. This causes the set of operations done by the
        // deferred materializer to exceed the WAL max size (which is 1000 pages that are 4KB each,
        // so about 4MB of data), which causes SQLite to write the WAL to the database file, which
        // is the only circumstance under which SQLite does a `fsync` when WAL is enabled, and then
        // we completely stall I/O for a little while.
        connection.pragma_update(None, "synchronous", "OFF")?;

        let connection = Arc::new(Mutex::new(connection));
        let materializer_state_table = MaterializerStateSqliteTable::new(connection.dupe());
        let versions_table = KeyValueSqliteTable::new("versions".to_owned(), connection.dupe());
        let created_by_table = KeyValueSqliteTable::new("created_by".to_owned(), connection.dupe());
        let last_read_by_table = KeyValueSqliteTable::new("last_read_by".to_owned(), connection);

        Ok(Self {
            materializer_state_table,
            versions_table,
            created_by_table,
            last_read_by_table,
        })
    }

    fn create_all_tables(&self) -> anyhow::Result<()> {
        self.materializer_state_table.create_table()?;
        self.versions_table.create_table()?;
        self.created_by_table.create_table()?;
        self.last_read_by_table.create_table()?;
        Ok(())
    }
}

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

    use assert_matches::assert_matches;
    use buck2_common::file_ops::FileMetadata;
    use buck2_common::file_ops::TrackedFileDigest;
    use buck2_core::directory::DirectoryEntry;
    use buck2_core::fs::project::ProjectRoot;
    use buck2_core::fs::project::ProjectRootTemp;
    use buck2_core::fs::project_rel_path::ProjectRelativePath;
    use buck2_execute::digest_config::DigestConfig;
    use buck2_execute::directory::new_symlink;
    use buck2_execute::directory::ActionDirectoryMember;

    use super::*;

    #[test]
    fn test_artifact_metadata_dir_sqlite_entry_conversion_succeeds() {
        let digest_config = DigestConfig::testing_default();

        let digest =
            TrackedFileDigest::from_content(b"directory", digest_config.cas_digest_config());
        let metadata = ArtifactMetadata(DirectoryEntry::Dir(DirectoryMetadata {
            fingerprint: digest,
            total_size: 32,
        }));
        let entry = ArtifactMetadataSqliteEntry::from(&metadata);
        assert_eq!(
            metadata,
            convert_artifact_metadata(entry, digest_config).unwrap()
        );
    }

    #[test]
    fn test_artifact_metadata_file_sqlite_entry_conversion_succeeds() {
        let digest_config = DigestConfig::testing_default();

        let digest = TrackedFileDigest::from_content(b"file", digest_config.cas_digest_config());
        let metadata = ArtifactMetadata(DirectoryEntry::Leaf(ActionDirectoryMember::File(
            FileMetadata {
                digest,
                is_executable: false,
            },
        )));
        let entry = ArtifactMetadataSqliteEntry::from(&metadata);
        assert_eq!(
            metadata,
            convert_artifact_metadata(entry, digest_config).unwrap()
        );
    }

    #[test]
    fn test_artifact_metadata_symlink_sqlite_entry_conversion_succeeds() {
        let digest_config = DigestConfig::testing_default();

        let symlink = new_symlink("foo/bar").unwrap();
        // We use `new_symlink` as a helper but it can technically create both Symlink and ExternalSymlink.
        // Verify that we have a Symlink here.
        assert_matches!(symlink, ActionDirectoryMember::Symlink(..));

        let metadata = ArtifactMetadata(DirectoryEntry::Leaf(symlink));
        let entry = ArtifactMetadataSqliteEntry::from(&metadata);
        assert_eq!(
            metadata,
            convert_artifact_metadata(entry, digest_config).unwrap()
        );
    }

    #[test]
    fn test_artifact_metadata_external_symlink_sqlite_entry_conversion_succeeds() {
        let digest_config = DigestConfig::testing_default();

        let external_symlink = new_symlink(if cfg!(windows) {
            // Not sure if we actually support any external symlink on windows, but better
            // to just check anyways.
            "C:\\external\\symlink\\to\\artifact"
        } else {
            "/mnt/gvfs/third-party2/zstd/28def025ee38919d509596da7d09e7a5262cbf32/1.4.x/platform010/64091f4/share"
        }).unwrap();
        // We use `new_symlink` as a helper but it can technically create both Symlink and ExternalSymlink.
        // Verify that we have an ExternalSymlink here.
        assert_matches!(external_symlink, ActionDirectoryMember::ExternalSymlink(..));

        let metadata = ArtifactMetadata(DirectoryEntry::Leaf(external_symlink));
        let entry = ArtifactMetadataSqliteEntry::from(&metadata);
        assert_eq!(
            metadata,
            convert_artifact_metadata(entry, digest_config).unwrap()
        );
    }

    fn now_seconds() -> DateTime<Utc> {
        Utc.timestamp_opt(Utc::now().timestamp(), 0)
            .single()
            .unwrap()
    }

    #[test]
    fn test_materializer_state_sqlite_table() {
        let digest_config = DigestConfig::testing_default();

        let fs = ProjectRootTemp::new().unwrap();
        let connection = Connection::open(
            fs.path()
                .resolve(ProjectRelativePath::unchecked_new("test.db")),
        )
        .unwrap();
        let table = MaterializerStateSqliteTable::new(Arc::new(Mutex::new(connection)));

        table.create_table().unwrap();

        let dir_metadata = DirectoryMetadata {
            fingerprint: TrackedFileDigest::from_content(
                b"directory",
                digest_config.cas_digest_config(),
            ),
            total_size: 32,
        };
        let file = ActionDirectoryMember::File(FileMetadata {
            digest: TrackedFileDigest::from_content(b"file", digest_config.cas_digest_config()),
            is_executable: false,
        });
        let symlink = new_symlink("foo/bar").unwrap();
        let external_symlink = new_symlink(if cfg!(windows) {
            // Not sure if we actually support any external symlink on windows, but better
            // to just check anyways.
            "C:\\external\\symlink\\to\\artifact"
        } else {
            "/mnt/gvfs/third-party2/zstd/28def025ee38919d509596da7d09e7a5262cbf32/1.4.x/platform010/64091f4/share"
        }).unwrap();
        // We use `new_symlink` as a helper but it can technically create both Symlink and ExternalSymlink.
        // Verify that we have proper symlink and external symlink.
        assert_matches!(symlink, ActionDirectoryMember::Symlink(..));
        assert_matches!(external_symlink, ActionDirectoryMember::ExternalSymlink(..));

        let mut artifacts = HashMap::from([
            (
                ProjectRelativePath::unchecked_new("a").to_owned(),
                (
                    ArtifactMetadata(DirectoryEntry::Dir(dir_metadata)),
                    now_seconds(),
                ),
            ),
            (
                ProjectRelativePath::unchecked_new("b/c").to_owned(),
                (ArtifactMetadata(DirectoryEntry::Leaf(file)), now_seconds()),
            ),
            (
                ProjectRelativePath::unchecked_new("d").to_owned(),
                (
                    ArtifactMetadata(DirectoryEntry::Leaf(symlink)),
                    now_seconds(),
                ),
            ),
            (
                ProjectRelativePath::unchecked_new("e").to_owned(),
                (
                    ArtifactMetadata(DirectoryEntry::Leaf(external_symlink)),
                    now_seconds(),
                ),
            ),
        ]);

        for (path, metadata) in artifacts.iter() {
            table.insert(path, &metadata.0, metadata.1).unwrap();
        }

        let state = table.read_all(digest_config).unwrap();
        assert_eq!(artifacts, state.into_iter().collect::<HashMap<_, _>>());

        let paths_to_remove = vec![
            ProjectRelativePath::unchecked_new("d").to_owned(),
            ProjectRelativePath::unchecked_new("doesnt/exist").to_owned(),
        ];
        let rows_deleted = table.delete(paths_to_remove.clone()).unwrap();
        assert_eq!(rows_deleted, 1);

        for path in paths_to_remove.iter() {
            artifacts.remove(path);
        }
        let state = table.read_all(digest_config).unwrap();
        assert_eq!(artifacts, state.into_iter().collect::<HashMap<_, _>>());
    }

    fn testing_materializer_state_sqlite_db(
        fs: &ProjectRoot,
        versions: HashMap<String, String>,
        metadata: HashMap<String, String>,
        reject_identity: Option<&MaterializerStateIdentity>,
    ) -> anyhow::Result<(MaterializerStateSqliteDb, anyhow::Result<MaterializerState>)> {
        MaterializerStateSqliteDb::initialize_impl(
            fs.resolve(ProjectRelativePath::unchecked_new(
                "buck-out/v2/cache/materializer_state",
            )),
            versions,
            metadata,
            DigestConfig::testing_default(),
            reject_identity,
        )
    }

    // Only implementing for tests, actual code should use `matches_entry` (and not check total_size)
    impl PartialEq for DirectoryMetadata {
        fn eq(&self, other: &DirectoryMetadata) -> bool {
            self.fingerprint == other.fingerprint && self.total_size == other.total_size
        }
    }

    impl PartialEq for ArtifactMetadata {
        fn eq(&self, other: &ArtifactMetadata) -> bool {
            self.0 == other.0
        }
    }

    #[test]
    fn test_initialize_sqlite_db() -> anyhow::Result<()> {
        fn testing_metadatas() -> Vec<HashMap<String, String>> {
            let metadata = buck2_events::metadata::collect();
            let mut metadatas = vec![metadata; 5];
            for (i, metadata) in metadatas.iter_mut().enumerate() {
                metadata.insert("version".to_owned(), i.to_string());
            }
            metadatas
        }

        fn assert_metadata_matches(
            mut have: HashMap<String, String>,
            want: &HashMap<String, String>,
        ) {
            // Remove the key we inject (and check it's there).
            have.remove(IDENTITY_KEY).unwrap();
            assert_eq!(have, *want);
        }

        let digest_config = DigestConfig::testing_default();

        let fs = ProjectRootTemp::new()?;

        let path = ProjectRelativePath::unchecked_new("foo").to_owned();
        let artifact_metadata = ArtifactMetadata(DirectoryEntry::Dir(DirectoryMetadata {
            fingerprint: TrackedFileDigest::from_content(
                b"directory",
                digest_config.cas_digest_config(),
            ),
            total_size: 32,
        }));
        let timestamp = now_seconds();
        let metadatas = testing_metadatas();

        let v0 = HashMap::from([("version".to_owned(), "0".to_owned())]);
        let v1 = HashMap::from([("version".to_owned(), "1".to_owned())]);

        {
            let (mut db, loaded_state) = testing_materializer_state_sqlite_db(
                fs.path(),
                v0.clone(),
                metadatas[0].clone(),
                None,
            )
            .unwrap();
            assert_matches!(
                loaded_state,
                Err(e) => {
                    assert_matches!(
                        e.downcast_ref::<MaterializerStateSqliteDbError>(),
                        Some(MaterializerStateSqliteDbError::PathDoesNotExist(_path)));
                }
            );
            assert_metadata_matches(db.tables.created_by_table.read_all()?, &metadatas[0]);
            assert_metadata_matches(db.tables.last_read_by_table.read_all()?, &metadatas[0]);

            db.materializer_state_table()
                .insert(&path, &artifact_metadata, timestamp)
                .unwrap();
        }

        {
            let (db, loaded_state) =
                testing_materializer_state_sqlite_db(fs.path(), v0, metadatas[1].clone(), None)
                    .unwrap();
            assert_matches!(
                loaded_state,
                Ok(v) => {
                    assert_eq!(v, vec![(path.clone(), (artifact_metadata.clone(), timestamp))]);
                }
            );
            assert_metadata_matches(db.tables.created_by_table.read_all()?, &metadatas[0]);
            assert_metadata_matches(db.tables.last_read_by_table.read_all()?, &metadatas[1]);
        }

        {
            let (mut db, loaded_state) = testing_materializer_state_sqlite_db(
                fs.path(),
                v1.clone(),
                metadatas[2].clone(),
                None,
            )
            .unwrap();
            assert_matches!(
                loaded_state,
                Err(e) => {
                    assert_matches!(
                        e.downcast_ref::<MaterializerStateSqliteDbError>(),
                        Some(MaterializerStateSqliteDbError::VersionMismatch {
                            expected: _expected,
                            found: _found,
                            path: _path,
                    }));
                }
            );
            assert_metadata_matches(db.tables.created_by_table.read_all()?, &metadatas[2]);
            assert_metadata_matches(db.tables.last_read_by_table.read_all()?, &metadatas[2]);

            db.materializer_state_table()
                .insert(&path, &artifact_metadata, timestamp)
                .unwrap();
        }

        let identity = {
            let (db, loaded_state) = testing_materializer_state_sqlite_db(
                fs.path(),
                v1.clone(),
                metadatas[3].clone(),
                None,
            )
            .unwrap();
            assert_matches!(
                loaded_state,
                Ok(v) => {
                    assert_eq!(v, vec![(path, (artifact_metadata, timestamp))]);
                }
            );
            assert_metadata_matches(db.tables.created_by_table.read_all()?, &metadatas[2]);
            assert_metadata_matches(db.tables.last_read_by_table.read_all()?, &metadatas[3]);

            db.identity
        };

        {
            let (db, loaded_state) = testing_materializer_state_sqlite_db(
                fs.path(),
                v1,
                metadatas[4].clone(),
                Some(&identity),
            )
            .unwrap();
            assert_matches!(
                loaded_state,
                Err(e) => {
                    assert_matches!(
                        e.downcast_ref::<MaterializerStateSqliteDbError>(),
                        Some(MaterializerStateSqliteDbError::RejectedIdentity {
                            ..
                    }));
                }
            );
            assert_metadata_matches(db.tables.created_by_table.read_all()?, &metadatas[4]);
            assert_metadata_matches(db.tables.last_read_by_table.read_all()?, &metadatas[4]);
        }

        Ok(())
    }

    #[test]
    fn test_delete_many() -> anyhow::Result<()> {
        let conn = Connection::open_in_memory()?;

        let table = MaterializerStateSqliteTable::new(Arc::new(Mutex::new(conn)));
        table.create_table()?;

        // Sqlite has limits on how many variables you can use at once. Check we don't run into
        // those.
        let paths = (0..50000)
            .map(|i| ProjectRelativePathBuf::unchecked_new(format!("foo/{}", i)))
            .collect();

        table.delete(paths)?;

        Ok(())
    }
}