/*
 * 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::borrow::Cow;
use std::collections::HashMap;
use std::collections::HashSet;

use buck2_common::file_ops::FileMetadata;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::insert_file;
use buck2_execute::directory::ActionDirectoryBuilder;
use buck2_execute::materialize::materializer::DeferredMaterializerSubscription;
use dupe::Dupe;

use super::Version;
use super::VersionTracker;
use super::*;

#[test]
fn test_find_artifacts() -> anyhow::Result<()> {
    let artifact1 = ProjectRelativePathBuf::unchecked_new("foo/bar/baz".to_owned());
    let artifact2 = ProjectRelativePathBuf::unchecked_new("foo/bar/bar/qux".to_owned());
    let artifact3 = ProjectRelativePathBuf::unchecked_new("foo/bar/bar/quux".to_owned());
    let artifact4 = ProjectRelativePathBuf::unchecked_new("foo/bar/qux/quuz".to_owned());
    let non_artifact1 = ProjectRelativePathBuf::unchecked_new("foo/bar/qux".to_owned());
    let non_artifact2 = ProjectRelativePathBuf::unchecked_new("foo/bar/bar/corge".to_owned());

    let file = FileMetadata::empty(DigestConfig::testing_default().cas_digest_config());

    // Build deps with artifacts 1-3, and non-artifacts 1-2
    let mut builder = ActionDirectoryBuilder::empty();
    insert_file(&mut builder, &artifact1.join_normalized("f1")?, file.dupe())?;
    insert_file(
        &mut builder,
        &artifact2.join_normalized("d/f1")?,
        file.dupe(),
    )?;
    insert_file(&mut builder, &artifact3, file.dupe())?;
    insert_file(&mut builder, &non_artifact2, file.dupe())?;
    builder.mkdir(&non_artifact1)?;

    // Build tree with artifacts 1-4
    let mut tree: FileTree<()> = FileTree::new();
    tree.insert(artifact1.iter().map(|f| f.to_owned()), ());
    tree.insert(artifact2.iter().map(|f| f.to_owned()), ());
    tree.insert(artifact3.iter().map(|f| f.to_owned()), ());
    tree.insert(artifact4.iter().map(|f| f.to_owned()), ());

    let expected_artifacts: HashSet<_> =
        vec![artifact1, artifact2, artifact3].into_iter().collect();
    let found_artifacts: HashSet<_> = tree.find_artifacts(&builder).into_iter().collect();
    assert_eq!(found_artifacts, expected_artifacts);
    Ok(())
}

#[test]
fn test_remove_path() {
    fn insert(tree: &mut FileTree<String>, path: &str) {
        tree.insert(
            ProjectRelativePath::unchecked_new(path)
                .iter()
                .map(|f| f.to_owned()),
            path.to_owned(),
        );
    }

    let mut tree: FileTree<String> = FileTree::new();
    insert(&mut tree, "a/b/c/d");
    insert(&mut tree, "a/b/c/e");
    insert(&mut tree, "a/c");

    let removed_subtree = tree.remove_path(ProjectRelativePath::unchecked_new("a/b"));
    // Convert to HashMap<String, String> so it's easier to test
    let removed_subtree: HashMap<String, String> = removed_subtree
        .map(|(k, v)| (k.as_str().to_owned(), v))
        .collect();

    assert_eq!(removed_subtree.len(), 2);
    assert_eq!(removed_subtree.get("a/b/c/d"), Some(&"a/b/c/d".to_owned()));
    assert_eq!(removed_subtree.get("a/b/c/e"), Some(&"a/b/c/e".to_owned()));
}

mod state_machine {
    use std::path::Path;

    use assert_matches::assert_matches;
    use buck2_execute::directory::Symlink;
    use buck2_execute::directory::INTERNER;
    use parking_lot::Mutex;
    use tokio::time::sleep;
    use tokio::time::Duration as TokioDuration;

    use super::*;

    #[derive(Debug, Eq, PartialEq)]
    enum Op {
        Clean,
        Materialize,
    }

    #[derive(Default)]
    struct StubIoHandler {
        log: Mutex<Vec<(Op, ProjectRelativePathBuf)>>,
        fail: Mutex<bool>,
        // If set, add a sleep when materializing to simulate a long materialization period
        materialization_config: HashMap<ProjectRelativePathBuf, TokioDuration>,
    }

    impl StubIoHandler {
        fn take_log(&self) -> Vec<(Op, ProjectRelativePathBuf)> {
            std::mem::take(&mut *self.log.lock())
        }

        fn set_fail(&self, fail: bool) {
            *self.fail.lock() = fail;
        }

        pub fn new(materialization_config: HashMap<ProjectRelativePathBuf, TokioDuration>) -> Self {
            Self {
                log: Default::default(),
                fail: Default::default(),
                materialization_config,
            }
        }
    }

    #[async_trait]
    impl IoHandler for StubIoHandler {
        fn write<'a>(
            self: &Arc<Self>,
            _path: ProjectRelativePathBuf,
            _write: Arc<WriteFile>,
            _version: Version,
            _command_sender: MaterializerSender<Self>,
            _cancellations: &'a CancellationContext<'a>,
        ) -> BoxFuture<'a, Result<(), SharedMaterializingError>> {
            unimplemented!()
        }

        fn clean_path<'a>(
            self: &Arc<Self>,
            path: ProjectRelativePathBuf,
            version: Version,
            command_sender: MaterializerSender<Self>,
            _cancellations: &'a CancellationContext,
        ) -> BoxFuture<'a, Result<(), SharedError>> {
            self.log.lock().push((Op::Clean, path.clone()));

            async move {
                let _ignored = command_sender.send_low_priority(
                    LowPriorityMaterializerCommand::CleanupFinished {
                        path,
                        version,
                        result: Ok(()),
                    },
                );
                Ok(())
            }
            .boxed()
        }

        async fn materialize_entry(
            self: &Arc<Self>,
            path: ProjectRelativePathBuf,
            _method: Arc<ArtifactMaterializationMethod>,
            _entry: ActionDirectoryEntry<ActionSharedDirectory>,
            _event_dispatcher: EventDispatcher,
            _cancellations: &CancellationContext,
        ) -> Result<(), MaterializeEntryError> {
            // Simulate a non-immediate materialization if configured
            match self.materialization_config.get(&path) {
                Some(duration) => {
                    sleep(*duration).await;
                }
                None => (),
            }
            self.log.lock().push((Op::Materialize, path));

            if *self.fail.lock() {
                Err(anyhow::anyhow!("Injected error").into())
            } else {
                Ok(())
            }
        }

        fn create_ttl_refresh(
            self: &Arc<Self>,
            _tree: &ArtifactTree,
            _min_ttl: Duration,
            _digest_config: DigestConfig,
        ) -> Option<BoxFuture<'static, anyhow::Result<()>>> {
            unimplemented!()
        }
    }

    /// A stub command sender. We are calling materializer methods directly so that's all we need.
    fn channel() -> (
        MaterializerSender<StubIoHandler>,
        MaterializerReceiver<StubIoHandler>,
    ) {
        // We don't use those counts in tests.
        static SENT: AtomicUsize = AtomicUsize::new(0);
        static RECEIVED: AtomicUsize = AtomicUsize::new(0);

        let (hi_send, hi_recv) = mpsc::unbounded_channel();
        let (lo_send, lo_recv) = mpsc::unbounded_channel();
        let counters = MaterializerCounters {
            sent: &SENT,
            received: &RECEIVED,
        };

        (
            MaterializerSender {
                high_priority: Cow::Owned(hi_send),
                low_priority: Cow::Owned(lo_send),
                counters,
            },
            MaterializerReceiver {
                high_priority: hi_recv,
                low_priority: lo_recv,
                counters,
            },
        )
    }

    fn make_path(p: &str) -> ProjectRelativePathBuf {
        ProjectRelativePath::new(p).unwrap().to_owned()
    }

    fn make_processor(
        digest_config: DigestConfig,
        materialization_config: HashMap<ProjectRelativePathBuf, TokioDuration>,
    ) -> (
        DeferredMaterializerCommandProcessor<StubIoHandler>,
        MaterializerReceiver<StubIoHandler>,
    ) {
        let (command_sender, command_receiver) = channel();

        (
            DeferredMaterializerCommandProcessor {
                io: Arc::new(StubIoHandler::new(materialization_config)),
                sqlite_db: None,
                rt: Handle::current(),
                defer_write_actions: true,
                log_buffer: LogBuffer::new(1),
                digest_config,
                version_tracker: VersionTracker::new(),
                command_sender,
                tree: ArtifactTree::new(),
                subscriptions: MaterializerSubscriptions::new(),
                ttl_refresh_history: Default::default(),
                ttl_refresh_instance: Default::default(),
                cancellations: CancellationContext::testing(),
                stats: Arc::new(DeferredMaterializerStats::default()),
            },
            command_receiver,
        )
    }

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

        let (mut dm, _) = make_processor(digest_config, Default::default());

        let path = make_path("foo/bar");
        let value = ArtifactValue::file(digest_config.empty_file());

        dm.declare(
            &path,
            value.dupe(),
            Box::new(ArtifactMaterializationMethod::Test),
        );
        assert_eq!(dm.io.take_log(), &[(Op::Clean, path.clone())]);

        let res = dm
            .materialize_artifact(&path, EventDispatcher::null())
            .context("Expected a future")?
            .await;
        assert_eq!(dm.io.take_log(), &[(Op::Materialize, path.clone())]);

        dm.materialization_finished(path.clone(), Utc::now(), dm.version_tracker.current(), res);
        assert_eq!(dm.io.take_log(), &[]);

        // When redeclaring the same artifact nothing happens.
        dm.declare(
            &path,
            value.dupe(),
            Box::new(ArtifactMaterializationMethod::Test),
        );
        assert_eq!(dm.io.take_log(), &[]);

        // When declaring the same artifact but under it, we clean it and it's a new artifact.
        let path2 = make_path("foo/bar/baz");
        dm.declare(
            &path2,
            value.dupe(),
            Box::new(ArtifactMaterializationMethod::Test),
        );
        assert_eq!(dm.io.take_log(), &[(Op::Clean, path2.clone())]);

        let _ignore = dm
            .materialize_artifact(&path2, EventDispatcher::null())
            .context("Expected a future")?
            .await;
        assert_eq!(dm.io.take_log(), &[(Op::Materialize, path2.clone())]);

        Ok(())
    }

    fn make_artifact_value_with_symlink_dep(
        target_path: &ProjectRelativePathBuf,
        target_from_symlink: &RelativePathBuf,
        digest_config: DigestConfig,
    ) -> anyhow::Result<ArtifactValue> {
        let mut deps = ActionDirectoryBuilder::empty();
        let target = ActionDirectoryEntry::Leaf(ActionDirectoryMember::File(FileMetadata::empty(
            digest_config.cas_digest_config(),
        )));
        deps.insert(target_path.as_forward_relative_path(), target)?;
        let symlink_value = ArtifactValue::new(
            ActionDirectoryEntry::Leaf(ActionDirectoryMember::Symlink(Arc::new(Symlink::new(
                target_from_symlink.clone(),
            )))),
            Some(
                deps.fingerprint(digest_config.as_directory_serializer())
                    .shared(&*INTERNER),
            ),
        );
        Ok(symlink_value)
    }

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

        // Construct a tree with a symlink and its target, materialize both at once
        let symlink_path = make_path("foo/bar_symlink");
        let target_path = make_path("foo/bar_target");
        let target_from_symlink = RelativePathBuf::from_path(Path::new("bar_target"))?;

        let mut materialization_config = HashMap::new();
        // Materialize the symlink target slowly so that we actually hit the logic point where we
        // await for symlink targets and the entry materialization
        materialization_config.insert(target_path.clone(), TokioDuration::from_millis(100));

        let (mut dm, _) = make_processor(digest_config, materialization_config);

        // Declare symlink target
        dm.declare(
            &target_path,
            ArtifactValue::file(digest_config.empty_file()),
            Box::new(ArtifactMaterializationMethod::Test),
        );
        assert_eq!(dm.io.take_log(), &[(Op::Clean, target_path.clone())]);

        // Declare symlink
        let symlink_value = make_artifact_value_with_symlink_dep(
            &target_path,
            &target_from_symlink,
            digest_config,
        )?;
        dm.declare(
            &symlink_path,
            symlink_value,
            Box::new(ArtifactMaterializationMethod::Test),
        );
        assert_eq!(dm.io.take_log(), &[(Op::Clean, symlink_path.clone())]);

        dm.materialize_artifact(&symlink_path, EventDispatcher::null())
            .context("Expected a future")?
            .await
            .map_err(|_| anyhow::anyhow!("error materializing"))?;

        let logs = dm.io.take_log();
        if cfg!(unix) {
            assert_eq!(
                logs,
                &[
                    (Op::Materialize, symlink_path.clone()),
                    (Op::Materialize, target_path.clone())
                ]
            );
        } else {
            assert_eq!(
                logs,
                &[
                    (Op::Materialize, target_path.clone()),
                    (Op::Materialize, symlink_path.clone())
                ]
            );
        }
        Ok(())
    }

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

        // Materialize a symlink, then materialize the target. Test that we still
        // materialize deps if the main artifact has already been materialized.
        let symlink_path = make_path("foo/bar_symlink");
        let target_path = make_path("foo/bar_target");
        let target_from_symlink = RelativePathBuf::from_path(Path::new("bar_target"))?;

        let mut materialization_config = HashMap::new();
        // Materialize the symlink target slowly so that we actually hit the logic point where we
        // await for symlink targets and the entry materialization
        materialization_config.insert(target_path.clone(), TokioDuration::from_millis(100));

        let (mut dm, _) = make_processor(digest_config, materialization_config);

        // Declare symlink
        let symlink_value = make_artifact_value_with_symlink_dep(
            &target_path,
            &target_from_symlink,
            digest_config,
        )?;
        dm.declare(
            &symlink_path,
            symlink_value,
            Box::new(ArtifactMaterializationMethod::Test),
        );
        assert_eq!(dm.io.take_log(), &[(Op::Clean, symlink_path.clone())]);

        // Materialize the symlink, at this point the target is not in the tree so it's ignored
        let res = dm
            .materialize_artifact(&symlink_path, EventDispatcher::null())
            .context("Expected a future")?
            .await;

        let logs = dm.io.take_log();
        assert_eq!(logs, &[(Op::Materialize, symlink_path.clone())]);

        // Mark the symlink as materialized
        dm.materialization_finished(
            symlink_path.clone(),
            Utc::now(),
            dm.version_tracker.current(),
            res,
        );
        assert_eq!(dm.io.take_log(), &[]);

        // Declare symlink target
        dm.declare(
            &target_path,
            ArtifactValue::file(digest_config.empty_file()),
            Box::new(ArtifactMaterializationMethod::Test),
        );
        assert_eq!(dm.io.take_log(), &[(Op::Clean, target_path.clone())]);

        // Materialize the symlink again.
        // This time, we don't re-materialize the symlink as that's already been done.
        // But we still materialize the target as that has not been materialized yet.
        dm.materialize_artifact(&symlink_path, EventDispatcher::null())
            .context("Expected a future")?
            .await
            .map_err(|_| anyhow::anyhow!("error materializing"))?;

        let logs = dm.io.take_log();
        assert_eq!(logs, &[(Op::Materialize, target_path.clone())]);

        Ok(())
    }

    #[tokio::test]
    async fn test_subscription_create_destroy() {
        let digest_config = DigestConfig::testing_default();

        let (mut dm, mut channel) = make_processor(digest_config, Default::default());

        let handle = {
            let (sender, recv) = oneshot::channel();
            MaterializerSubscriptionOperation::Create { sender }.execute(&mut dm);
            recv.await.unwrap()
        };

        assert!(dm.subscriptions.has_subscription(&handle));

        drop(handle);

        while let Ok(cmd) = channel.high_priority.try_recv() {
            dm.process_one_command(cmd);
        }

        assert!(!dm.subscriptions.has_any_subscriptions());
    }

    #[tokio::test]
    async fn test_subscription_notifications() {
        let digest_config = DigestConfig::testing_default();
        let value = ArtifactValue::file(digest_config.empty_file());

        let (mut dm, mut channel) = make_processor(digest_config, Default::default());

        let mut handle = {
            let (sender, recv) = oneshot::channel();
            MaterializerSubscriptionOperation::Create { sender }.execute(&mut dm);
            recv.await.unwrap()
        };

        let foo_bar = make_path("foo/bar");
        let foo_bar_baz = make_path("foo/bar/baz");
        let bar = make_path("bar");
        let qux = make_path("qux");

        dm.declare_existing(&foo_bar, value.dupe());

        handle.subscribe_to_paths(vec![foo_bar_baz.clone(), bar.clone()]);
        while let Ok(cmd) = channel.high_priority.try_recv() {
            dm.process_one_command(cmd);
        }

        dm.declare_existing(&bar, value.dupe());
        dm.declare_existing(&foo_bar_baz, value.dupe());
        dm.declare_existing(&qux, value.dupe());

        let mut paths = Vec::new();
        while let Ok(path) = handle.receiver().try_recv() {
            paths.push(path);
        }

        assert_eq!(paths, vec![foo_bar_baz.clone(), bar, foo_bar_baz]);
    }

    #[tokio::test]
    async fn test_subscription_unsubscribe() {
        let digest_config = DigestConfig::testing_default();
        let value1 = ArtifactValue::file(digest_config.empty_file());
        let value2 = ArtifactValue::dir(digest_config.empty_directory());

        let (mut dm, mut channel) = make_processor(digest_config, Default::default());

        let mut handle = {
            let (sender, recv) = oneshot::channel();
            MaterializerSubscriptionOperation::Create { sender }.execute(&mut dm);
            recv.await.unwrap()
        };

        let path = make_path("foo/bar");

        handle.subscribe_to_paths(vec![path.clone()]);
        while let Ok(cmd) = channel.high_priority.try_recv() {
            dm.process_one_command(cmd);
        }

        dm.declare_existing(&path, value1.dupe());

        handle.unsubscribe_from_paths(vec![path.clone()]);
        while let Ok(cmd) = channel.high_priority.try_recv() {
            dm.process_one_command(cmd);
        }

        dm.declare_existing(&path, value2.dupe());

        let mut paths = Vec::new();
        while let Ok(path) = handle.receiver().try_recv() {
            paths.push(path);
        }

        // Expect only one notification
        assert_eq!(paths, vec![path]);
    }

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

        let (mut dm, _) = make_processor(digest_config, Default::default());

        let path = make_path("test/invalidate/failure");
        let value1 = ArtifactValue::file(digest_config.empty_file());
        let value2 = ArtifactValue::dir(digest_config.empty_directory());

        // Start from having something.
        dm.declare_existing(&path, value1);

        // This will collect the existing future and invalidate, and then fail in doing so.
        dm.declare(&path, value2, Box::new(ArtifactMaterializationMethod::Test));

        // Now we check that materialization fails. This needs to wait on the previous clean.
        let res = dm
            .materialize_artifact(&path, EventDispatcher::null())
            .context("Expected a future")?
            .await;

        assert_matches!(
            res,
            Err(SharedMaterializingError::Error(e)) if format!("{:#}", e).contains("Injected error")
        );

        // We do not actually get to materializing or cleaning.
        assert_eq!(dm.io.take_log(), &[]);

        Ok(())
    }

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

        let (mut dm, mut channel) = make_processor(digest_config, Default::default());

        let path = make_path("test");
        let value1 = ArtifactValue::file(digest_config.empty_file());

        // Declare a value.
        dm.declare(&path, value1, Box::new(ArtifactMaterializationMethod::Test));

        // Make materializations fail
        dm.io.set_fail(true);

        // Materializing it fails.
        let res = dm
            .materialize_artifact(&path, EventDispatcher::null())
            .context("Expected a future")?
            .await;

        assert_matches!(
            res,
            Err(SharedMaterializingError::Error(e)) if format!("{:#}", e).contains("Injected error")
        );

        // Unset fail, but we haven't processed materialization_finished yet so this does nothing.
        dm.io.set_fail(false);

        // Rejoining the existing future fails.
        let res = dm
            .materialize_artifact(&path, EventDispatcher::null())
            .context("Expected a future")?
            .await;

        assert_matches!(
            res,
            Err(SharedMaterializingError::Error(e)) if format!("{:#}", e).contains("Injected error")
        );

        // Now process cleanup_finished_vacant and materialization_finished.
        let mut processed = 0;

        while let Ok(cmd) = channel.low_priority.try_recv() {
            eprintln!("got cmd = {:?}", cmd);
            dm.process_one_low_priority_command(cmd);
            processed += 1;
        }

        assert_eq!(processed, 2);

        // Materializing works now:
        let res = dm
            .materialize_artifact(&path, EventDispatcher::null())
            .context("Expected a future")?
            .await;

        assert_matches!(res, Ok(()));

        Ok(())
    }
}


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

use std::convert::Infallible;
use std::ops::ControlFlow;
use std::ops::FromResidual;
use std::sync::Arc;

use anyhow::Context as _;
use buck2_common::executor_config::RemoteExecutorUseCase;
use buck2_common::file_ops::FileDigest;
use buck2_common::file_ops::FileMetadata;
use buck2_common::file_ops::TrackedFileDigest;
use buck2_core::directory::DirectoryEntry;
use buck2_core::env_helper::EnvHelper;
use buck2_core::fs::paths::forward_rel_path::ForwardRelativePath;
use buck2_core::fs::project_rel_path::ProjectRelativePathBuf;
use buck2_execute::artifact_value::ArtifactValue;
use buck2_execute::digest::CasDigestFromReExt;
use buck2_execute::digest_config::DigestConfig;
use buck2_execute::directory::extract_artifact_value;
use buck2_execute::directory::re_tree_to_directory;
use buck2_execute::directory::ActionDirectoryMember;
use buck2_execute::execute::action_digest::ActionDigest;
use buck2_execute::execute::action_digest::TrackedActionDigest;
use buck2_execute::execute::executor_stage_async;
use buck2_execute::execute::manager::CommandExecutionManager;
use buck2_execute::execute::manager::CommandExecutionManagerExt;
use buck2_execute::execute::manager::CommandExecutionManagerWithClaim;
use buck2_execute::execute::output::CommandStdStreams;
use buck2_execute::execute::request::CommandExecutionOutput;
use buck2_execute::execute::request::CommandExecutionOutputRef;
use buck2_execute::execute::request::CommandExecutionPaths;
use buck2_execute::execute::request::CommandExecutionRequest;
use buck2_execute::execute::result::CommandExecutionResult;
use buck2_execute::materialize::materializer::CasDownloadInfo;
use buck2_execute::materialize::materializer::Materializer;
use buck2_execute::re::manager::ManagedRemoteExecutionClient;
use buck2_execute::re::remote_action_result::RemoteActionResult;
use chrono::DateTime;
use chrono::Duration;
use chrono::Utc;
use dupe::Dupe;
use futures::future;
use futures::FutureExt;
use gazebo::prelude::*;
use indexmap::IndexMap;
use more_futures::cancellation::CancellationContext;
use remote_execution as RE;
use thiserror::Error;

pub async fn download_action_results<'a>(
    request: &CommandExecutionRequest,
    materializer: &dyn Materializer,
    re_client: &ManagedRemoteExecutionClient,
    re_use_case: RemoteExecutorUseCase,
    digest_config: DigestConfig,
    manager: CommandExecutionManager,
    stage: buck2_data::executor_stage_start::Stage,
    paths: &CommandExecutionPaths,
    requested_outputs: impl Iterator<Item = CommandExecutionOutputRef<'a>>,
    action_digest: &ActionDigest,
    response: &dyn RemoteActionResult,
    cancellations: &CancellationContext<'_>,
) -> DownloadResult {
    let downloader = CasDownloader {
        materializer,
        re_client,
        re_use_case,
        digest_config,
    };

    let download = downloader.download(
        manager,
        stage,
        paths,
        requested_outputs,
        action_digest,
        response,
        cancellations,
    );

    let std_streams = response.std_streams(re_client, re_use_case, digest_config);
    let std_streams = async {
        if request.prefetch_lossy_stderr() {
            std_streams.prefetch_lossy_stderr().await
        } else {
            std_streams
        }
    };

    let (download, std_streams) = future::join(download, std_streams).await;
    let (manager, outputs) = download?;

    DownloadResult::Result(manager.success(
        response.execution_kind(action_digest.dupe()),
        outputs,
        CommandStdStreams::Remote(std_streams),
        response.timing(),
    ))
}

pub struct CasDownloader<'a> {
    pub materializer: &'a dyn Materializer,
    pub re_client: &'a ManagedRemoteExecutionClient,
    pub re_use_case: RemoteExecutorUseCase,
    pub digest_config: DigestConfig,
}

impl CasDownloader<'_> {
    async fn download<'a>(
        &self,
        manager: CommandExecutionManager,
        stage: buck2_data::executor_stage_start::Stage,
        paths: &CommandExecutionPaths,
        requested_outputs: impl Iterator<Item = CommandExecutionOutputRef<'a>>,
        action_digest: &ActionDigest,
        output_spec: &dyn RemoteActionResult,
        cancellations: &CancellationContext<'_>,
    ) -> ControlFlow<
        DownloadResult,
        (
            CommandExecutionManagerWithClaim,
            IndexMap<CommandExecutionOutput, ArtifactValue>,
        ),
    > {
        executor_stage_async(stage, async {
            let artifacts = self
                .extract_artifacts(paths, requested_outputs, output_spec)
                .await;

            let artifacts = match artifacts {
                Ok(artifacts) => artifacts,
                Err(e) => {
                    return ControlFlow::Break(DownloadResult::Result(manager.error(
                        "extract_artifacts",
                        e.context(format!("action_digest={}", action_digest)),
                    )));
                }
            };

            // Claim the request before starting the download.
            let manager = manager.claim().await;

            let outputs = self
                .materialize_outputs(artifacts, action_digest, cancellations)
                .await;

            let outputs = match outputs {
                Ok(outputs) => outputs,
                Err(e) => {
                    return ControlFlow::Break(DownloadResult::Result(manager.error(
                        "materialize_outputs",
                        e.context(format!("action_digest={}", action_digest)),
                    )));
                }
            };

            ControlFlow::Continue((manager, outputs))
        })
        .await
    }

    async fn extract_artifacts<'a>(
        &self,
        paths: &CommandExecutionPaths,
        requested_outputs: impl Iterator<Item = CommandExecutionOutputRef<'a>>,
        output_spec: &dyn RemoteActionResult,
    ) -> anyhow::Result<ExtractedArtifacts> {
        static FAIL_RE_DOWNLOADS: EnvHelper<bool> = EnvHelper::new("BUCK2_TEST_FAIL_RE_DOWNLOADS");
        if FAIL_RE_DOWNLOADS.get()?.copied().unwrap_or_default() {
            return Err(anyhow::anyhow!("Injected error"));
        }

        let now = Utc::now();
        let ttl = Duration::seconds(output_spec.ttl());
        let expires = now + ttl;

        // Download process:
        // 1. merges all the outputs (files and trees) into the inputs structure
        // 2. computes the ArtifactValue for all outputs from that merged structure
        // 3. pass those new ArtifactValue to the materializer
        let input_dir = paths.input_directory();
        let output_paths = paths.output_paths();
        let mut input_dir = input_dir.clone().into_builder();

        for x in output_spec.output_files() {
            let digest = FileDigest::from_re(&x.digest.digest, self.digest_config)?;
            let digest = TrackedFileDigest::new_expires(
                digest,
                expires,
                self.digest_config.cas_digest_config(),
            );

            let entry = DirectoryEntry::Leaf(ActionDirectoryMember::File(FileMetadata {
                digest,
                is_executable: x.executable,
            }));

            input_dir.insert(re_forward_path(x.name.as_str())?, entry)?;
        }

        // Compute the re_outputs from the output_directories
        // This requires traversing the trees to find symlinks that point outside such trees
        let trees = self
            .re_client
            .download_typed_blobs::<RE::Tree>(
                output_spec
                    .output_directories()
                    .map(|x| x.tree_digest.clone()),
                self.re_use_case,
            )
            .boxed()
            .await
            .context(DownloadError::DownloadTrees)?;

        for (dir, tree) in output_spec.output_directories().iter().zip(trees) {
            let entry = re_tree_to_directory(&tree, &expires, self.digest_config)?;
            input_dir.insert(
                re_forward_path(dir.path.as_str())?,
                DirectoryEntry::Dir(entry),
            )?;
        }

        let mut to_declare = Vec::with_capacity(output_paths.len());
        let mut mapped_outputs = IndexMap::with_capacity(output_paths.len());

        for (requested, (path, _)) in requested_outputs.zip(output_paths.iter()) {
            let value = extract_artifact_value(&input_dir, path, self.digest_config)?;
            if let Some(value) = value {
                to_declare.push((path.to_owned(), value.dupe()));
                mapped_outputs.insert(requested.cloned(), value);
            }
        }

        Ok(ExtractedArtifacts {
            to_declare,
            mapped_outputs,
            now,
            expires,
            ttl,
        })
    }

    async fn materialize_outputs<'a>(
        &self,
        artifacts: ExtractedArtifacts,
        action_digest: &ActionDigest,
        cancellations: &CancellationContext<'_>,
    ) -> anyhow::Result<IndexMap<CommandExecutionOutput, ArtifactValue>> {
        // Declare the outputs to the materializer
        self.materializer
            .declare_cas_many(
                Arc::new(CasDownloadInfo::new_execution(
                    TrackedActionDigest::new_expires(
                        action_digest.dupe(),
                        artifacts.expires,
                        self.digest_config.cas_digest_config(),
                    ),
                    self.re_use_case,
                    artifacts.now,
                    artifacts.ttl,
                )),
                artifacts.to_declare,
                cancellations,
            )
            .boxed()
            .await
            .context(DownloadError::Materialization)?;

        Ok(artifacts.mapped_outputs)
    }
}

/// Takes a path that came from RE and tries to convert it to
/// a `ForwardRelativePath`. These paths are supposed to be forward relative,
/// so if the conversion fails, RE is broken.
fn re_forward_path(re_path: &str) -> anyhow::Result<&ForwardRelativePath> {
    // RE sends us paths with trailing slash.
    ForwardRelativePath::new_trim_trailing_slashes(re_path)
        .context(DownloadError::InvalidPathFromRe)
}

#[derive(Error, Debug)]
enum DownloadError {
    #[error("Failed to declare in materializer")]
    Materialization,

    #[error("Failed to download trees")]
    DownloadTrees,

    #[error("Path received from RE is not normalized.")]
    InvalidPathFromRe,
}

struct ExtractedArtifacts {
    to_declare: Vec<(ProjectRelativePathBuf, ArtifactValue)>,
    mapped_outputs: IndexMap<CommandExecutionOutput, ArtifactValue>,
    now: DateTime<Utc>,
    expires: DateTime<Utc>,
    ttl: Duration,
}

/// Did this download work out?
pub enum DownloadResult {
    /// Got a result: might be a success, might be a failure. Caller needs to deal with this
    /// result.
    Result(CommandExecutionResult),
}

impl FromResidual<ControlFlow<Self, Infallible>> for DownloadResult {
    fn from_residual(residual: ControlFlow<Self, Infallible>) -> Self {
        match residual {
            ControlFlow::Break(v) => v,
            ControlFlow::Continue(_) => unreachable!(),
        }
    }
}


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

pub mod download;


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

use buck2_util::late_binding::LateBinding;

pub static FLUSH_DEP_FILES: LateBinding<fn()> = LateBinding::new("FLUSH_DEP_FILES");

/// Forget about all dep files. This isn't really meant to be commonly used, but if an invalid dep
/// file was produced and the user wants unblocking, this will provide it.
pub fn flush_dep_files() {
    (FLUSH_DEP_FILES.get().unwrap())();
}


/*
 * 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 async_trait::async_trait;
use buck2_common::ignores::ignore_set::IgnoreSet;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_core::cells::name::CellName;
use buck2_core::cells::CellResolver;
use buck2_core::fs::project::ProjectRoot;
use buck2_core::is_open_source;
use dice::DiceTransactionUpdater;

use crate::mergebase::Mergebase;
use crate::notify::NotifyFileWatcher;
use crate::watchman::interface::WatchmanFileWatcher;

#[async_trait]
pub trait FileWatcher: Allocative + Send + Sync + 'static {
    async fn sync(
        &self,
        dice: DiceTransactionUpdater,
    ) -> anyhow::Result<(DiceTransactionUpdater, Mergebase)>;
}

impl dyn FileWatcher {
    /// Create a new FileWatcher. Note that this is not async, since it's called during daemon
    /// startup and shouldn't be doing any work that could warrant suspending.
    pub fn new(
        project_root: &ProjectRoot,
        root_config: &LegacyBuckConfig,
        cells: CellResolver,
        ignore_specs: HashMap<CellName, IgnoreSet>,
    ) -> anyhow::Result<Arc<dyn FileWatcher>> {
        let default = if is_open_source() {
            "notify"
        } else {
            "watchman"
        };

        match root_config.get("buck2", "file_watcher").unwrap_or(default) {
            "watchman" => Ok(Arc::new(
                WatchmanFileWatcher::new(project_root.root(), root_config, cells, ignore_specs)
                    .context("Creating watchman file watcher")?,
            )),
            "notify" => Ok(Arc::new(
                NotifyFileWatcher::new(project_root, cells, ignore_specs)
                    .context("Creating notify file watcher")?,
            )),
            other => Err(anyhow::anyhow!("Invalid buck2.file_watcher: {}", other)),
        }
    }
}


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

pub mod dep_files;
pub mod file_watcher;
pub mod mergebase;
mod notify;
mod stats;
mod watchman;


/*
 * 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 dice::UserComputationData;
use dupe::Dupe;

#[derive(Clone, Default, Dupe)]
pub struct Mergebase(pub Arc<Option<String>>);

pub trait SetMergebase {
    fn set_mergebase(&mut self, mergebase: Mergebase);
}

pub trait GetMergebase {
    fn get_mergebase(&self) -> Mergebase;
}

impl SetMergebase for UserComputationData {
    fn set_mergebase(&mut self, mergebase: Mergebase) {
        self.data.set(mergebase);
    }
}

impl GetMergebase for UserComputationData {
    fn get_mergebase(&self) -> Mergebase {
        self.data
            .get::<Mergebase>()
            .expect("mergebase should be set")
            .dupe()
    }
}


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

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

use allocative::Allocative;
use async_trait::async_trait;
use buck2_common::dice::file_ops::FileChangeTracker;
use buck2_common::ignores::ignore_set::IgnoreSet;
use buck2_common::invocation_paths::InvocationPaths;
use buck2_core::cells::cell_path::CellPath;
use buck2_core::cells::name::CellName;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::project::ProjectRoot;
use buck2_events::dispatch::span_async;
use buck2_util::collections::ordered_set::OrderedSet;
use dice::DiceTransactionUpdater;
use dupe::Dupe;
use notify::event::CreateKind;
use notify::event::MetadataKind;
use notify::event::ModifyKind;
use notify::event::RemoveKind;
use notify::EventKind;
use notify::RecommendedWatcher;
use notify::Watcher;
use tracing::info;

use crate::file_watcher::FileWatcher;
use crate::mergebase::Mergebase;
use crate::stats::FileWatcherStats;

#[derive(Debug, Clone, Copy, Dupe, PartialEq, Eq, Hash, Allocative)]
enum ChangeType {
    None,
    FileContents,
    FileExistence,
    DirExistence,
    SomeExistence,
    Unknown,
}

impl ChangeType {
    fn new(x: EventKind) -> Self {
        match x {
            EventKind::Access(_) => Self::None,
            EventKind::Create(x) => match x {
                CreateKind::File => Self::FileExistence,
                CreateKind::Folder => Self::DirExistence,
                CreateKind::Any | CreateKind::Other => Self::SomeExistence,
            },
            EventKind::Modify(x) => match x {
                ModifyKind::Data(_) => Self::FileContents,
                ModifyKind::Metadata(x) => match x {
                    MetadataKind::Ownership | MetadataKind::Permissions => Self::FileContents,
                    _ => Self::None,
                },
                ModifyKind::Name(_) => Self::SomeExistence,
                ModifyKind::Any | ModifyKind::Other => Self::Unknown,
            },
            EventKind::Remove(x) => match x {
                RemoveKind::File => Self::FileExistence,
                RemoveKind::Folder => Self::DirExistence,
                RemoveKind::Any | RemoveKind::Other => Self::SomeExistence,
            },
            EventKind::Any | EventKind::Other => Self::Unknown,
        }
    }
}

/// Buffer containing the events that have happened since we last got a message.
/// Used to dedupe events, since notify sends a notification on every change.
#[derive(Allocative)]
struct NotifyFileData {
    ignored: u64,
    events: OrderedSet<(CellPath, ChangeType)>,
}

impl NotifyFileData {
    fn new() -> Self {
        Self {
            ignored: 0,
            events: OrderedSet::new(),
        }
    }

    fn process(
        &mut self,
        event: notify::Result<notify::Event>,
        root: &ProjectRoot,
        cells: &CellResolver,
        ignore_specs: &HashMap<CellName, IgnoreSet>,
    ) -> anyhow::Result<()> {
        let event = event?;
        let change_type = ChangeType::new(event.kind);
        for path in event.paths {
            // Testing shows that we get absolute paths back from the `notify` library.
            // It's not documented though.
            let path = root.relativize(AbsNormPath::new(&path)?)?;

            // We ignore the buck-out prefix, as those are uninteresting events caused by us.
            // We also ignore other buck-out directories, as if you have two isolation dirs running at once, they are not interesting.
            // We do this in the notify-watcher, rather than a generic layer, as watchman users should configure
            // to ignore buck-out, to reduce the number of events, rather than hiding them later.
            if path.starts_with(InvocationPaths::buck_out_dir_prefix()) {
                // We don't want to event add them as ignored events, since they are super common
                // and very boring
                continue;
            }

            let cell_path = cells.get_cell_path(&path)?;
            let ignore = ignore_specs
                .get(&cell_path.cell())
                .expect("unexpected cell name mismatch")
                .is_match(cell_path.path());

            info!(
                "FileWatcher: {:?} {:?} (ignore = {})",
                path, change_type, ignore
            );

            if ignore || change_type == ChangeType::None {
                self.ignored += 1;
            } else {
                self.events.insert((cell_path, change_type));
            }
        }
        Ok(())
    }

    fn sync(self) -> (buck2_data::FileWatcherStats, FileChangeTracker) {
        // The changes that go into the DICE transaction
        let mut changed = FileChangeTracker::new();
        // The files that were changed for accumulating the stats
        let mut changed_paths = OrderedSet::new();

        for (cell_path, change_type) in self.events {
            let cell_path_str = cell_path.to_string();
            match change_type {
                ChangeType::None => {}
                ChangeType::FileContents => changed.file_changed(cell_path),
                ChangeType::FileExistence => changed.file_added_or_removed(cell_path),
                ChangeType::DirExistence => changed.dir_added_or_removed(cell_path),
                ChangeType::SomeExistence | ChangeType::Unknown => {
                    changed.dir_added_or_removed(cell_path.clone());
                    changed.file_added_or_removed(cell_path)
                }
            }
            // We use changed_paths to deduplicate
            changed_paths.insert(cell_path_str);
        }

        let mut stats = FileWatcherStats::new(changed_paths.len(), None, None);
        stats.add_ignored(self.ignored);
        for path in changed_paths {
            // The event type and watcher kind are just made up, but that's not a big deal
            // since we only use this path open source, where we don't log the information to Scuba anyway.
            // The path is right, which is probably what matters most
            stats.add(
                path,
                buck2_data::FileWatcherEventType::Modify,
                buck2_data::FileWatcherKind::File,
            );
        }

        (stats.finish(), changed)
    }
}

#[derive(Allocative)]
pub struct NotifyFileWatcher {
    #[allocative(skip)]
    watcher: RecommendedWatcher,
    data: Arc<Mutex<anyhow::Result<NotifyFileData>>>,
}

impl NotifyFileWatcher {
    pub fn new(
        root: &ProjectRoot,
        cells: CellResolver,
        ignore_specs: HashMap<CellName, IgnoreSet>,
    ) -> anyhow::Result<Self> {
        let data = Arc::new(Mutex::new(Ok(NotifyFileData::new())));
        let data2 = data.dupe();
        let root2 = root.dupe();
        let mut watcher = notify::recommended_watcher(move |event| {
            let mut guard = data2.lock().unwrap();
            if let Ok(state) = &mut *guard {
                if let Err(e) = state.process(event, &root2, &cells, &ignore_specs) {
                    *guard = Err(e);
                }
            }
        })?;
        watcher.watch(root.root().as_path(), notify::RecursiveMode::Recursive)?;
        Ok(Self { watcher, data })
    }

    fn sync2(
        &self,
        mut dice: DiceTransactionUpdater,
    ) -> anyhow::Result<(buck2_data::FileWatcherStats, DiceTransactionUpdater)> {
        let mut guard = self.data.lock().unwrap();
        let old = mem::replace(&mut *guard, Ok(NotifyFileData::new()));
        let (stats, changes) = old?.sync();
        changes.write_to_dice(&mut dice)?;
        Ok((stats, dice))
    }
}

#[async_trait]
impl FileWatcher for NotifyFileWatcher {
    async fn sync(
        &self,
        dice: DiceTransactionUpdater,
    ) -> anyhow::Result<(DiceTransactionUpdater, Mergebase)> {
        span_async(
            buck2_data::FileWatcherStart {
                provider: buck2_data::FileWatcherProvider::RustNotify as i32,
            },
            async {
                let (stats, res) = match self.sync2(dice) {
                    Ok((stats, dice)) => {
                        let mergebase = Mergebase(Arc::new(stats.branched_from_revision.clone()));
                        ((Some(stats)), Ok((dice, mergebase)))
                    }
                    Err(e) => (None, Err(e)),
                };
                (res, buck2_data::FileWatcherEnd { stats })
            },
        )
        .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 allocative::Allocative;

/// We limit the number of file change records so we don't use too much memory
/// or too much space in scribe.
///
/// 100 entries covers everything required for 96% of updates, which seems sufficient.
/// Number needs to be < 850 or it is often bigger than a scribe message.
const MAX_FILE_CHANGE_RECORDS: usize = 100;

#[derive(Allocative)]
pub(crate) struct FileWatcherStats {
    stats: buck2_data::FileWatcherStats,
    // Bounded by MAX_FILE_CHANGE_RECORDS
    changes: Vec<buck2_data::FileWatcherEvent>,
    // Did we not insert things into changes
    changes_missed: bool,
}

impl FileWatcherStats {
    pub(crate) fn new(
        min_count: usize,
        mergebase: Option<&str>,
        watchman_version: Option<String>,
    ) -> Self {
        let stats = buck2_data::FileWatcherStats {
            branched_from_revision: mergebase.map(ToOwned::to_owned),
            watchman_version,
            ..Default::default()
        };

        let changes = Vec::with_capacity(std::cmp::min(MAX_FILE_CHANGE_RECORDS, min_count));

        Self {
            stats,
            changes,
            changes_missed: false,
        }
    }

    /// I have seen an event that I am ignoring
    pub(crate) fn add_ignored(&mut self, count: u64) {
        self.stats.events_total += count;
    }

    /// I have seen an event that I am processing
    pub(crate) fn add(
        &mut self,
        path: String,
        event: buck2_data::FileWatcherEventType,
        kind: buck2_data::FileWatcherKind,
    ) {
        self.stats.events_total += 1;
        self.stats.events_processed += 1;

        if self.changes.len() < MAX_FILE_CHANGE_RECORDS {
            self.changes.push(buck2_data::FileWatcherEvent {
                event: event as i32,
                kind: kind as i32,
                path,
            });
        } else {
            self.changes_missed = true;
        }
    }

    pub(crate) fn finish(self) -> buck2_data::FileWatcherStats {
        let Self {
            mut stats,
            changes,
            changes_missed,
        } = self;

        stats.events = changes;
        if changes_missed {
            let reason = format!(
                "Too many files changed ({}, max {})",
                stats.events_processed, MAX_FILE_CHANGE_RECORDS
            );
            stats.incomplete_events_reason = Some(reason);
        }

        stats
    }
}


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

use std::fmt::Debug;
use std::fmt::Display;
use std::path::Path;
use std::path::PathBuf;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context as _;
use async_trait::async_trait;
use buck2_core::env_helper::EnvHelper;
use dupe::Dupe;
use futures::future::Future;
use serde::Deserialize;
use tokio::sync::mpsc::UnboundedReceiver;
use tokio::sync::mpsc::UnboundedSender;
use tokio::sync::oneshot;
use watchman_client::prelude::*;

// We use the "new" field. This is marked as deprecated, but buck1 uses it and
// I'm unaware of issues due to its use there.
//
// Putting this in it own mod was the best way to scope the allow(deprecated).
#[allow(deprecated)]
mod types {
    use super::*;
    query_result_type! {
        pub(crate) struct BuckQueryResult {
            name: NameField,
            file_type: FileTypeField,
            exists: ExistsField,
            new: NewField,
        }
    }

    impl BuckQueryResult {
        pub fn into_event(self) -> Option<WatchmanEvent> {
            let kind = match *self.file_type {
                FileType::BlockSpecial
                | FileType::CharSpecial
                | FileType::Fifo
                | FileType::Socket
                | FileType::SolarisDoor
                | FileType::Unknown => {
                    return None;
                }
                FileType::Directory => WatchmanKind::Directory,
                FileType::Regular => WatchmanKind::File,
                FileType::Symlink => WatchmanKind::Symlink,
            };

            let event = match (*self.exists, *self.new) {
                (true, true) => WatchmanEventType::Create,
                (false, _) => WatchmanEventType::Delete,
                _ => WatchmanEventType::Modify,
            };

            Some(WatchmanEvent {
                kind,
                event,
                path: self.name.into_inner(),
            })
        }
    }
}

use types::*;

#[derive(Debug)]
pub enum WatchmanEventType {
    Create,
    Modify,
    Delete,
}

#[derive(Debug)]
pub enum WatchmanKind {
    File,
    Directory,
    Symlink,
}

impl WatchmanKind {
    fn symbol(&self) -> &'static str {
        match self {
            WatchmanKind::File => "f",
            WatchmanKind::Directory => "d",
            WatchmanKind::Symlink => "l",
        }
    }
}

#[derive(Debug)]
pub struct WatchmanEvent {
    pub kind: WatchmanKind,
    pub event: WatchmanEventType,
    pub path: PathBuf,
}

impl Display for WatchmanEvent {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(
            f,
            "{:?}{}<{}>",
            &self.event,
            self.kind.symbol(),
            self.path.to_string_lossy()
        )
    }
}

#[derive(Dupe, Clone)]
pub struct WatchmanClient(Arc<(watchman_client::Client, ResolvedRoot)>);

impl Debug for WatchmanClient {
    fn fmt(&self, f: &mut std::fmt::Formatter<'_>) -> std::fmt::Result {
        write!(f, "watchman<{:?}>", self.0.1)
    }
}

impl WatchmanClient {
    pub async fn connect(
        connector: &Connector,
        path: CanonicalPath,
    ) -> anyhow::Result<WatchmanClient> {
        let client = connector.connect().await?;
        let root = client.resolve_root(path).await?;
        Ok(Self(Arc::new((client, root))))
    }

    pub async fn query<
        F: serde::de::DeserializeOwned + std::fmt::Debug + Clone + QueryFieldList,
    >(
        &self,
        query: QueryRequestCommon,
    ) -> anyhow::Result<QueryResult<F>> {
        static WATCHMAN_TIMEOUT: EnvHelper<u64> = EnvHelper::new("BUCK2_WATCHMAN_TIMEOUT");
        let fut = self.client().query(self.root(), query);

        Ok(match WATCHMAN_TIMEOUT.get_copied()? {
            Some(timeout) => tokio::time::timeout(Duration::from_secs(timeout), fut)
                .await
                .context("Watchman request timed out")?,
            None => fut.await,
        }?)
    }

    fn root(&self) -> &ResolvedRoot {
        &self.0.1
    }

    fn client(&self) -> &watchman_client::Client {
        &self.0.0
    }
}

#[async_trait]
pub trait SyncableQueryProcessor: Send + Sync {
    type Output;
    type Payload;

    /// Process a set of filesystem change events.
    async fn process_events(
        &mut self,
        payload: Self::Payload,
        events: Vec<WatchmanEvent>,
        mergebase: &Option<String>,
        watchman_version: Option<String>,
    ) -> anyhow::Result<(Self::Output, Self::Payload)>;

    /// Indicates that all derived data should be invalidated. This could happen, for example, if the watchman server restarts.
    async fn on_fresh_instance(
        &mut self,
        dice: Self::Payload,
        mergebase: &Option<String>,
        watchman_version: Option<String>,
    ) -> anyhow::Result<(Self::Output, Self::Payload)>;
}

/// commands to be sent to the SyncableQueryHandler.
enum SyncableQueryCommand<T, P> {
    Sync(P, oneshot::Sender<anyhow::Result<(T, P)>>),
}

/// A SyncableQuery is similar to a subscription. When created, it accepts a query expression
/// and a SyncableQueryProcessor. It can be used to maintain some state that depends on the
/// results of the query expression. As changes are made to the filesystem, the SyncableQueryProcessor
/// will be called to process the events.
///
/// It provides a single function `sync()` that will ensure that all file events seen by watchman have
/// been processed by the SyncableQueryProcessor.
///
/// In the background, the SyncableQuery may use a subscription to eagerly process updates, but this is
/// only an optimization and users should use `sync()` when they want events to have been processed.
pub struct SyncableQuery<T, P> {
    control_tx: UnboundedSender<SyncableQueryCommand<T, P>>,
}

pub enum WatchmanSyncResult {
    FreshInstance {
        merge_base: Option<String>,
        clock: ClockSpec,
        watchman_version: Option<String>,
    },
    Events {
        events: Vec<WatchmanEvent>,
        merge_base: Option<String>,
        clock: ClockSpec,
        watchman_version: Option<String>,
    },
}

/// The SyncableQueryHandler runs a processing loop that communicates with watchman and invokes the SyncableQueryProcessor. As
/// these only happen within the handler's run loop, for a particular SyncableQuery all watchman invocations and processor
/// processing will happen in a linear order.
///
/// The SyncableQueryHandler maintains the clock and last mergebase and updates them with each request.
struct SyncableQueryHandler<T, P> {
    connector: Connector,
    path: CanonicalPath,
    processor: Box<dyn SyncableQueryProcessor<Output = T, Payload = P>>,
    query: QueryRequestCommon,
    last_clock: ClockSpec,
    last_mergebase: Option<String>,
    mergebase_with: Option<String>,
    control_rx: UnboundedReceiver<SyncableQueryCommand<T, P>>,
}

impl<T, P> SyncableQueryHandler<T, P>
where
    T: Send + 'static,
    P: Send + 'static,
{
    async fn run_loop(&mut self) {
        // We discard the first error here, if there is one. It's a bit unfortunate but it makes
        // everything a lot simpler below, and kicking it off earlier is desirable because it can
        // give Watchman time to warm up.
        let mut client = None;
        if let Err(e) = self.reconnect(&mut client).await {
            tracing::warn!("Connecting to Watchman failed (will re-attempt): {:#}", e);
        };

        loop {
            match self.control_rx.recv().await {
                Some(SyncableQueryCommand::Sync(dice, sync_tx)) => {
                    let res = self.sync(dice, &mut client).await;

                    // NOTE: If the receiver is gone, then they won't be told we finished their
                    // job. That's fine.
                    let _ignore = sync_tx.send(res);
                }
                None => {
                    // This indicates the controlling SyncableQuery has been dropped.
                    return;
                }
            }
        }
    }

    /// sync() will send a since query to watchman and invoke the processor
    /// with either the received events or a fresh instance call.
    async fn sync(
        &mut self,
        payload: P,
        client: &mut Option<WatchmanClient>,
    ) -> anyhow::Result<(T, P)> {
        let sync_res = match self.sync_query(client).await {
            Ok(res) => Ok(res),
            Err(e) => self.reconnect_and_sync_query(client).await.context(e),
        }?;

        let (res, new_mergebase, clock) = match sync_res {
            WatchmanSyncResult::Events {
                events,
                merge_base,
                clock,
                watchman_version,
            } => {
                if self.mergebase_with.is_none()
                    || self.last_mergebase.is_some() && self.last_mergebase == merge_base
                {
                    (
                        self.processor
                            .process_events(payload, events, &merge_base, watchman_version)
                            .await?,
                        merge_base,
                        clock,
                    )
                } else {
                    (
                        self.processor
                            .on_fresh_instance(payload, &merge_base, watchman_version)
                            .await?,
                        merge_base,
                        clock,
                    )
                }
            }
            WatchmanSyncResult::FreshInstance {
                merge_base,
                clock,
                watchman_version,
            } => (
                self.processor
                    .on_fresh_instance(payload, &merge_base, watchman_version)
                    .await?,
                merge_base,
                clock,
            ),
        };

        self.last_mergebase = new_mergebase;
        self.last_clock = clock;

        Ok(res)
    }

    async fn reconnect(&mut self, client: &mut Option<WatchmanClient>) -> anyhow::Result<()> {
        self.last_clock = Default::default();
        self.last_mergebase = None;
        *client = Some(
            WatchmanClient::connect(&self.connector, self.path.clone())
                .await
                .context("Error reconnecting to Watchman")?,
        );
        Ok(())
    }

    async fn reconnect_and_sync_query(
        &mut self,
        client: &mut Option<WatchmanClient>,
    ) -> anyhow::Result<WatchmanSyncResult> {
        self.reconnect(client)
            .await
            .context("Error reconnecting to Watchman")?;

        let out = self.sync_query(client).await?;

        Ok(out)
    }

    async fn sync_query(
        &mut self,
        client: &mut Option<WatchmanClient>,
    ) -> anyhow::Result<WatchmanSyncResult> {
        let client = client.as_mut().context("No Watchman connection")?;

        let mut query = self.query.clone();
        query.since = if let Some(mergebase_with) = self.mergebase_with.as_ref() {
            Some(Clock::ScmAware(FatClockData {
                clock: self.last_clock.clone(),
                scm: Some(ScmAwareClockData {
                    mergebase: self.last_mergebase.clone(),
                    mergebase_with: Some(mergebase_with.clone()),
                    saved_state: None,
                }),
            }))
        } else {
            Some(Clock::Spec(self.last_clock.clone()))
        };

        let QueryResult {
            version,
            is_fresh_instance,
            files,
            clock,
            // subscription_canceled,
            // state_enter,
            // state_leave,
            // state_metadata,
            // saved_state_info,
            ..
        } = client.query::<BuckQueryResult>(query).await?;

        // While we use scm-based queries, the processor api doesn't really support them yet so we just treat it as a fresh instance.
        let (new_mergebase, clock) = unpack_clock(clock);

        Ok(if is_fresh_instance {
            WatchmanSyncResult::FreshInstance {
                merge_base: new_mergebase,
                clock,
                watchman_version: Some(version),
            }
        } else {
            WatchmanSyncResult::Events {
                events: files
                    .ok_or_else(|| anyhow::anyhow!(""))?
                    .into_iter()
                    .filter_map(|f| f.into_event())
                    .collect(),
                merge_base: new_mergebase,
                clock,
                watchman_version: Some(version),
            }
        })
    }
}

/// Unpacks the clock returned for an scm-aware query into a tuple of the mergebase and the clockspec.
fn unpack_clock(clock: Clock) -> (Option<String>, ClockSpec) {
    match clock {
        Clock::Spec(clock_spec) => (None, clock_spec),
        Clock::ScmAware(FatClockData {
            clock: clock_spec,
            scm:
                Some(ScmAwareClockData {
                    mergebase: Some(mergebase),
                    ..
                }),
        }) => (Some(mergebase), clock_spec),
        clock => panic!(
            "requested watchman query, got unexpected clock `{:?}`",
            clock
        ),
    }
}

impl<T, P> SyncableQuery<T, P>
where
    T: Send + 'static,
    P: Send + 'static,
{
    /// Ensures that the processor has been sent all changes that watchman has seen.
    pub fn sync(&self, dice: P) -> impl Future<Output = anyhow::Result<(T, P)>> + Send + 'static {
        let (sync_done_tx, sync_done_rx) = tokio::sync::oneshot::channel();
        let tx_res = self
            .control_tx
            .send(SyncableQueryCommand::Sync(dice, sync_done_tx));

        async move {
            tx_res.ok().context("SyncableQueryHandler has exited")?;

            let out = sync_done_rx
                .await
                .context("SyncableQueryHandler did not return a response for sync request")?
                .context("SyncableQueryHandler returned an error")?;

            Ok(out)
        }
    }

    pub fn new(
        connector: Connector,
        path: impl AsRef<Path>,
        expr: Expr,
        processor: Box<dyn SyncableQueryProcessor<Output = T, Payload = P>>,
        mergebase_with: Option<String>,
    ) -> anyhow::Result<SyncableQuery<T, P>> {
        let path = path.as_ref();
        let path = CanonicalPath::canonicalize(path)
            .with_context(|| format!("Error canonicalizing: `{}`", path.display()))?;

        let query = QueryRequestCommon {
            expression: Some(expr),
            fields: vec!["name"],
            empty_on_fresh_instance: true,
            relative_root: None,
            case_sensitive: true,
            dedup_results: false,
            // Required or we miss directory events
            always_include_directories: true,
            // TODO(cjhopman): Figure out reasonable timeouts.
            // sync_timeout: ???,
            // lock_timeout: ???,
            ..QueryRequestCommon::default()
        };

        let (control_tx, control_rx) =
            tokio::sync::mpsc::unbounded_channel::<SyncableQueryCommand<T, P>>();

        tokio::spawn(async move {
            let mut handler = SyncableQueryHandler {
                connector,
                path,
                query,
                last_clock: ClockSpec::default(),
                last_mergebase: None,
                mergebase_with,
                processor,
                control_rx,
            };
            handler.run_loop().await
        });

        Ok(Self { control_tx })
    }
}


/*
 * 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::path::Path;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context as _;
use async_trait::async_trait;
use buck2_common::dice::file_ops::FileChangeTracker;
use buck2_common::ignores::ignore_set::IgnoreSet;
use buck2_common::legacy_configs::LegacyBuckConfig;
use buck2_core::cells::name::CellName;
use buck2_core::cells::CellResolver;
use buck2_core::fs::paths::abs_norm_path::AbsNormPath;
use buck2_core::fs::project_rel_path::ProjectRelativePath;
use buck2_core::rollout_percentage::RolloutPercentage;
use buck2_events::dispatch::span_async;
use dice::DiceTransactionUpdater;
use tracing::info;
use tracing::warn;
use watchman_client::expr::Expr;
use watchman_client::prelude::Connector;
use watchman_client::prelude::FileType;

use crate::file_watcher::FileWatcher;
use crate::mergebase::Mergebase;
use crate::stats::FileWatcherStats;
use crate::watchman::core::SyncableQuery;
use crate::watchman::core::SyncableQueryProcessor;
use crate::watchman::core::WatchmanEvent;
use crate::watchman::core::WatchmanEventType;
use crate::watchman::core::WatchmanKind;

struct WatchmanQueryProcessor {
    cells: CellResolver,
    ignore_specs: HashMap<CellName, IgnoreSet>,
    retain_dep_files_on_watchman_fresh_instance: bool,
    last_mergebase: Option<String>,
}

/// Used in process_one_change
#[derive(Debug)]
enum ChangeEvent<'a> {
    Watchman(&'a WatchmanEvent),
    SyntheticDirectoryChange,
}

impl WatchmanQueryProcessor {
    async fn process_events_impl(
        &self,
        mut ctx: DiceTransactionUpdater,
        events: Vec<WatchmanEvent>,
        mergebase: &Option<String>,
        watchman_version: Option<String>,
    ) -> anyhow::Result<(buck2_data::FileWatcherStats, DiceTransactionUpdater)> {
        let mut handler = FileChangeTracker::new();
        let mut stats = FileWatcherStats::new(events.len(), mergebase.as_deref(), watchman_version);

        for ev in events {
            // If the path is invalid, then walk up all the way until you find a valid dir to
            // invalidate listings. We don't need to invalidate the file itself, as we can't
            // read invalid files.

            let (path, event) = match ProjectRelativePath::new(&ev.path) {
                Ok(path) => (path, ChangeEvent::Watchman(&ev)),
                Err(_) => {
                    // If we error out here then we might miss other changes. This seems like
                    // it shouldn't happen, since the empty path should always be a valid path.
                    let path = find_first_valid_parent(&ev.path)
                        .with_context(|| {
                            format!("Invalid path had no valid parent: `{}`", ev.path.display())
                        })
                        .unwrap();

                    (path, ChangeEvent::SyntheticDirectoryChange)
                }
            };

            self.process_one_change(path, event, &mut handler, &mut stats)?;
        }

        let stats = stats.finish();
        handler.write_to_dice(&mut ctx)?;

        Ok((stats, ctx))
    }

    fn process_one_change(
        &self,
        path: &ProjectRelativePath,
        ev: ChangeEvent<'_>,
        handler: &mut FileChangeTracker,
        stats: &mut FileWatcherStats,
    ) -> anyhow::Result<()> {
        let cell_path = self.cells.get_cell_path(path)?;

        let ignore = self
            .ignore_specs
            .get(&cell_path.cell())
            .expect("unexpected cell name mismatch")
            .is_match(cell_path.path());

        info!("Watchman: {:?} (ignore = {})", ev, ignore);

        if ignore {
            stats.add_ignored(1);
        } else {
            let cell_path_str = cell_path.to_string();
            let log_kind;
            let log_event;

            match ev {
                ChangeEvent::Watchman(ev) => {
                    match (&ev.kind, &ev.event) {
                        (WatchmanKind::File, typ) => {
                            log_kind = buck2_data::FileWatcherKind::File;
                            match typ {
                                WatchmanEventType::Modify => {
                                    handler.file_changed(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Modify;
                                }
                                WatchmanEventType::Create => {
                                    handler.file_added(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Create;
                                }
                                WatchmanEventType::Delete => {
                                    handler.file_removed(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Delete;
                                }
                            }
                        }
                        (WatchmanKind::Directory, typ) => {
                            log_kind = buck2_data::FileWatcherKind::Directory;
                            match typ {
                                WatchmanEventType::Modify => {
                                    // We can safely ignore this, as it corresponds to files being added or removed,
                                    // but there are always file add/remove notifications sent too.
                                    // See https://fb.workplace.com/groups/watchman.users/permalink/2858842194433249
                                    return Ok(());
                                }
                                WatchmanEventType::Create => {
                                    handler.dir_added(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Create;
                                }
                                WatchmanEventType::Delete => {
                                    handler.dir_removed(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Delete;
                                }
                            }
                        }
                        (WatchmanKind::Symlink, typ) => {
                            log_kind = buck2_data::FileWatcherKind::Symlink;
                            match typ {
                                WatchmanEventType::Modify => {
                                    handler.file_changed(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Modify;
                                }
                                WatchmanEventType::Create => {
                                    warn!(
                                        "New symlink detected (source symlinks are not supported): {}",
                                        cell_path
                                    );
                                    handler.file_added(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Create;
                                }
                                WatchmanEventType::Delete => {
                                    handler.file_removed(cell_path);
                                    log_event = buck2_data::FileWatcherEventType::Delete;
                                }
                            }
                        }
                    }
                }
                ChangeEvent::SyntheticDirectoryChange => {
                    log_kind = buck2_data::FileWatcherKind::Directory;
                    log_event = buck2_data::FileWatcherEventType::Modify;
                    handler.dir_changed(cell_path);
                }
            };

            stats.add(cell_path_str, log_event, log_kind);
        }

        Ok(())
    }
}

fn find_first_valid_parent(mut path: &Path) -> Option<&ProjectRelativePath> {
    loop {
        path = path.parent()?;

        match ProjectRelativePath::new(path) {
            Ok(path) => return Some(path),
            Err(_) => {}
        }
    }
}

#[async_trait]
impl SyncableQueryProcessor for WatchmanQueryProcessor {
    type Output = buck2_data::FileWatcherStats;
    type Payload = DiceTransactionUpdater;

    async fn process_events(
        &mut self,
        dice: DiceTransactionUpdater,
        events: Vec<WatchmanEvent>,
        mergebase: &Option<String>,
        watchman_version: Option<String>,
    ) -> anyhow::Result<(Self::Output, DiceTransactionUpdater)> {
        self.last_mergebase = mergebase.clone();
        self.process_events_impl(dice, events, mergebase, watchman_version)
            .await
    }

    async fn on_fresh_instance(
        &mut self,
        ctx: DiceTransactionUpdater,
        mergebase: &Option<String>,
        watchman_version: Option<String>,
    ) -> anyhow::Result<(Self::Output, DiceTransactionUpdater)> {
        let has_new_mergebase = self.last_mergebase.as_ref() != mergebase.as_ref();

        let clear_dep_files =
            has_new_mergebase || !self.retain_dep_files_on_watchman_fresh_instance;

        // We'll clear dep files if we're configured to do so on all fresh instances. Otherwise,
        // we'll drop them if the mergebase has changed, which means our dep files are likely
        // irrelevant.
        //
        // This is imperfect. If the user rebased from yesterday's stable to today's stable, then
        // flushing is the right thing to do. In contrast, if they rebased from X to X's parent,
        // then it probably isn't. The consequences of flushing in the latter case aren't as bad
        // (where we'll rebuild things our dep files *could* have avoided) as not flushing in the
        // former (where we'll fetch loads of dep files that all miss), so we err on the side of
        // being safe and drop them when the mergebase changes.
        if clear_dep_files {
            crate::dep_files::flush_dep_files();
        }

        self.last_mergebase = mergebase.clone();

        // TODO(cjhopman): could probably get away with just invalidating all fs things, but that's not supported.
        // Dropping the entire DICE map can be somewhat computationally expensive as there
        // are a lot of destructors to run. On the other hand, we don't have to wait for
        // it. So, we just send it off to its own thread.
        let ctx = ctx.unstable_take();

        Ok((
            buck2_data::FileWatcherStats {
                fresh_instance: true,
                branched_from_revision: mergebase.clone(),
                incomplete_events_reason: Some("Fresh instance".to_owned()),
                watchman_version,
                fresh_instance_data: Some(buck2_data::FreshInstance {
                    new_mergebase: has_new_mergebase,
                    cleared_dice: true,
                    cleared_dep_files: clear_dep_files,
                }),
                ..Default::default()
            },
            ctx,
        ))
    }
}

#[derive(Allocative)]
pub(crate) struct WatchmanFileWatcher {
    #[allocative(skip)]
    query: SyncableQuery<buck2_data::FileWatcherStats, DiceTransactionUpdater>,
}

/// The watchman query is constructed once on daemon startup. It is an unfiltered watchman query
/// over the entire project (though, like all watchman queries, still filtered by the
/// .watchmanconfig itself). Before any new computation request is started, it will be synced to
/// ensure that any recent changes are flushed and visible to the computation.
impl WatchmanFileWatcher {
    pub(crate) fn new(
        project_root: &AbsNormPath,
        root_config: &LegacyBuckConfig,
        cells: CellResolver,
        ignore_specs: HashMap<CellName, IgnoreSet>,
    ) -> anyhow::Result<Self> {
        let watchman_merge_base = root_config
            .get("project", "watchman_merge_base")
            .map(|s| s.to_owned());

        let retain_dep_files_on_watchman_fresh_instance = root_config
            .parse::<RolloutPercentage>("buck2", "retain_dep_files_on_watchman_fresh_instance")?
            .unwrap_or_else(RolloutPercentage::always)
            .roll();

        let query = SyncableQuery::new(
            Connector::new(),
            project_root,
            Expr::Any(vec![
                Expr::FileType(FileType::Regular),
                Expr::FileType(FileType::Directory),
                Expr::FileType(FileType::Symlink),
            ]),
            Box::new(WatchmanQueryProcessor {
                cells,
                ignore_specs,
                retain_dep_files_on_watchman_fresh_instance,
                last_mergebase: None,
            }),
            watchman_merge_base,
        )?;

        Ok(Self { query })
    }
}

#[async_trait]
impl FileWatcher for WatchmanFileWatcher {
    async fn sync(
        &self,
        dice: DiceTransactionUpdater,
    ) -> anyhow::Result<(DiceTransactionUpdater, Mergebase)> {
        span_async(
            buck2_data::FileWatcherStart {
                provider: buck2_data::FileWatcherProvider::Watchman as i32,
            },
            async {
                let (stats, res) = match self.query.sync(dice).await {
                    Ok((stats, dice)) => {
                        let mergebase = Mergebase(Arc::new(stats.branched_from_revision.clone()));
                        ((Some(stats)), Ok((dice, mergebase)))
                    }
                    Err(e) => (None, Err(e)),
                };
                (res, buck2_data::FileWatcherEnd { stats })
            },
        )
        .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.
 */

#[cfg(all(test, not(windows)))]
mod test;

pub(crate) mod core;
pub(crate) mod interface;


/*
 * 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::fs;
use std::fs::File;
use std::path::Path;
use std::path::PathBuf;
use std::time::Duration;

use anyhow::Context as _;
use assert_matches::assert_matches;
use async_trait::async_trait;
use buck2_util::process::async_background_command;
use gazebo::prelude::*;
use tokio::io::AsyncWriteExt;
use tokio::process::Child;
use watchman_client::expr::Expr;
use watchman_client::prelude::Connector;
use watchman_client::prelude::FileType;

use crate::watchman::core::SyncableQuery;
use crate::watchman::core::SyncableQueryProcessor;
use crate::watchman::core::WatchmanEvent;

struct TestQueryProcessor;

#[derive(PartialEq, Eq, Debug)]
enum Out {
    FreshInstance,
    Files(Vec<String>),
}

#[async_trait]
impl SyncableQueryProcessor for TestQueryProcessor {
    type Output = Out;
    type Payload = ();

    async fn process_events(
        &mut self,
        payload: Self::Payload,
        events: Vec<WatchmanEvent>,
        _mergebase: &Option<String>,
        _watchman_version: Option<String>,
    ) -> anyhow::Result<(Self::Output, Self::Payload)> {
        Ok((
            Out::Files(events.into_map(|e| e.path.display().to_string())),
            payload,
        ))
    }

    async fn on_fresh_instance(
        &mut self,
        payload: Self::Payload,
        _mergebase: &Option<String>,
        _watchman_version: Option<String>,
    ) -> anyhow::Result<(Self::Output, Self::Payload)> {
        Ok((Out::FreshInstance, payload))
    }
}

async fn wait_for_watchman(watchman_sock: &Path) -> anyhow::Result<()> {
    let connector = Connector::default().unix_domain_socket(watchman_sock);

    let mut i = 0;
    loop {
        i += 1;

        match connector.connect().await {
            Ok(..) => return Ok(()),
            Err(e) if i >= 100 => return Err(anyhow::anyhow!("{}", e)),
            _ => {
                tokio::time::sleep(Duration::from_millis(50)).await;
            }
        }
    }
}

struct WatchmanInstance {
    child: Option<Child>,
    sock: PathBuf,
    log: PathBuf,
}

impl WatchmanInstance {
    pub async fn shutdown(&mut self) -> anyhow::Result<()> {
        let child = self
            .child
            .as_mut()
            .context("Watchman was already shutdown")?;

        child.kill().await.context("Failed to kill Watchman")?;

        // Everything went well. Remove the child so we don't log on Drop.
        self.child.take();

        Ok(())
    }
}

impl Drop for WatchmanInstance {
    fn drop(&mut self) {
        let child = match self.child.as_mut() {
            Some(c) => c,
            None => return,
        };

        // If we get here, something went wrong and we didn't stop Watchman properly. Log debug
        // info.
        eprintln!("WatchmanInstance did not exit cleanly!");
        let log = std::fs::read_to_string(&self.log)
            .with_context(|| format!("Failed to read log file at {}", self.log.display()));
        match log {
            Ok(log) => eprintln!("Watchman logs follow\n{}", log),
            Err(e) => eprintln!("Failed to read logs: {:#}", e),
        };

        // Try to see if Watchman had exited or not.
        match child.try_wait() {
            Ok(Some(status)) => eprintln!("Watchman had exited with status {:?}", status),
            Ok(None) => eprintln!("Watchan is still running"),
            Err(e) => eprintln!("Failed to access Watchman status: {:#}", e),
        }
    }
}

async fn spawn_watchman(watchman_dir: &Path) -> anyhow::Result<WatchmanInstance> {
    // This config might make Watchman a bit more efficient on Mac and make tests less flaky there
    let watchman_config = watchman_dir.join("config");
    let watchman_config_text = r#"{"prefer_split_fsevents_watcher": true}"#;

    {
        let mut conf = tokio::fs::File::create(&watchman_config).await?;
        conf.write_all(watchman_config_text.as_bytes()).await?;
        conf.flush().await?;
    }

    let watchman_sock = watchman_dir.join("sock");
    let watchman_log = watchman_dir.join("log");
    let watchman_pid = watchman_dir.join("pid");

    let watchman = async_background_command("watchman")
        .arg("--unix-listener-path")
        .arg(&watchman_sock)
        .arg("--foreground")
        .arg("--no-site-spawner")
        .arg("--logfile")
        .arg(&watchman_log)
        .arg("--no-save-state")
        .arg("--pidfile")
        .arg(&watchman_pid)
        .env("WATCHMAN_CONFIG_FILE", watchman_config)
        .kill_on_drop(true)
        .spawn()?;

    wait_for_watchman(&watchman_sock)
        .await
        .context("Waiting for Watchman to start")?;

    Ok(WatchmanInstance {
        child: Some(watchman),
        sock: watchman_sock,
        log: watchman_log,
    })
}

#[tokio::test]
async fn test_syncable_query() -> anyhow::Result<()> {
    // This test doesn't work unless Watchman is working, so let's
    // over-approximate that as fbcode_build for now.
    if !cfg!(fbcode_build) {
        return Ok(());
    }

    let tempdir = tempfile::tempdir()?;

    // NOTE: This isn't async (and so is tempfile), but this is all tests.
    let root = tempdir.path().join("root");
    let watchman_dir = tempdir.path().join("watchman");
    fs::create_dir(&watchman_dir)?;
    fs::create_dir(&root)?;

    let mut watchman_instance = spawn_watchman(&watchman_dir).await?;

    let connector = Connector::default().unix_domain_socket(&watchman_instance.sock);

    let watchman_query = SyncableQuery::new(
        connector,
        &root,
        Expr::Any(vec![Expr::FileType(FileType::Regular)]),
        Box::new(TestQueryProcessor),
        None,
    )?;

    // Startup
    assert_eq!(watchman_query.sync(()).await?.0, Out::FreshInstance);
    assert_eq!(watchman_query.sync(()).await?.0, Out::Files(vec![]));

    // Create a file, see that we receive it.
    let test = root.join("test");
    File::create(&test)?;
    assert_eq!(
        watchman_query.sync(()).await?.0,
        Out::Files(vec!["test".into()])
    );

    // Kill Watchman, see that we're broken now
    watchman_instance.shutdown().await?;
    assert_matches!(watchman_query.sync(()).await, Err(..));

    // Restart Watchman, we should be fixed now, and get a fresh instance.
    let mut watchman_instance = spawn_watchman(&watchman_dir).await?;
    assert_eq!(watchman_query.sync(()).await?.0, Out::FreshInstance);
    assert_eq!(watchman_query.sync(()).await?.0, Out::Files(vec![]));

    // Clean up
    watchman_instance.shutdown().await?;

    Ok(())
}


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

use std::io;
use std::process::ExitStatus;
use std::sync::Arc;

use allocative::Allocative;
use anyhow::Context;
use arc_swap::ArcSwapOption;
use buck2_common::result::SharedError;
use buck2_core::tag_error;
use dupe::Dupe;
use futures::future;
use futures::future::Future;
use futures::future::FutureExt;
use futures::stream;
use futures::stream::StreamExt;
use tokio::process::Child;
use tonic::transport::Channel;
use tonic::Request;

use crate::convert::decode_event_stream;
use crate::run::decode_command_event_stream;
use crate::run::GatherOutputStatus;

#[derive(Clone, Dupe, Allocative)]
pub struct ForkserverClient {
    inner: Arc<ForkserverClientInner>,
}

#[derive(thiserror::Error, Debug)]
enum ForkserverError {
    #[error("Error on Forkserver wait()")]
    WaitError(#[source] io::Error),
    #[error("Forkserver exited {}", _0)]
    Exited(ExitStatus),
}

#[derive(Allocative)]
struct ForkserverClientInner {
    /// Error from the forkserver process, if any.
    #[allocative(skip)]
    error: Arc<ArcSwapOption<anyhow::Error>>,
    pid: u32,
    #[allocative(skip)]
    rpc: buck2_forkserver_proto::forkserver_client::ForkserverClient<Channel>,
}

impl ForkserverClient {
    #[allow(unused)] // Unused on Windows
    pub(crate) fn new(mut child: Child, channel: Channel) -> Self {
        let rpc = buck2_forkserver_proto::forkserver_client::ForkserverClient::new(channel)
            .max_encoding_message_size(usize::MAX)
            .max_decoding_message_size(usize::MAX);

        let pid = child.id().expect("Child has not been polled");

        let error = Arc::new(ArcSwapOption::empty());

        tokio::task::spawn({
            let error = error.clone();

            async move {
                let err = match child.wait().await {
                    Ok(status) => ForkserverError::Exited(status),
                    Err(e) => ForkserverError::WaitError(e),
                };

                let err = anyhow::Error::new(err).context("Forkserver is unavailable");

                error.swap(Some(Arc::new(err)));
            }
        });

        Self {
            inner: Arc::new(ForkserverClientInner { error, pid, rpc }),
        }
    }

    pub fn pid(&self) -> u32 {
        self.inner.pid
    }

    pub async fn execute<C>(
        &self,
        req: buck2_forkserver_proto::CommandRequest,
        cancel: C,
    ) -> anyhow::Result<(GatherOutputStatus, Vec<u8>, Vec<u8>)>
    where
        C: Future<Output = ()> + Send + 'static,
    {
        if let Some(err) = &*self.inner.error.load() {
            return Err(tag_error!(
                "forkserver_exit",
                SharedError::from(err.dupe()).into(),
                quiet: true,
                task: false,
                daemon_in_memory_state_is_corrupted: true,
            ));
        }

        let stream = stream::once(future::ready(buck2_forkserver_proto::RequestEvent {
            data: Some(req.into()),
        }))
        .chain(stream::once(cancel.map(|()| {
            buck2_forkserver_proto::RequestEvent {
                data: Some(buck2_forkserver_proto::CancelRequest {}.into()),
            }
        })));

        let stream = self
            .inner
            .rpc
            .clone()
            .run(stream)
            .await
            .context("Error dispatching command to Forkserver")?
            .into_inner();
        let stream = decode_event_stream(stream);
        decode_command_event_stream(stream).await
    }

    pub async fn set_log_filter(&self, log_filter: String) -> anyhow::Result<()> {
        self.inner
            .rpc
            .clone()
            .set_log_filter(Request::new(buck2_forkserver_proto::SetLogFilterRequest {
                log_filter,
            }))
            .await?;

        Ok(())
    }
}


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

use anyhow::Context as _;
use buck2_common::convert::ProstDurationExt;
use futures::stream::Stream;
use futures::stream::StreamExt;

use crate::run::CommandEvent;
use crate::run::GatherOutputStatus;

pub fn encode_event_stream<S>(
    s: S,
) -> impl Stream<Item = Result<buck2_forkserver_proto::CommandEvent, tonic::Status>>
where
    S: Stream<Item = anyhow::Result<CommandEvent>>,
{
    fn convert_event(e: CommandEvent) -> buck2_forkserver_proto::CommandEvent {
        use buck2_forkserver_proto::command_event::Data;

        let data = match e {
            CommandEvent::Stdout(bytes) => Data::Stdout(buck2_forkserver_proto::StreamEvent {
                data: bytes.to_vec(),
            }),
            CommandEvent::Stderr(bytes) => Data::Stderr(buck2_forkserver_proto::StreamEvent {
                data: bytes.to_vec(),
            }),
            CommandEvent::Exit(GatherOutputStatus::Finished {
                exit_code,
                execution_stats,
            }) => Data::Exit(buck2_forkserver_proto::ExitEvent {
                exit_code,
                execution_stats,
            }),
            CommandEvent::Exit(GatherOutputStatus::TimedOut(duration)) => {
                Data::Timeout(buck2_forkserver_proto::TimeoutEvent {
                    duration: duration.try_into().ok(),
                })
            }
            CommandEvent::Exit(GatherOutputStatus::Cancelled) => {
                Data::Cancel(buck2_forkserver_proto::CancelEvent {})
            }
            CommandEvent::Exit(GatherOutputStatus::SpawnFailed(reason)) => {
                Data::SpawnFailed(buck2_forkserver_proto::SpawnFailedEvent { reason })
            }
        };

        buck2_forkserver_proto::CommandEvent { data: Some(data) }
    }

    fn convert_err(e: anyhow::Error) -> tonic::Status {
        tonic::Status::unknown(format!("{:#}", e))
    }

    s.map(|r| r.map(convert_event).map_err(convert_err))
}

pub fn decode_event_stream<S>(s: S) -> impl Stream<Item = anyhow::Result<CommandEvent>>
where
    S: Stream<Item = Result<buck2_forkserver_proto::CommandEvent, tonic::Status>>,
{
    fn convert_event(e: buck2_forkserver_proto::CommandEvent) -> anyhow::Result<CommandEvent> {
        use buck2_forkserver_proto::command_event::Data;

        let event = match e.data.context("Missing `data`")? {
            Data::Stdout(buck2_forkserver_proto::StreamEvent { data }) => {
                CommandEvent::Stdout(data.into())
            }
            Data::Stderr(buck2_forkserver_proto::StreamEvent { data }) => {
                CommandEvent::Stderr(data.into())
            }
            Data::Exit(buck2_forkserver_proto::ExitEvent {
                exit_code,
                execution_stats,
            }) => CommandEvent::Exit(GatherOutputStatus::Finished {
                exit_code,
                execution_stats,
            }),
            Data::Timeout(buck2_forkserver_proto::TimeoutEvent { duration }) => {
                CommandEvent::Exit(GatherOutputStatus::TimedOut(
                    duration
                        .context("Missing `duration`")?
                        .try_into_duration()
                        .context("Invalid `duration`")?,
                ))
            }
            Data::Cancel(buck2_forkserver_proto::CancelEvent {}) => {
                CommandEvent::Exit(GatherOutputStatus::Cancelled)
            }
            Data::SpawnFailed(buck2_forkserver_proto::SpawnFailedEvent { reason }) => {
                CommandEvent::Exit(GatherOutputStatus::SpawnFailed(reason))
            }
        };

        Ok(event)
    }

    fn convert_err(e: tonic::Status) -> anyhow::Error {
        anyhow::anyhow!("forkserver error: {}", e.message())
    }

    s.map(|r| r.map_err(convert_err).and_then(convert_event))
}


/*
 * 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_attr(unix, allow(stable_features))]
#![cfg_attr(unix, feature(process_set_process_group))]

pub mod client;
pub mod convert;
pub mod run;

#[cfg(unix)]
pub mod unix;