/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

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

use anyhow::Error;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkUpdateLogEntry;
use bookmarks::BookmarkUpdateReason;
use cloned::cloned;
use commit_graph::CommitGraph;
use commit_graph::CommitGraphRef;
use context::CoreContext;
use futures::future::try_join;
use futures::future::try_join_all;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::future::TryFutureExt;
use futures::lock::Mutex;
use futures_watchdog::WatchdogExt;
use getbundle_response::SessionLfsParams;
use itertools::Itertools;
use mercurial_derivation::DeriveHgChangeset;
use mercurial_types::HgChangesetId;
use metaconfig_types::RepoConfigRef;
use mononoke_hg_sync_job_helper_lib::save_bytes_to_temp_file;
use mononoke_hg_sync_job_helper_lib::write_to_named_temp_file;
use mononoke_types::datetime::Timestamp;
use mononoke_types::ChangesetId;
use mononoke_types::Generation;
use regex::Regex;
use slog::info;
use slog::warn;
use tempfile::NamedTempFile;

use crate::bind_sync_err;
use crate::bundle_generator::BookmarkChange;
use crate::bundle_generator::FilenodeVerifier;
use crate::bundle_generator::FilterExistingChangesets;
use crate::errors::ErrorKind::BookmarkMismatchInBundleCombining;
use crate::errors::ErrorKind::UnexpectedBookmarkMove;
use crate::errors::PipelineError;
use crate::BookmarkOverlay;
use crate::CombinedBookmarkUpdateLogEntry;
use crate::CommitsInBundle;
use crate::Repo;

#[derive(Clone)]
pub struct PreparedBookmarkUpdateLogEntry {
    pub log_entry: BookmarkUpdateLogEntry,
    pub bundle_file: Arc<NamedTempFile>,
    pub timestamps_file: Arc<NamedTempFile>,
    pub cs_id: Option<(ChangesetId, HgChangesetId)>,
}

pub struct BundlePreparer {
    repo: Repo,
    base_retry_delay_ms: u64,
    retry_num: usize,
    filenode_verifier: FilenodeVerifier,
    bookmark_regex_force_lfs: Option<Regex>,
    push_vars: Option<HashMap<String, bytes::Bytes>>,
}

impl BundlePreparer {
    pub async fn new_generate_bundles(
        repo: Repo,
        base_retry_delay_ms: u64,
        retry_num: usize,
        filenode_verifier: FilenodeVerifier,
        bookmark_regex_force_lfs: Option<Regex>,
        push_vars: Option<HashMap<String, bytes::Bytes>>,
    ) -> Result<BundlePreparer, Error> {
        Ok(BundlePreparer {
            repo,
            base_retry_delay_ms,
            retry_num,
            filenode_verifier,
            bookmark_regex_force_lfs,
            push_vars,
        })
    }

    pub async fn prepare_batches(
        &self,
        ctx: &CoreContext,
        overlay: Arc<Mutex<BookmarkOverlay>>,
        entries: Vec<BookmarkUpdateLogEntry>,
        commit_limit: u64,
    ) -> Result<Vec<BookmarkLogEntryBatch>, Error> {
        use BookmarkUpdateReason::*;

        for log_entry in &entries {
            match log_entry.reason {
                Pushrebase | Backsyncer | ManualMove | ApiRequest | XRepoSync | Push | TestMove => {
                }
                Blobimport => {
                    return Err(UnexpectedBookmarkMove(format!("{}", log_entry.reason)).into());
                }
            };
        }

        split_in_batches(
            ctx,
            self.repo.commit_graph(),
            overlay,
            entries,
            commit_limit,
        )
        .await
    }

    pub fn prepare_bundles(
        &self,
        ctx: CoreContext,
        batches: Vec<BookmarkLogEntryBatch>,
        filter_changesets: Arc<dyn FilterExistingChangesets>,
    ) -> BoxFuture<'static, Result<Vec<CombinedBookmarkUpdateLogEntry>, PipelineError>> {
        let mut futs = vec![];

        for batch in batches {
            if batch.is_empty() {
                continue;
            }
            let session_lfs_params = self.session_lfs_params(&ctx, &batch.bookmark_name);
            let entries = batch.entries.clone();
            let f = self.prepare_single_bundle(
                ctx.clone(),
                batch,
                session_lfs_params,
                filter_changesets.clone(),
            );
            futs.push((f, entries));
        }

        let futs = futs
            .into_iter()
            .map(|(f, entries)| {
                let ctx = ctx.clone();
                async move {
                    let f = tokio::spawn(f);
                    let res = f.map_err(Error::from).watched(ctx.logger()).await;
                    let res = match res {
                        Ok(Ok(res)) => Ok(res),
                        Ok(Err(err)) => Err(err),
                        Err(err) => Err(err),
                    };
                    res.map_err(|err| bind_sync_err(&entries, err))
                }
            })
            .collect::<Vec<_>>();
        async move { try_join_all(futs).await }.boxed()
    }

    // Prepares a bundle that might be a result of combining a few BookmarkUpdateLogEntry.
    // Note that these entries should all move the same bookmark.
    fn prepare_single_bundle(
        &self,
        ctx: CoreContext,
        batch: BookmarkLogEntryBatch,
        session_lfs_params: SessionLfsParams,
        filter_changesets: Arc<dyn FilterExistingChangesets>,
    ) -> BoxFuture<'static, Result<CombinedBookmarkUpdateLogEntry, Error>> {
        cloned!(self.repo, self.push_vars, self.filenode_verifier);

        let base_retry_delay_ms = self.base_retry_delay_ms;
        let retry_num = self.retry_num;
        async move {
            let entry_ids = batch
                .entries
                .iter()
                .map(|log_entry| log_entry.id)
                .collect::<Vec<_>>();
            info!(ctx.logger(), "preparing log entry ids #{:?} ...", entry_ids);
            // Check that all entries modify bookmark_name
            for entry in &batch.entries {
                if entry.bookmark_name != batch.bookmark_name {
                    return Err(BookmarkMismatchInBundleCombining {
                        ids: entry_ids,
                        entry_id: entry.id,
                        entry_bookmark_name: entry.bookmark_name.clone(),
                        bundle_bookmark_name: batch.bookmark_name,
                    }
                    .into());
                }
            }

            let bookmark_change = BookmarkChange::new(batch.from_cs_id, batch.to_cs_id)?;
            let bundle_timestamps_commits = retry::retry_always(
                ctx.logger(),
                {
                    |_| {
                        Self::try_prepare_bundle_and_timestamps_file(
                            &ctx,
                            &repo,
                            &filenode_verifier,
                            session_lfs_params.clone(),
                            &batch.server_heads,
                            &bookmark_change,
                            &batch.bookmark_name,
                            push_vars.clone(),
                            filter_changesets.clone(),
                        )
                    }
                },
                base_retry_delay_ms,
                retry_num,
            )
            .map_ok(|(res, _)| res);

            let cs_id = async {
                match batch.to_cs_id {
                    Some(to_changeset_id) => {
                        let hg_cs_id = repo.derive_hg_changeset(&ctx, to_changeset_id).await?;
                        Ok(Some((to_changeset_id, hg_cs_id)))
                    }
                    None => Ok(None),
                }
            };

            let ((bundle_file, timestamps_file, commits), cs_id) =
                try_join(bundle_timestamps_commits, cs_id).await?;

            info!(
                ctx.logger(),
                "successful prepare of entries #{:?}", entry_ids
            );

            Ok(CombinedBookmarkUpdateLogEntry {
                components: batch.entries,
                bundle_file: Arc::new(bundle_file),
                timestamps_file: Arc::new(timestamps_file),
                cs_id,
                bookmark: batch.bookmark_name,
                commits,
            })
        }
        .boxed()
    }

    async fn try_prepare_bundle_and_timestamps_file<'a>(
        ctx: &'a CoreContext,
        repo: &'a Repo,
        filenode_verifier: &'a FilenodeVerifier,
        session_lfs_params: SessionLfsParams,
        hg_server_heads: &'a [ChangesetId],
        bookmark_change: &'a BookmarkChange,
        bookmark_name: &'a BookmarkKey,
        push_vars: Option<HashMap<String, bytes::Bytes>>,
        filter_changesets: Arc<dyn FilterExistingChangesets>,
    ) -> Result<(NamedTempFile, NamedTempFile, CommitsInBundle), Error> {
        let (bytes, timestamps) = crate::bundle_generator::create_bundle(
            ctx,
            repo,
            bookmark_name,
            bookmark_change,
            hg_server_heads.to_vec(),
            session_lfs_params,
            filenode_verifier,
            push_vars,
            filter_changesets,
        )
        .await?;

        let mut bcs_ids = vec![];
        for (hg_cs_id, (bcs_id, _)) in &timestamps {
            bcs_ids.push((*hg_cs_id, *bcs_id));
        }

        let timestamps = timestamps
            .into_iter()
            .map(|(hg_cs_id, (_, timestamp))| (hg_cs_id, timestamp))
            .collect();
        let (bundle, timestamps) = try_join(
            save_bytes_to_temp_file(&bytes),
            save_timestamps_to_file(&timestamps),
        )
        .await?;
        Ok((bundle, timestamps, CommitsInBundle::Commits(bcs_ids)))
    }

    fn session_lfs_params(&self, ctx: &CoreContext, bookmark: &BookmarkKey) -> SessionLfsParams {
        let lfs_params = &self.repo.repo_config().lfs;

        if let Some(regex) = &self.bookmark_regex_force_lfs {
            if regex.is_match(bookmark.as_str()) {
                info!(ctx.logger(), "force generating lfs bundle for {}", bookmark);
                return SessionLfsParams {
                    threshold: lfs_params.threshold,
                };
            }
        }

        if lfs_params.generate_lfs_blob_in_hg_sync_job {
            SessionLfsParams {
                threshold: lfs_params.threshold,
            }
        } else {
            SessionLfsParams { threshold: None }
        }
    }
}

async fn save_timestamps_to_file(
    timestamps: &HashMap<HgChangesetId, Timestamp>,
) -> Result<NamedTempFile, Error> {
    let encoded_timestamps = timestamps
        .iter()
        .map(|(key, value)| {
            let timestamp = value.timestamp_seconds();
            format!("{}={}", key, timestamp)
        })
        .join("\n");

    write_to_named_temp_file(encoded_timestamps).await
}

#[derive(Clone, Debug, Eq, PartialEq)]
pub struct BookmarkLogEntryBatch {
    entries: Vec<BookmarkUpdateLogEntry>,
    bookmark_name: BookmarkKey,
    from_cs_id: Option<ChangesetId>,
    to_cs_id: Option<ChangesetId>,
    server_old_value: Option<ChangesetId>,
    server_heads: Vec<ChangesetId>,
}

impl BookmarkLogEntryBatch {
    pub fn new(overlay: &mut BookmarkOverlay, log_entry: BookmarkUpdateLogEntry) -> Self {
        let bookmark_name = log_entry.bookmark_name.clone();
        let from_cs_id = log_entry.from_changeset_id;
        let to_cs_id = log_entry.to_changeset_id;
        let server_heads = overlay.all_values();
        let server_old_value = overlay.get(&bookmark_name);
        overlay.update(bookmark_name.clone(), to_cs_id.clone());
        Self {
            entries: vec![log_entry],
            bookmark_name,
            from_cs_id,
            to_cs_id,
            server_old_value,
            server_heads,
        }
    }

    fn is_empty(&self) -> bool {
        self.entries.is_empty()
    }

    // Outer result's error means that some infrastructure error happened.
    // Inner result's error means that it wasn't possible to append entry,
    // and this entry is returned as the error.
    pub async fn try_append(
        &mut self,
        ctx: &CoreContext,
        commit_graph: &CommitGraph,
        overlay: &mut BookmarkOverlay,
        entry: BookmarkUpdateLogEntry,
        commit_limit: u64,
    ) -> Result<Result<(), BookmarkUpdateLogEntry>, Error> {
        // Combine two bookmark update log entries only if bookmark names are the same
        if self.bookmark_name != entry.bookmark_name {
            return Ok(Err(entry));
        }

        // if it's a non-fast forward move then put it in the separate batch.
        // Otherwise some of the commits might not be synced to hg servers.
        // Consider this case:
        // C
        // |
        // B
        // |
        // A
        //
        // 1 entry - moves a bookmark from A to C
        // 2 entry - moves a bookmark from C to B
        //
        // if we combine them together then we get a batch that
        // moves a bookmark from A to B and commit C won't be synced
        // to hg servers. To prevent that let's put non-fast forward
        // moves to a separate branch
        match (entry.from_changeset_id, entry.to_changeset_id) {
            (Some(from_cs_id), Some(to_cs_id)) => {
                let is_ancestor = commit_graph
                    .is_ancestor(ctx, from_cs_id, to_cs_id)
                    .watched(ctx.logger())
                    .await?;
                if !is_ancestor {
                    // Force non-forward moves to go to a separate batch
                    return Ok(Err(entry));
                }
            }
            _ => {}
        };

        // If we got a move where new from_cs_id is not equal to latest to_cs_id then
        // put it in a separate batch. This shouldn't normally happen though
        if self.to_cs_id != entry.from_changeset_id {
            return Ok(Err(entry));
        }
        match entry.to_changeset_id {
            Some(to_cs_id) => {
                let (from_gen, to_gen) = try_join(
                    async {
                        match self.from_cs_id {
                            Some(from_cs_id) => {
                                commit_graph
                                    .changeset_generation_required(ctx, from_cs_id)
                                    .await
                            }
                            None => Ok(Generation::new(0)),
                        }
                    },
                    commit_graph.changeset_generation_required(ctx, to_cs_id),
                )
                .await?;
                if let Some(diff) = to_gen.difference_from(from_gen) {
                    if diff > commit_limit {
                        return Ok(Err(entry));
                    }
                }
            }
            _ => {}
        }

        self.push(overlay, entry);
        Ok(Ok(()))
    }

    fn push(&mut self, overlay: &mut BookmarkOverlay, entry: BookmarkUpdateLogEntry) {
        self.to_cs_id = entry.to_changeset_id;
        self.entries.push(entry);

        // Side-effect: successfully appending the entry to the batch updates
        // the bookmark value in the overlay.  This means the overlay is
        // correct when computing the server heads for the next batch.
        overlay.update(self.bookmark_name.clone(), self.to_cs_id);
    }
}

async fn split_in_batches(
    ctx: &CoreContext,
    commit_graph: &CommitGraph,
    overlay: Arc<Mutex<BookmarkOverlay>>,
    entries: Vec<BookmarkUpdateLogEntry>,
    commit_limit: u64,
) -> Result<Vec<BookmarkLogEntryBatch>, Error> {
    let mut batches: Vec<BookmarkLogEntryBatch> = vec![];
    let mut overlay = overlay.lock().await;

    for entry in entries {
        let entry = match batches.last_mut() {
            Some(batch) => match batch
                .try_append(ctx, commit_graph, &mut overlay, entry, commit_limit)
                .watched(ctx.logger())
                .await?
            {
                Ok(()) => {
                    continue;
                }
                Err(entry) => entry,
            },
            None => entry,
        };
        batches.push(BookmarkLogEntryBatch::new(&mut overlay, entry));
    }

    Ok(batches)
}

// This function might remove a few first bookmark log entries from the batch.
// It can be useful to modify the very first batch sync job is trying to sync. There can be
// mismatches in what was actually synced to hg servers and the value of "latest-replayed-request"
// counter, and this function can help removing these mismatches.
//
// To be precise:
// 1) If a value of the bookmark in overlay points to from_cs_id in the batch, then batch is not
//    modified. Everything is ok, counter is correct.
// 2) If a value of the bookmark in overlay points to to_changeset_id for the middle bookmark
//    log entry in the batch (say, entry X), then all bookmark log entries up to and including
//    entry X are removed.
//    Usually it means that the everything up to entry X was successfully synced, but sync job failed to
//    update "latest-replayed-request" counter.
// 3) If overlay doesn't point to any commit in the batch, then the batch is not modified.
//    Usually it means that hg server is out of date with hgsql, and we don't need to do anything
impl BookmarkLogEntryBatch {
    pub fn maybe_adjust(&mut self, ctx: &CoreContext) {
        if self.server_old_value == self.from_cs_id {
            // No adjustment needed
            return;
        }

        info!(
            ctx.logger(),
            concat!(
                "trying to adjust first batch for bookmark {} - ",
                "first batch starts points to {:?} but server points to {:?}",
            ),
            self.bookmark_name,
            self.from_cs_id,
            self.server_old_value,
        );

        let mut found = false;
        let mut entries_to_skip = vec![];
        for entry in &self.entries {
            entries_to_skip.push(entry.id);
            if self.server_old_value == entry.to_changeset_id {
                found = true;
                break;
            }
        }

        if found {
            warn!(
                ctx.logger(),
                "adjusting first batch - skipping first entries: {:?}", entries_to_skip
            );
            self.entries.splice(..entries_to_skip.len(), None);
            self.from_cs_id = self.server_old_value;
        } else {
            // The server bookmark doesn't point to any commit from the batch. This might
            // be expected in a repo with high commit rate.
            warn!(
                ctx.logger(),
                concat!(
                    "could not adjust first batch - ",
                    "the server bookmark ({:?}) does not point to any commit in the batch",
                ),
                self.server_old_value
            );
        }
    }
}

#[cfg(test)]
mod test {

    use bonsai_hg_mapping::BonsaiHgMapping;
    use bookmarks::Bookmarks;
    use changeset_fetcher::ChangesetFetcher;
    use changesets::Changesets;
    use fbinit::FacebookInit;
    use filestore::FilestoreConfig;
    use maplit::hashmap;
    use mononoke_types::RepositoryId;
    use repo_blobstore::RepoBlobstore;
    use repo_derived_data::RepoDerivedData;
    use tests_utils::drawdag::create_from_dag;

    use super::*;

    #[facet::container]
    #[derive(Clone)]
    pub struct TestRepo {
        #[facet]
        pub repo_blobstore: RepoBlobstore,

        #[facet]
        pub changesets: dyn Changesets,

        #[facet]
        pub changeset_fetcher: dyn ChangesetFetcher,

        #[facet]
        pub bonsai_hg_mapping: dyn BonsaiHgMapping,

        #[facet]
        pub bookmarks: dyn Bookmarks,

        #[facet]
        pub filestore_config: FilestoreConfig,

        #[facet]
        pub repo_derived_data: RepoDerivedData,

        #[facet]
        pub commit_graph: CommitGraph,
    }

    #[fbinit::test]
    async fn test_split_in_batches_simple(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let commits = create_from_dag(
            &ctx,
            &repo,
            r##"
                A-B-C
            "##,
        )
        .await?;

        let main = BookmarkKey::new("main")?;
        let commit = commits.get("A").cloned().unwrap();
        let entries = vec![create_bookmark_log_entry(
            0,
            main.clone(),
            None,
            Some(commit),
        )];
        let overlay = Arc::new(Mutex::new(BookmarkOverlay::new(Arc::new(hashmap! {}))));
        let res = split_in_batches(
            &ctx,
            repo.commit_graph(),
            overlay.clone(),
            entries.clone(),
            1000,
        )
        .await?;

        assert_eq!(res.len(), 1);
        assert_eq!(res[0].entries, entries);
        assert_eq!(res[0].bookmark_name, main);
        assert_eq!(res[0].from_cs_id, None);
        assert_eq!(res[0].to_cs_id, Some(commit));

        // The overlay should've been mutated so that the bookmark points to
        // the latest value.
        let overlay = overlay.lock().await;
        assert_eq!(overlay.get(&main), Some(commit));

        Ok(())
    }

    #[fbinit::test]
    async fn test_split_in_batches_all_in_one_batch(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let commits = create_from_dag(
            &ctx,
            &repo,
            r##"
                A-B-C
            "##,
        )
        .await?;

        let main = BookmarkKey::new("main")?;
        let commit_a = commits.get("A").cloned().unwrap();
        let commit_b = commits.get("B").cloned().unwrap();
        let commit_c = commits.get("C").cloned().unwrap();
        let entries = vec![
            create_bookmark_log_entry(0, main.clone(), None, Some(commit_a)),
            create_bookmark_log_entry(1, main.clone(), Some(commit_a), Some(commit_b)),
            create_bookmark_log_entry(2, main.clone(), Some(commit_b), Some(commit_c)),
        ];
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {}));
        let res = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?;

        assert_eq!(res.len(), 1);
        assert_eq!(res[0].entries, entries);
        assert_eq!(res[0].bookmark_name, main);
        assert_eq!(res[0].from_cs_id, None);
        assert_eq!(res[0].to_cs_id, Some(commit_c));

        Ok(())
    }

    #[fbinit::test]
    async fn test_split_in_batches_commit_limit(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let commits = create_from_dag(
            &ctx,
            &repo,
            r##"
                A-B-C
            "##,
        )
        .await?;

        let main = BookmarkKey::new("main")?;
        let commit_a = commits.get("A").cloned().unwrap();
        let commit_b = commits.get("B").cloned().unwrap();
        let commit_c = commits.get("C").cloned().unwrap();
        let entries = vec![
            create_bookmark_log_entry(0, main.clone(), None, Some(commit_a)),
            create_bookmark_log_entry(1, main.clone(), Some(commit_a), Some(commit_b)),
            create_bookmark_log_entry(2, main.clone(), Some(commit_b), Some(commit_c)),
        ];
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {}));
        let res = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            2,
        )
        .await?;

        assert_eq!(res.len(), 2);
        assert_eq!(res[0].entries[..], entries[..=1]);
        assert_eq!(res[0].bookmark_name, main);
        assert_eq!(res[0].from_cs_id, None);
        assert_eq!(res[0].to_cs_id, Some(commit_b));

        assert_eq!(res[1].entries[..], entries[2..]);
        assert_eq!(res[1].bookmark_name, main);
        assert_eq!(res[1].from_cs_id, Some(commit_b));
        assert_eq!(res[1].to_cs_id, Some(commit_c));

        Ok(())
    }

    #[fbinit::test]
    async fn test_split_in_batches_different_bookmarks(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let commits = create_from_dag(
            &ctx,
            &repo,
            r##"
                A-B-C
            "##,
        )
        .await?;

        let main = BookmarkKey::new("main")?;
        let another = BookmarkKey::new("another")?;
        let commit_a = commits.get("A").cloned().unwrap();
        let commit_b = commits.get("B").cloned().unwrap();
        let commit_c = commits.get("C").cloned().unwrap();
        let log_entry_1 = create_bookmark_log_entry(0, main.clone(), None, Some(commit_a));
        let log_entry_2 = create_bookmark_log_entry(1, another.clone(), None, Some(commit_b));
        let log_entry_3 =
            create_bookmark_log_entry(2, main.clone(), Some(commit_a), Some(commit_c));
        let entries = vec![
            log_entry_1.clone(),
            log_entry_2.clone(),
            log_entry_3.clone(),
        ];
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {}));
        let res = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?;

        assert_eq!(res.len(), 3);
        assert_eq!(res[0].entries, vec![log_entry_1]);
        assert_eq!(res[0].bookmark_name, main);
        assert_eq!(res[0].from_cs_id, None);
        assert_eq!(res[0].to_cs_id, Some(commit_a));

        assert_eq!(res[1].entries, vec![log_entry_2]);
        assert_eq!(res[1].bookmark_name, another);
        assert_eq!(res[1].from_cs_id, None);
        assert_eq!(res[1].to_cs_id, Some(commit_b));

        assert_eq!(res[2].entries, vec![log_entry_3]);
        assert_eq!(res[2].bookmark_name, main);
        assert_eq!(res[2].from_cs_id, Some(commit_a));
        assert_eq!(res[2].to_cs_id, Some(commit_c));

        Ok(())
    }

    #[fbinit::test]
    async fn test_split_in_batches_non_forward_move(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let commits = create_from_dag(
            &ctx,
            &repo,
            r##"
                A-B-C
            "##,
        )
        .await?;

        let main = BookmarkKey::new("main")?;
        let commit_a = commits.get("A").cloned().unwrap();
        let commit_b = commits.get("B").cloned().unwrap();
        let commit_c = commits.get("C").cloned().unwrap();
        let log_entry_1 = create_bookmark_log_entry(0, main.clone(), None, Some(commit_a));
        let log_entry_2 =
            create_bookmark_log_entry(1, main.clone(), Some(commit_a), Some(commit_c));
        let log_entry_3 =
            create_bookmark_log_entry(2, main.clone(), Some(commit_c), Some(commit_b));
        let entries = vec![
            log_entry_1.clone(),
            log_entry_2.clone(),
            log_entry_3.clone(),
        ];
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {}));
        let res = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?;

        assert_eq!(res.len(), 2);
        assert_eq!(res[0].entries, vec![log_entry_1, log_entry_2]);
        assert_eq!(res[0].bookmark_name, main);
        assert_eq!(res[0].from_cs_id, None);
        assert_eq!(res[0].to_cs_id, Some(commit_c));

        assert_eq!(res[1].entries, vec![log_entry_3]);
        assert_eq!(res[1].bookmark_name, main);
        assert_eq!(res[1].from_cs_id, Some(commit_c));
        assert_eq!(res[1].to_cs_id, Some(commit_b));

        Ok(())
    }

    #[fbinit::test]
    async fn test_split_in_batches_weird_move(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let commits = create_from_dag(
            &ctx,
            &repo,
            r##"
                A-B-C
            "##,
        )
        .await?;

        let main = BookmarkKey::new("main")?;
        let commit_a = commits.get("A").cloned().unwrap();
        let commit_b = commits.get("B").cloned().unwrap();
        let commit_c = commits.get("C").cloned().unwrap();
        let log_entry_1 = create_bookmark_log_entry(0, main.clone(), None, Some(commit_a));
        let log_entry_2 =
            create_bookmark_log_entry(1, main.clone(), Some(commit_b), Some(commit_c));
        let entries = vec![log_entry_1.clone(), log_entry_2.clone()];
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {}));
        let res = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?;

        assert_eq!(res.len(), 2);
        assert_eq!(res[0].entries, vec![log_entry_1]);
        assert_eq!(res[0].bookmark_name, main);
        assert_eq!(res[0].from_cs_id, None);
        assert_eq!(res[0].to_cs_id, Some(commit_a));

        assert_eq!(res[1].entries, vec![log_entry_2]);
        assert_eq!(res[1].bookmark_name, main);
        assert_eq!(res[1].from_cs_id, Some(commit_b));
        assert_eq!(res[1].to_cs_id, Some(commit_c));

        Ok(())
    }

    #[fbinit::test]
    async fn test_maybe_adjust_batch(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let repo: TestRepo = test_repo_factory::build_empty(ctx.fb).await?;

        let commits = create_from_dag(
            &ctx,
            &repo,
            r##"
                A-B-C
                D
            "##,
        )
        .await?;

        let main = BookmarkKey::new("main")?;
        let commit_a = commits.get("A").cloned().unwrap();
        let commit_b = commits.get("B").cloned().unwrap();
        let commit_c = commits.get("C").cloned().unwrap();
        let commit_d = commits.get("D").cloned().unwrap();
        let log_entry_1 = create_bookmark_log_entry(0, main.clone(), None, Some(commit_a));
        let log_entry_2 =
            create_bookmark_log_entry(1, main.clone(), Some(commit_a), Some(commit_b));
        let log_entry_3 =
            create_bookmark_log_entry(1, main.clone(), Some(commit_b), Some(commit_c));
        let entries = vec![log_entry_1, log_entry_2.clone(), log_entry_3.clone()];

        // Default case: no adjustment
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {}));
        let mut batch = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?
        .into_iter()
        .next()
        .unwrap();
        let original = batch.clone();
        batch.maybe_adjust(&ctx);
        assert_eq!(batch, original);

        // Skip a single entry
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {
          main.clone() => commit_a,
        }));
        let mut batch = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?
        .into_iter()
        .next()
        .unwrap();
        let original = batch.clone();
        batch.maybe_adjust(&ctx);
        assert_ne!(batch, original);
        assert_eq!(batch.from_cs_id, Some(commit_a));
        assert_eq!(batch.to_cs_id, Some(commit_c));
        assert_eq!(batch.entries, vec![log_entry_2, log_entry_3.clone()]);

        // Skip two entries
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {
          main.clone() => commit_b,
        }));
        let mut batch = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?
        .into_iter()
        .next()
        .unwrap();
        let original = batch.clone();
        batch.maybe_adjust(&ctx);
        assert_ne!(batch, original);
        assert_eq!(batch.from_cs_id, Some(commit_b));
        assert_eq!(batch.to_cs_id, Some(commit_c));
        assert_eq!(batch.entries, vec![log_entry_3]);

        // The whole batch was already synced - nothing to do!
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {
          main.clone() => commit_c,
        }));
        let mut batch = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?
        .into_iter()
        .next()
        .unwrap();
        let original = batch.clone();
        batch.maybe_adjust(&ctx);
        assert_ne!(batch, original);
        assert!(batch.is_empty());

        // Bookmark is not in the batch at all - in that case just do nothing and
        // return existing bundle
        let overlay = BookmarkOverlay::new(Arc::new(hashmap! {
          main => commit_d,
        }));
        let mut batch = split_in_batches(
            &ctx,
            repo.commit_graph(),
            Arc::new(Mutex::new(overlay)),
            entries.clone(),
            1000,
        )
        .await?
        .into_iter()
        .next()
        .unwrap();
        let original = batch.clone();
        batch.maybe_adjust(&ctx);
        assert_eq!(batch, original);
        Ok(())
    }

    fn create_bookmark_log_entry(
        id: i64,
        bookmark_name: BookmarkKey,
        from_changeset_id: Option<ChangesetId>,
        to_changeset_id: Option<ChangesetId>,
    ) -> BookmarkUpdateLogEntry {
        BookmarkUpdateLogEntry {
            id,
            repo_id: RepositoryId::new(0),
            bookmark_name,
            from_changeset_id,
            to_changeset_id,
            reason: BookmarkUpdateReason::TestMove,
            timestamp: Timestamp::now(),
        }
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use anyhow::Context;
use anyhow::Error;
use blobstore::Blobstore;
use context::CoreContext;
use filestore::FetchKey;
use filestore::FilestoreConfig;
use filestore::StoreRequest;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mononoke_types::hash::Sha256;

#[derive(Clone)]
pub struct DarkstormVerifier {
    origin_blobstore: Arc<dyn Blobstore>,
    backup_blobstore: Arc<dyn Blobstore>,
    backup_filestore_config: FilestoreConfig,
}

impl DarkstormVerifier {
    pub fn new(
        origin_blobstore: Arc<dyn Blobstore>,
        backup_blobstore: Arc<dyn Blobstore>,
        backup_filestore_config: FilestoreConfig,
    ) -> Self {
        Self {
            origin_blobstore,
            backup_blobstore,
            backup_filestore_config,
        }
    }

    pub async fn upload(&self, ctx: &CoreContext, blobs: &[(Sha256, u64)]) -> Result<(), Error> {
        stream::iter(blobs.iter().copied())
            .map(Ok)
            .try_for_each_concurrent(50, async move |(key, size)| -> Result<(), Error> {
                let blob = filestore::fetch_stream(
                    self.origin_blobstore.clone(),
                    ctx,
                    FetchKey::from(key),
                );
                let request = StoreRequest::with_sha256(size, key);
                filestore::store(
                    &self.backup_blobstore.clone(),
                    self.backup_filestore_config,
                    ctx,
                    &request,
                    blob,
                )
                .await
                .with_context(|| format!("while syncing LFS entry {:?}", key))?;
                Ok(())
            })
            .await?;
        Ok(())
    }
}

#[cfg(test)]
mod test {
    #![allow(warnings)]
    use fbinit::FacebookInit;
    use futures::stream;
    use futures::TryFutureExt;
    use futures::TryStreamExt;
    use memblob::Memblob;

    use super::*;

    #[fbinit::test]
    async fn test_upload(fb: FacebookInit) -> Result<(), anyhow::Error> {
        let ctx = CoreContext::test_mock(fb);

        let origin = Arc::new(Memblob::default());
        let backup = Arc::new(Memblob::default());
        let filestore = FilestoreConfig::no_chunking_filestore();

        let verifier = DarkstormVerifier::new(origin.clone(), backup.clone(), filestore);

        let bytes = b"hello world";
        let size = bytes.len().try_into().unwrap();
        let req = StoreRequest::new(size);

        let res = filestore::store(
            &origin,
            filestore,
            &ctx,
            &req,
            stream::once(async { Ok(bytes::Bytes::from(&bytes[..])) }),
        )
        .await?;

        let alias = res.sha256;

        verifier.upload(&ctx, &vec![(alias, size)][..]).await?;

        let uploaded_bytes = filestore::fetch(backup, ctx, &FetchKey::from(alias))
            .map_ok(|maybe_stream| async move {
                let res: Vec<bytes::Bytes> = maybe_stream.unwrap().try_collect().await?;
                Result::<_, Error>::Ok(res)
            })
            .try_flatten()
            .await?;
        assert_eq!(vec![bytes::Bytes::from(&bytes[..])], uploaded_bytes);
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use anyhow::Error;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkUpdateLogEntry;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("unexpected bookmark move: {0}")]
    UnexpectedBookmarkMove(String),
    #[error("sync failed for ids {ids:?}")]
    SyncFailed {
        ids: Vec<i64>,
        #[source]
        cause: Error,
    },
    #[error(
        "Programming error: cannot prepare a bundle for entry ids #{ids:?}: \
        entry {entry_id} modifies bookmark {entry_bookmark_name}, while bundle moves {bundle_bookmark_name}"
    )]
    BookmarkMismatchInBundleCombining {
        ids: Vec<i64>,
        entry_id: i64,
        entry_bookmark_name: BookmarkKey,
        bundle_bookmark_name: BookmarkKey,
    },
}

#[derive(Debug, Error)]
pub enum PipelineError {
    #[error("error processing entries {entries:?}")]
    EntryError {
        entries: Vec<BookmarkUpdateLogEntry>,
        #[source]
        cause: Error,
    },

    #[allow(dead_code)]
    #[error("error without tracking entry")]
    AnonymousError {
        #[source]
        cause: Error,
    },
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

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

use anyhow::format_err;
use anyhow::Error;
use bonsai_globalrev_mapping::BonsaiGlobalrevMappingEntry;
use bonsai_globalrev_mapping::BonsaiGlobalrevMappingRef;
use cloned::cloned;
use context::CoreContext;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use sql::Connection;
use sql_construct::SqlConstruct;
use sql_ext::SqlConnections;

use crate::CommitsInBundle;
use crate::Repo;

#[derive(Clone)]
pub enum GlobalrevSyncer {
    Noop,
    Darkstorm(Arc<DarkstormGlobalrevSyncer>),
}

pub struct DarkstormGlobalrevSyncer {
    orig_repo: Repo,
    darkstorm_repo: Repo,
}

#[derive(Clone)]
struct HgsqlConnection {
    #[allow(dead_code)]
    connection: Connection,
}

impl SqlConstruct for HgsqlConnection {
    const LABEL: &'static str = "globalrev-syncer";

    const CREATION_QUERY: &'static str = include_str!("../schemas/hgsql.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self {
            connection: connections.write_connection,
        }
    }
}

impl GlobalrevSyncer {
    pub fn darkstorm(orig_repo: &Repo, darkstorm_repo: &Repo) -> Self {
        Self::Darkstorm(Arc::new(DarkstormGlobalrevSyncer {
            orig_repo: orig_repo.clone(),
            darkstorm_repo: darkstorm_repo.clone(),
        }))
    }

    pub async fn sync(&self, ctx: &CoreContext, commits: &CommitsInBundle) -> Result<(), Error> {
        match self {
            Self::Noop => Ok(()),
            Self::Darkstorm(syncer) => syncer.sync(ctx, commits).await,
        }
    }
}

impl DarkstormGlobalrevSyncer {
    pub async fn sync(&self, ctx: &CoreContext, commits: &CommitsInBundle) -> Result<(), Error> {
        let commits = match commits {
            CommitsInBundle::Commits(commits) => commits,
            CommitsInBundle::Unknown => {
                return Err(format_err!(
                    "can't use darkstorm globalrev syncer because commits that were \
                    sent in the bundle are not known"
                ));
            }
        };

        let commits = commits.clone().into_iter().map(|(_, bcs_id)| bcs_id);
        let bcs_id_to_globalrev = stream::iter(commits.map(|bcs_id| {
            cloned!(ctx, self.orig_repo);
            async move {
                let maybe_globalrev = orig_repo
                    .bonsai_globalrev_mapping()
                    .get_globalrev_from_bonsai(&ctx, bcs_id)
                    .await?;
                Result::<_, Error>::Ok((bcs_id, maybe_globalrev))
            }
        }))
        .map(Ok)
        .try_buffer_unordered(100)
        .try_filter_map(|(bcs_id, maybe_globalrev)| {
            let bcs_id = bcs_id.clone();
            let maybe_globalrev = maybe_globalrev.clone();
            async move { Ok(maybe_globalrev.map(|globalrev| (bcs_id, globalrev))) }
        })
        .try_collect::<HashMap<_, _>>()
        .await?;

        let entries = bcs_id_to_globalrev
            .into_iter()
            .map(|(bcs_id, globalrev)| BonsaiGlobalrevMappingEntry { bcs_id, globalrev })
            .collect::<Vec<_>>();

        self.darkstorm_repo
            .bonsai_globalrev_mapping()
            .bulk_import(ctx, &entries[..])
            .await?;

        Ok(())
    }
}

#[cfg(test)]
mod test {
    use bonsai_globalrev_mapping::BonsaiGlobalrevMappingEntry;
    use fbinit::FacebookInit;
    use mercurial_types_mocks::globalrev::GLOBALREV_ONE;
    use mercurial_types_mocks::globalrev::GLOBALREV_TWO;
    use mercurial_types_mocks::nodehash::ONES_CSID as ONES_HG_CSID;
    use mercurial_types_mocks::nodehash::TWOS_CSID as TWOS_HG_CSID;
    use mononoke_types::RepositoryId;
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::changesetid::TWOS_CSID;
    use test_repo_factory::TestRepoFactory;

    use super::*;

    #[fbinit::test]
    async fn test_sync_darkstorm(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);

        let orig_repo: Repo = TestRepoFactory::new(fb)?
            .with_id(RepositoryId::new(0))
            .build()
            .await?;
        let darkstorm_repo: Repo = TestRepoFactory::new(fb)?
            .with_id(RepositoryId::new(1))
            .build()
            .await?;

        let e1 = BonsaiGlobalrevMappingEntry {
            bcs_id: ONES_CSID,
            globalrev: GLOBALREV_ONE,
        };

        let e2 = BonsaiGlobalrevMappingEntry {
            bcs_id: TWOS_CSID,
            globalrev: GLOBALREV_TWO,
        };
        orig_repo
            .bonsai_globalrev_mapping()
            .bulk_import(&ctx, &[e1, e2])
            .await?;

        let syncer = DarkstormGlobalrevSyncer {
            orig_repo,
            darkstorm_repo: darkstorm_repo.clone(),
        };

        assert!(
            darkstorm_repo
                .bonsai_globalrev_mapping()
                .get_globalrev_from_bonsai(&ctx, ONES_CSID)
                .await?
                .is_none()
        );
        assert!(
            darkstorm_repo
                .bonsai_globalrev_mapping()
                .get_globalrev_from_bonsai(&ctx, TWOS_CSID)
                .await?
                .is_none()
        );
        syncer
            .sync(
                &ctx,
                &CommitsInBundle::Commits(vec![
                    (ONES_HG_CSID, ONES_CSID),
                    (TWOS_HG_CSID, TWOS_CSID),
                ]),
            )
            .await?;

        assert_eq!(
            Some(GLOBALREV_ONE),
            darkstorm_repo
                .bonsai_globalrev_mapping()
                .get_globalrev_from_bonsai(&ctx, ONES_CSID)
                .await?
        );
        assert_eq!(
            Some(GLOBALREV_TWO),
            darkstorm_repo
                .bonsai_globalrev_mapping()
                .get_globalrev_from_bonsai(&ctx, TWOS_CSID)
                .await?
        );
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::collections::HashMap;
use std::ffi::OsStr;
use std::fs;
use std::fs::File;
use std::process::Stdio;
use std::str::FromStr;
use std::sync::Arc;
use std::time::Duration;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context as _;
use anyhow::Error;
use anyhow::Result;
use bookmarks::BookmarkKey;
use futures::future;
use futures::future::FutureExt;
use futures::future::TryFuture;
use futures::future::TryFutureExt;
use futures_ext::future::FbFutureExt;
use futures_ext::future::FbTryFutureExt;
use futures_watchdog::WatchdogExt;
use itertools::Itertools;
use mercurial_types::HgChangesetId;
use mononoke_hg_sync_job_helper_lib::lines_after;
use mononoke_hg_sync_job_helper_lib::read_file_contents;
use mononoke_hg_sync_job_helper_lib::wait_till_more_lines;
use mononoke_hg_sync_job_helper_lib::write_to_named_temp_file;
use slog::debug;
use slog::info;
use slog::Logger;
use tempfile::NamedTempFile;
use tokio::io::AsyncWriteExt;
use tokio::process::Child;
use tokio::process::ChildStdin;
use tokio::process::Command;
use tokio::sync::Mutex;

use crate::CommitsInBundle;

const BOOKMARK_LOCATION_LOOKUP_TIMEOUT_MS: u64 = 10_000;
const LIST_SERVER_BOOKMARKS_EXTENSION: &str = include_str!("listserverbookmarks.py");
const SEND_UNBUNDLE_REPLAY_EXTENSION: &str = include_str!("sendunbundlereplay.py");

pub async fn list_hg_server_bookmarks(
    hg_repo_path: String,
) -> Result<HashMap<BookmarkKey, HgChangesetId>, Error> {
    let extension_file = tempfile::Builder::new()
        .suffix(".py")
        .tempfile()
        .context("Error in creating extension file")?;
    let file_path = extension_file
        .path()
        .to_str()
        .ok_or_else(|| Error::msg("Temp file path contains non-unicode chars"))?;
    fs::write(file_path, LIST_SERVER_BOOKMARKS_EXTENSION)?;
    let ext = format!("extensions.listserverbookmarks={}", file_path);

    let full_args = vec![
        "--config",
        &ext,
        "listserverbookmarks",
        "--path",
        &hg_repo_path,
    ];

    let output = get_hg_command(&full_args)
        .output()
        .await
        .context("Error listing server bookmarks")?;

    let mut res = HashMap::new();
    for keyvalue in output.stdout.split(|x| x == &0) {
        if keyvalue.is_empty() {
            continue;
        }
        let mut iter = keyvalue.split(|x| x == &1);
        match (iter.next(), iter.next()) {
            (Some(key), Some(value)) => {
                let key = String::from_utf8(key.to_vec()).map_err(Error::from)?;
                let value = String::from_utf8(value.to_vec()).map_err(Error::from)?;
                res.insert(BookmarkKey::new(key)?, HgChangesetId::from_str(&value)?);
            }
            _ => {
                let stdout = String::from_utf8_lossy(&output.stdout);
                bail!("invalid format returned from server: {}", stdout);
            }
        }
    }

    Ok(res)
}

fn expected_location_string_arg(maybe_hgcsid: Option<HgChangesetId>) -> String {
    match maybe_hgcsid {
        Some(hash) => hash.to_string(),
        None => "DELETED".into(),
    }
}

fn get_hg_command<I, S>(args: I) -> Command
where
    I: IntoIterator<Item = S>,
    S: AsRef<OsStr>,
{
    let full_args = vec![
        "--traceback",
        "--config",
        "extensions.clienttelemetry=",
        "--config",
        "clienttelemetry.announceremotehostname=on",
    ]
    .into_iter()
    .map(|item| item.into())
    .chain(args.into_iter().map(|item| item.as_ref().to_os_string()));
    let mut child = Command::new("hg");
    child.args(full_args);
    child
}

struct AsyncProcess {
    child: Child,
    stdin: ChildStdin,
}

impl AsyncProcess {
    pub fn new<I, S>(args: I) -> Result<Self>
    where
        I: IntoIterator<Item = S>,
        S: AsRef<OsStr>,
    {
        Self::from_command(get_hg_command(args))
    }

    fn from_command(mut command: Command) -> Result<Self> {
        let mut child = command
            .stdin(Stdio::piped())
            .spawn()
            .context("Failed to spawn hg command")?;
        let stdin = child
            .stdin
            .take()
            .ok_or_else(|| Error::msg("ChildStdin unexpectedly not captured"))?;
        Ok(Self { child, stdin })
    }

    pub async fn write_line(&mut self, line: Vec<u8>) -> Result<(), Error> {
        self.stdin
            .write_all(&line)
            .await
            .context("Failed to write")?;
        Ok(())
    }

    pub fn is_valid(&mut self) -> bool {
        !self.child_is_finished()
    }

    pub async fn kill(&mut self, logger: &Logger) {
        if let Err(e) = self.child.kill().await {
            debug!(logger, "failed to kill the hg process: {}", e);
        }
    }

    fn child_is_finished(&mut self) -> bool {
        self.child.try_wait().transpose().is_some()
    }

    /// Make sure child is still alive while provided future is being executed
    /// If `grace_period` is specified, future will be given additional time
    /// to resolve even if peer has already been terminated.
    pub async fn ensure_alive<F>(
        &mut self,
        fut: F,
        grace_period: Option<Duration>,
    ) -> Result<<F as TryFuture>::Ok, <F as TryFuture>::Error>
    where
        F: TryFuture<Error = Error> + Unpin + Send,
    {
        if self.child_is_finished() {
            return Err(Error::msg("hg peer is dead"));
        }

        let child = self.child.wait();

        let watchdog = async {
            let res = child.await;
            if let Some(grace_period) = grace_period {
                tokio::time::sleep(grace_period).await;
            }
            Err(format_err!("hg peer has died unexpectedly: {:?}", res))
        }
        .boxed();

        // NOTE: Right can't actually return an Ok variant here (which is why this compiles at
        // all), but it still needs to be in the match clause.
        match future::try_select(fut, watchdog).await {
            Ok(future::Either::Left((res, _))) | Ok(future::Either::Right((res, _))) => Ok(res),
            Err(future::Either::Left((e, _))) | Err(future::Either::Right((e, _))) => Err(e),
        }
    }
}

struct HgPeer {
    process: AsyncProcess,
    reports_file: Arc<NamedTempFile>,
    bundle_applied: usize,
    max_bundles_allowed: usize,
    baseline_bundle_timeout_ms: u64,
    invalidated: bool,
    // The extension_file needs to be kept around while we have running instances of the process.
    #[allow(unused)]
    extension_file: Arc<File>,
}

impl HgPeer {
    pub fn new(
        repo_path: &str,
        max_bundles_allowed: usize,
        baseline_bundle_timeout_ms: u64,
    ) -> Result<Self> {
        let reports_file = NamedTempFile::new()?;
        let file_path = reports_file
            .path()
            .to_str()
            .ok_or_else(|| Error::msg("Temp file path contains non-unicode chars"))?;

        let extension_file = tempfile::Builder::new()
            .suffix(".py")
            .tempfile()
            .context("Error in creating extension file")?;
        // Persisting the file so it does not get deleted before its contents are read.
        let (extension_file, extension_path) = extension_file.keep()?;
        let path_string = extension_path.clone();
        fs::write(extension_path, SEND_UNBUNDLE_REPLAY_EXTENSION)
            .with_context(|| format!("Error in writing data to file {}", &path_string.display()))?;
        let args = &[
            "--config",
            &format!("extensions.sendunbundlereplay={}", &path_string.display()),
            "sendunbundlereplaybatch",
            "--debug",
            "--path",
            repo_path,
            "--reports",
            file_path,
        ];
        let process = AsyncProcess::new(args)?;

        Ok(HgPeer {
            process,
            reports_file: Arc::new(reports_file),
            bundle_applied: 0,
            max_bundles_allowed,
            baseline_bundle_timeout_ms,
            invalidated: false,
            extension_file: Arc::new(extension_file),
        })
    }

    pub fn arc_mutexed(self) -> Arc<Mutex<Self>> {
        Arc::new(Mutex::new(self))
    }

    pub fn still_good(&mut self, logger: &Logger) -> bool {
        let can_be_used: bool = !self.invalidated & self.process.is_valid();
        let can_write_more = self.bundle_applied < self.max_bundles_allowed;
        debug!(
            logger,
            "can be used: {}, bundle_applied: {}, max bundles allowed: {}",
            can_be_used,
            self.bundle_applied,
            self.max_bundles_allowed
        );
        can_be_used && can_write_more
    }

    pub async fn kill(&mut self, logger: &Logger) {
        self.process.kill(logger).await;
    }

    pub async fn apply_bundle<'a>(
        &'a mut self,
        bundle_path: &'a str,
        timestamps_path: &'a str,
        onto_bookmark: BookmarkKey,
        expected_bookmark_position: Option<HgChangesetId>,
        attempt: usize,
        logger: &Logger,
        commits_in_bundle: &CommitsInBundle,
    ) -> Result<(), Error> {
        let mut log_file = match NamedTempFile::new() {
            Ok(log_file) => log_file,
            Err(e) => {
                return Err(format_err!("could not create log file: {:?}", e));
            }
        };

        let log_path = match log_file.path().to_str() {
            Some(log_path) => log_path,
            None => {
                return Err(Error::msg("log_file path was not a valid string"));
            }
        };

        let onto_bookmark = onto_bookmark.to_string();
        let onto_bookmark = base64::encode(&onto_bookmark);
        let expected_hash = expected_location_string_arg(expected_bookmark_position);

        let hgbonsaimapping = match commits_in_bundle {
            CommitsInBundle::Commits(hgbonsaimapping) => {
                let encoded_hg_bonsai_mapping = hgbonsaimapping
                    .iter()
                    .map(|(hg_cs_id, bcs_id)| format!("{}={}", hg_cs_id, bcs_id))
                    .join("\n");

                encoded_hg_bonsai_mapping
            }
            CommitsInBundle::Unknown => "".to_string(),
        };

        let hgbonsaimappingfile = write_to_named_temp_file(hgbonsaimapping)
            .watched(logger)
            .await?;
        let hgbonsaimappingpath = match hgbonsaimappingfile.path().to_str() {
            Some(path) => path,
            None => {
                return Err(Error::msg("hg bonsai mapping path was not a valid string"));
            }
        };

        let input_line = format!(
            "{} {} {} {} {} {}\n",
            bundle_path,
            timestamps_path,
            hgbonsaimappingpath,
            onto_bookmark,
            expected_hash,
            log_path,
        );
        let path = self.reports_file.path().to_path_buf();
        let bundle_timeout_ms = self.baseline_bundle_timeout_ms * 2_u64.pow(attempt as u32 - 1);

        let line_num_in_reports_file = lines_after(&path, 0).watched(logger).await?.len();

        let res = async {
            self.process
                .write_line(input_line.into_bytes())
                .watched(logger)
                .await?;
            self.bundle_applied += 1;

            let report_lines = self
                .process
                .ensure_alive(
                    wait_till_more_lines(path, line_num_in_reports_file, bundle_timeout_ms).boxed(),
                    // even if peer process has died, lets wait for additional grace
                    // period, and try to collect the report if any.
                    Some(Duration::from_secs(1)),
                )
                .watched(logger)
                .await?;

            let full_report = report_lines.join("\n");
            let success = !full_report.contains("failed");
            debug!(logger, "sync report: {}", full_report);

            if success {
                return Ok(());
            }

            let log = match read_file_contents(&mut log_file) {
                Ok(log) => format!("hg logs follow:\n{}", log),
                Err(e) => format!("no hg logs available ({:?})", e),
            };

            Err(format_err!("sync failed: {}", log))
        }
        .watched(logger)
        .await;

        if res.is_err() {
            info!(logger, "sync failed. Invalidating process");
            self.invalidated = true;
        }

        res
    }
}

/// Struct that knows how to work with on-disk mercurial repository.
/// It shells out to `hg` cmd line tool.
#[derive(Clone)]
pub struct HgRepo {
    repo_path: Arc<String>,
    peer: Arc<Mutex<HgPeer>>,
    max_bundles_per_peer: usize,
    baseline_bundle_timeout_ms: u64,
    verify_server_bookmark_on_failure: bool,
}

impl HgRepo {
    pub fn new(
        repo_path: String,
        max_bundles_per_peer: usize,
        baseline_bundle_timeout_ms: u64,
        verify_server_bookmark_on_failure: bool,
    ) -> Result<Self> {
        let peer = HgPeer::new(&repo_path, max_bundles_per_peer, baseline_bundle_timeout_ms)?;
        Ok(Self {
            repo_path: Arc::new(repo_path),
            peer: peer.arc_mutexed(),
            max_bundles_per_peer,
            baseline_bundle_timeout_ms,
            verify_server_bookmark_on_failure,
        })
    }

    pub async fn apply_bundle(
        &self,
        bundle_filename: String,
        timestamps_path: String,
        onto_bookmark: BookmarkKey,
        expected_bookmark_position: Option<HgChangesetId>,
        attempt: usize,
        logger: &Logger,
        commits_in_bundle: &CommitsInBundle,
    ) -> Result<(), Error> {
        self.renew_peer_if_needed(logger).await?;

        let mut peer = self.peer.lock().await;

        let res = peer
            .apply_bundle(
                &bundle_filename,
                &timestamps_path,
                onto_bookmark.clone(),
                expected_bookmark_position.clone(),
                attempt,
                logger,
                commits_in_bundle,
            )
            .watched(logger.clone())
            .await;

        let err = match res {
            Ok(()) => return Ok(()),
            Err(err) => err,
        };

        if !self.verify_server_bookmark_on_failure {
            return Err(err);
        }

        info!(
            logger,
            "sync failed, let's check if the bookmark is where we want it to be anyway",
        );

        if self
            .verify_server_bookmark_location(&onto_bookmark, expected_bookmark_position)
            .watched(logger.clone())
            .await
            .is_ok()
        {
            return Ok(());
        }

        Err(err)
    }

    async fn renew_peer_if_needed(&self, logger: &Logger) -> Result<()> {
        let mut peer = self.peer.lock().watched(logger).await;

        if peer.still_good(logger) {
            return Ok(debug!(logger, "existing hg peer is still good"));
        }

        debug!(logger, "killing the old peer");
        peer.kill(logger).await;

        debug!(logger, "renewing hg peer");
        let new_peer = HgPeer::new(
            &self.repo_path.clone(),
            self.max_bundles_per_peer,
            self.baseline_bundle_timeout_ms,
        )?;
        *peer = new_peer;
        Ok(debug!(logger, "done renewing hg peer"))
    }

    async fn verify_server_bookmark_location(
        &self,
        name: &BookmarkKey,
        expected_bookmark_position: Option<HgChangesetId>,
    ) -> Result<(), Error> {
        let name = name.to_string();
        let mut args: Vec<String> = [
            "checkserverbookmark",
            // NB: we can't enable extensions.checkserverbookmark universally until it
            //     is deployed as part of the package. For now, let it be enabled only when
            //     the appropriate command line flag is present (e.g. when this function is
            //     called)
            "--config",
            "extensions.checkserverbookmark=",
            "--path",
            &self.repo_path.clone(),
            "--name",
        ]
        .iter()
        .map(|item| item.to_string())
        .collect();
        args.push(name);
        match expected_bookmark_position {
            Some(hash) => {
                args.push("--hash".into());
                args.push(hash.to_string());
            }
            None => args.push("--deleted".into()),
        };

        let mut cmd = get_hg_command(args)
            .stdin(Stdio::piped())
            .spawn()
            .context("failed to start a mercurial process")?;

        let exit_status = cmd
            .wait()
            .map_err(Error::from)
            .timeout(Duration::from_millis(BOOKMARK_LOCATION_LOOKUP_TIMEOUT_MS))
            .flatten_err()
            .await?;

        if exit_status.success() {
            Ok(())
        } else {
            Err(Error::msg(
                "server does not have a bookmark in the expected location",
            ))
        }
    }
}

#[cfg(test)]
mod tests {
    use assert_matches::assert_matches;

    use super::*;

    #[tokio::test]
    async fn ensure_alive_alive_process() -> Result<()> {
        let command = {
            let mut command = Command::new("sleep");
            command.args(vec!["2"]);
            command
        };
        let mut proc = AsyncProcess::from_command(command)?;

        let res = proc
            .ensure_alive(
                async {
                    tokio::time::sleep(Duration::from_millis(100)).await;
                    Ok(())
                }
                .boxed(),
                None,
            )
            .await;
        assert_matches!(res, Ok(()));

        assert!(proc.is_valid());
        Ok(())
    }

    #[tokio::test]
    async fn ensure_alive_dead_process() -> Result<()> {
        let mut proc = AsyncProcess::from_command(Command::new("false"))?;

        // Give the command a little time to finish
        tokio::time::sleep(Duration::from_millis(100)).await;
        let res = proc
            .ensure_alive(
                async {
                    tokio::time::sleep(Duration::from_secs(5)).await;
                    Ok(())
                }
                .boxed(),
                None,
            )
            .await;
        assert_matches!(res, Err(_));
        assert!(res.unwrap_err().to_string().starts_with("hg peer is dead"));

        assert!(!proc.is_valid());
        Ok(())
    }

    #[tokio::test]
    async fn ensure_alive_no_grace_period() -> Result<()> {
        let command = {
            let mut command = Command::new("sleep");
            command.args(vec!["0.1"]);
            command
        };
        let mut proc = AsyncProcess::from_command(command)?;

        let res = proc
            .ensure_alive(
                async {
                    tokio::time::sleep(Duration::from_secs(5)).await;
                    Ok(())
                }
                .boxed(),
                None,
            )
            .await;
        assert_matches!(res, Err(_));
        assert!(res.unwrap_err().to_string().starts_with("hg peer has died"));

        assert!(!proc.is_valid());
        Ok(())
    }

    #[tokio::test]
    async fn ensure_alive_grace_period() -> Result<()> {
        let command = {
            let mut command = Command::new("sleep");
            command.args(vec!["0.1"]);
            command
        };
        let mut proc = AsyncProcess::from_command(command)?;

        let res = proc
            .ensure_alive(
                async {
                    tokio::time::sleep(Duration::from_secs(1)).await;
                    Ok(())
                }
                .boxed(),
                Some(Duration::from_secs(10)),
            )
            .await;
        assert_matches!(res, Ok(()));

        assert!(!proc.is_valid());
        Ok(())
    }
}


/*
 * Copyright (c) Meta Platforms, Inc. and affiliates.
 *
 * This software may be used and distributed according to the terms of the
 * GNU General Public License version 2.
 */

use std::sync::Arc;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobstore::Blobstore;
use bytes::Bytes;
use cloned::cloned;
use context::CoreContext;
use filestore::fetch_stream;
use filestore::FetchKey;
use futures::future;
use futures::pin_mut;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use gotham_ext::body_ext::BodyExt;
use http::status::StatusCode;
use http::uri::Uri;
use hyper::client::HttpConnector;
use hyper::Body;
use hyper::Client;
use hyper::Request;
use hyper_openssl::HttpsConnector;
use lfs_protocol::ObjectAction;
use lfs_protocol::ObjectStatus;
use lfs_protocol::Operation;
use lfs_protocol::RequestBatch;
use lfs_protocol::RequestObject;
use lfs_protocol::ResponseBatch;
use lfs_protocol::ResponseObject;
use lfs_protocol::Sha256 as LfsSha256;
use lfs_protocol::Transfer;
use mononoke_types::hash::Sha256;
use slog::info;
use slog::warn;
use thiserror::Error;

pub type HttpsHyperClient = Client<HttpsConnector<HttpConnector>>;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Serializing a LFS batch failed")]
    SerializationFailed,
    #[error("Deserializating a LFS batch failed")]
    DeserializationFailed,
    #[error("Creating a request failed")]
    RequestCreationFailed,
    #[error("Submitting a batch request failed")]
    BatchRequestNoResponse,
    #[error("Submitting a batch request failed with status {0}")]
    BatchRequestFailed(StatusCode),
    #[error("Reading the response for a batch request failed")]
    BatchRequestReadFailed,
    #[error("An error ocurred receiving a response from upstream ({0}): {1}")]
    UpstreamError(StatusCode, String),
}

struct LfsVerifierInner {
    client: HttpsHyperClient,
    batch_uri: Uri,
    blobstore: Arc<dyn Blobstore>,
}

#[derive(Clone)]
pub struct LfsVerifier {
    inner: Arc<LfsVerifierInner>,
}

impl LfsVerifier {
    pub fn new(batch_uri: Uri, blobstore: Arc<dyn Blobstore>) -> Result<Self> {
        let connector = HttpsConnector::new()?;
        let client = Client::builder().build(connector);

        let inner = LfsVerifierInner {
            batch_uri,
            client,
            blobstore,
        };

        Ok(Self {
            inner: Arc::new(inner),
        })
    }

    pub async fn ensure_lfs_presence<'a>(
        &'a self,
        ctx: &'a CoreContext,
        blobs: &'a [(Sha256, u64)],
    ) -> Result<()> {
        let uri = self.inner.batch_uri.clone();
        let blobstore = self.inner.blobstore.clone();
        let client = self.inner.client.clone();

        let batch = build_upload_request_batch(blobs);

        let body = Bytes::from(serde_json::to_vec(&batch).context(ErrorKind::SerializationFailed)?);

        let req = Request::post(uri)
            .body(body.into())
            .context(ErrorKind::RequestCreationFailed)?;

        let response = client
            .request(req)
            .await
            .context(ErrorKind::BatchRequestNoResponse)?;

        let (head, body) = response.into_parts();

        if !head.status.is_success() {
            return Result::<_, Error>::Err(ErrorKind::BatchRequestFailed(head.status).into())?;
        }

        let body = body
            .try_concat_body(&head.headers)
            .context(ErrorKind::BatchRequestReadFailed)?
            .await
            .context(ErrorKind::BatchRequestReadFailed)?;

        let batch = serde_json::from_slice::<ResponseBatch>(&body)
            .context(ErrorKind::DeserializationFailed)?;

        let missing_objects = find_missing_objects(batch);

        if missing_objects.is_empty() {
            return Ok(());
        }

        for object in &missing_objects {
            warn!(ctx.logger(), "missing {:?} object, uploading", object);
        }

        stream::iter(missing_objects)
            .map(|object| upload(ctx, &client, object, blobstore.clone()))
            .buffer_unordered(100)
            .try_for_each(|()| future::ready(Ok(())))
            .await?;

        Result::<_, Error>::Ok(())
    }
}

fn build_upload_request_batch(blobs: &[(Sha256, u64)]) -> RequestBatch {
    let objects = blobs
        .iter()
        .map(|(oid, size)| RequestObject {
            oid: LfsSha256(oid.into_inner()),
            size: *size,
        })
        .collect();

    RequestBatch {
        operation: Operation::Upload,
        r#ref: None,
        transfers: vec![Transfer::Basic],
        objects,
    }
}

fn find_missing_objects(batch: ResponseBatch) -> Vec<ResponseObject> {
    batch
        .objects
        .into_iter()
        .filter_map(|object| match object.status {
            ObjectStatus::Ok { ref actions, .. } if !actions.contains_key(&Operation::Upload) => {
                None
            }
            _ => Some(object),
        })
        .collect()
}

async fn upload<'a>(
    ctx: &'a CoreContext,
    client: &'a HttpsHyperClient,
    resp_object: ResponseObject,
    blobstore: Arc<dyn Blobstore>,
) -> Result<()> {
    match resp_object.status {
        ObjectStatus::Ok { actions, .. } => match actions.get(&Operation::Upload) {
            Some(action) => {
                let ObjectAction { href, .. } = action;

                let key = FetchKey::from(Sha256::from_byte_array(resp_object.object.oid.0));

                let s = {
                    cloned!(ctx);
                    async_stream::stream! {
                        let s = fetch_stream(
                            &blobstore,
                            ctx,
                            key,
                        );

                        pin_mut!(s);
                        while let Some(value) = s.next().await {
                            yield value;
                        }
                    }
                };

                let body = Body::wrap_stream(s.map_ok(|b| Bytes::copy_from_slice(b.as_ref())));

                let req = Request::put(format!("{}", href))
                    .header("Content-Length", &resp_object.object.size.to_string())
                    .body(body)?;

                let res = client.request(req).await?;
                let (head, body) = res.into_parts();

                let body = body.try_concat_body(&head.headers)?.await?;

                if !head.status.is_success() {
                    return Err(ErrorKind::UpstreamError(
                        head.status,
                        String::from_utf8_lossy(&body).to_string(),
                    )
                    .into());
                }

                info!(
                    ctx.logger(),
                    "uploaded content for {:?}", resp_object.object
                );

                Ok(())
            }
            None => Err(anyhow!(
                "not found upload action for {:?}",
                resp_object.object
            )),
        },
        ObjectStatus::Err { error } => Err(anyhow!(
            "batch failed for {:?} {:?}",
            resp_object.object,
            error
        )),
    }
}