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

#![feature(trait_alias)]

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

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Error;
use blobrepo::save_bonsai_changesets;
use blobrepo_utils::convert_diff_result_into_file_change_for_diamond_merge;
use blobstore::Loadable;
use blobsync::copy_content;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use bookmarks::BookmarksRef;
use borrowed::borrowed;
use changeset_fetcher::ChangesetFetcherArc;
use changesets::ChangesetsRef;
use cloned::cloned;
use commit_graph::CommitGraphRef;
use context::CoreContext;
use filestore::FilestoreConfigRef;
use futures::future::try_join_all;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use manifest::get_implicit_deletes;
use megarepo_configs::types::SourceMappingRules;
use mercurial_derivation::DeriveHgChangeset;
use mercurial_types::HgManifestId;
use mononoke_types::mpath_element_iter;
use mononoke_types::BonsaiChangeset;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use mononoke_types::ContentId;
use mononoke_types::FileChange;
use mononoke_types::MPath;
use mononoke_types::TrackedFileChange;
use pushrebase::find_bonsai_diff;
use repo_blobstore::RepoBlobstoreArc;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedDataRef;
use repo_identity::RepoIdentityRef;
use sorted_vector_map::SortedVectorMap;
use thiserror::Error;

pub type MultiMover = Arc<dyn Fn(&MPath) -> Result<Vec<MPath>, Error> + Send + Sync + 'static>;
pub type DirectoryMultiMover =
    Arc<dyn Fn(&Option<MPath>) -> Result<Vec<Option<MPath>>, Error> + Send + Sync + 'static>;

pub trait Repo = RepoIdentityRef
    + RepoBlobstoreArc
    + ChangesetsRef
    + ChangesetFetcherArc
    + BookmarksRef
    + BonsaiHgMappingRef
    + RepoDerivedDataRef
    + RepoBlobstoreRef
    + CommitGraphRef
    + Send
    + Sync;

const SQUASH_DELIMITER_MESSAGE: &str = r#"

============================

This commit created by squashing the following git commits:
"#;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Remapped commit {0} expected in target repo, but not present")]
    MissingRemappedCommit(ChangesetId),
    #[error(
        "Can't reoder changesets parents to put {0} first because it's not a changeset's parent."
    )]
    MissingForcedParent(ChangesetId),
}

pub fn create_source_to_target_multi_mover(
    mapping_rules: SourceMappingRules,
) -> Result<MultiMover, Error> {
    // We apply the longest prefix first
    let mut overrides = mapping_rules.overrides.into_iter().collect::<Vec<_>>();
    overrides.sort_unstable_by_key(|(ref prefix, _)| prefix.len());
    overrides.reverse();
    let prefix = MPath::new_opt(mapping_rules.default_prefix)?;

    Ok(Arc::new(move |path: &MPath| -> Result<Vec<MPath>, Error> {
        for (override_prefix_src, dsts) in &overrides {
            let override_prefix_src = MPath::new(override_prefix_src.clone())?;
            if override_prefix_src.is_prefix_of(path) {
                let suffix: Vec<_> = path
                    .into_iter()
                    .skip(override_prefix_src.num_components())
                    .collect();

                return dsts
                    .iter()
                    .map(|dst| {
                        let override_prefix = MPath::new_opt(dst)?;
                        MPath::join_opt(override_prefix.as_ref(), suffix.clone())
                            .ok_or_else(|| anyhow!("unexpected empty path"))
                    })
                    .collect::<Result<_, _>>();
            }
        }

        Ok(vec![
            MPath::join_opt(prefix.as_ref(), path)
                .ok_or_else(|| anyhow!("unexpected empty path"))?,
        ])
    }))
}

pub fn create_directory_source_to_target_multi_mover(
    mapping_rules: SourceMappingRules,
) -> Result<DirectoryMultiMover, Error> {
    // We apply the longest prefix first
    let mut overrides = mapping_rules.overrides.into_iter().collect::<Vec<_>>();
    overrides.sort_unstable_by_key(|(ref prefix, _)| prefix.len());
    overrides.reverse();
    let prefix = MPath::new_opt(mapping_rules.default_prefix)?;

    Ok(Arc::new(
        move |path: &Option<MPath>| -> Result<Vec<Option<MPath>>, Error> {
            for (override_prefix_src, dsts) in &overrides {
                let override_prefix_src = MPath::new(override_prefix_src.clone())?;
                if override_prefix_src.is_prefix_of(mpath_element_iter(path)) {
                    let suffix: Vec<_> = mpath_element_iter(path)
                        .skip(override_prefix_src.num_components())
                        .collect();

                    return dsts
                        .iter()
                        .map(|dst| {
                            let override_prefix = MPath::new_opt(dst)?;
                            Ok(MPath::join_opt(override_prefix.as_ref(), suffix.clone()))
                        })
                        .collect::<Result<_, _>>();
                }
            }

            Ok(vec![MPath::join_opt(
                prefix.as_ref(),
                mpath_element_iter(path),
            )])
        },
    ))
}

/// Get `HgManifestId`s for a set of `ChangesetId`s
/// This is needed for the purposes of implicit delete detection
async fn get_manifest_ids<'a, I: IntoIterator<Item = ChangesetId>>(
    ctx: &'a CoreContext,
    repo: &'a impl Repo,
    bcs_ids: I,
) -> Result<Vec<HgManifestId>, Error> {
    try_join_all(bcs_ids.into_iter().map({
        |bcs_id| {
            cloned!(ctx, repo);
            async move {
                let cs_id = repo.derive_hg_changeset(&ctx, bcs_id).await?;
                let hg_blob_changeset = cs_id.load(&ctx, repo.repo_blobstore()).await?;
                Ok(hg_blob_changeset.manifestid())
            }
        }
    }))
    .await
}

/// Take an iterator of file changes, which may contain implicit deletes
/// and produce a `SortedVectorMap` suitable to be used in the `BonsaiChangeset`,
/// without any implicit deletes.
fn minimize_file_change_set<I: IntoIterator<Item = (MPath, FileChange)>>(
    file_changes: I,
) -> SortedVectorMap<MPath, FileChange> {
    let (adds, removes): (Vec<_>, Vec<_>) = file_changes
        .into_iter()
        .partition(|(_, fc)| fc.is_changed());
    let adds: HashMap<MPath, FileChange> = adds.into_iter().collect();

    let prefix_path_was_added = |removed_path: MPath| {
        removed_path
            .into_parent_dir_iter()
            .any(|parent_dir| adds.contains_key(&parent_dir))
    };

    let filtered_removes = removes
        .into_iter()
        .filter(|(ref mpath, _)| !prefix_path_was_added(mpath.clone()));
    let mut result: SortedVectorMap<_, _> = filtered_removes.collect();
    result.extend(adds.into_iter());
    result
}

/// Given a changeset and it's parents, get the list of file
/// changes, which arise from "implicit deletes" as opposed
/// to naive `MPath` rewriting in `cs.file_changes`. For
/// more information about implicit deletes, please see
/// `manifest/src/implici_deletes.rs`
async fn get_implicit_delete_file_changes<'a, I: IntoIterator<Item = ChangesetId>>(
    ctx: &'a CoreContext,
    cs: BonsaiChangesetMut,
    parent_changeset_ids: I,
    mover: MultiMover,
    source_repo: &'a impl Repo,
) -> Result<Vec<(MPath, FileChange)>, Error> {
    let parent_manifest_ids = get_manifest_ids(ctx, source_repo, parent_changeset_ids).await?;
    let file_adds: Vec<_> = cs
        .file_changes
        .iter()
        .filter_map(|(mpath, file_change)| file_change.is_changed().then(|| mpath.clone()))
        .collect();
    let store = source_repo.repo_blobstore().clone();
    let implicit_deletes: Vec<MPath> =
        get_implicit_deletes(ctx, store, file_adds, parent_manifest_ids)
            .try_collect()
            .await?;
    let maybe_renamed_implicit_deletes: Result<Vec<Vec<MPath>>, _> =
        implicit_deletes.iter().map(|mpath| mover(mpath)).collect();
    let maybe_renamed_implicit_deletes: Vec<Vec<MPath>> = maybe_renamed_implicit_deletes?;
    let implicit_delete_file_changes: Vec<_> = maybe_renamed_implicit_deletes
        .into_iter()
        .flatten()
        .map(|implicit_delete_mpath| (implicit_delete_mpath, FileChange::Deletion))
        .collect();

    Ok(implicit_delete_file_changes)
}

/// Determines what to do in commits rewriting to empty commit in small repo.
///
/// NOTE: The empty commits from large repo are kept regardless of this flag.
#[derive(PartialEq, Debug, Copy, Clone, Default)]
pub enum CommitRewrittenToEmpty {
    Keep,
    #[default]
    Discard,
}

/// Determines what to do with commits that are empty in large repo.  They may
/// be useful to keep them in small repo if they have some special meaning.
///
/// NOTE: This flag doesn't affect non-empty commits from large repo rewriting
/// to empty commits in small repo. Use CommitsRewrittenToEmpty to control that.
#[derive(PartialEq, Debug, Copy, Clone, Default)]
pub enum EmptyCommitFromLargeRepo {
    #[default]
    Keep,
    Discard,
}

#[derive(PartialEq, Debug, Copy, Clone, Default)]
pub struct RewriteOpts {
    pub commit_rewritten_to_empty: CommitRewrittenToEmpty,
    pub empty_commit_from_large_repo: EmptyCommitFromLargeRepo,
}

/// Create a version of `cs` with `Mover` applied to all changes
/// The return value can be:
/// - `Err` if the rewrite failed
/// - `Ok(None)` if the rewrite decided that this commit should
///              not be present in the rewrite target
/// - `Ok(Some(rewritten))` for a successful rewrite, which should be
///                         present in the rewrite target
/// The notion that the commit "should not be present in the rewrite
/// target" means that the commit is not a merge and all of its changes
/// were rewritten into nothingness by the `Mover`.
///
/// Precondition: this function expects all `cs` parents to be present
/// in `remapped_parents` as keys, and their remapped versions as values.
///
/// If `force_first_parent` is set commit parents are reordered to ensure that
/// the specified changeset comes first.
pub async fn rewrite_commit<'a>(
    ctx: &'a CoreContext,
    cs: BonsaiChangesetMut,
    remapped_parents: &'a HashMap<ChangesetId, ChangesetId>,
    mover: MultiMover,
    source_repo: &'a impl Repo,
    force_first_parent: Option<ChangesetId>,
    rewrite_opts: RewriteOpts,
) -> Result<Option<BonsaiChangesetMut>, Error> {
    let delete_file_changes = if !cs.file_changes.is_empty() {
        get_implicit_delete_file_changes(
            ctx,
            cs.clone(),
            remapped_parents.keys().cloned(),
            mover.clone(),
            source_repo,
        )
        .await?
    } else {
        vec![]
    };

    internal_rewrite_commit_with_implicit_deletes(
        cs,
        remapped_parents,
        mover,
        force_first_parent,
        delete_file_changes,
        rewrite_opts,
    )
}

pub async fn rewrite_as_squashed_commit<'a>(
    ctx: &'a CoreContext,
    source_repo: &'a impl Repo,
    source_cs_id: ChangesetId,
    (source_parent_cs_id, target_parent_cs_id): (ChangesetId, ChangesetId),
    mut cs: BonsaiChangesetMut,
    mover: MultiMover,
    side_commits_info: Vec<String>,
) -> Result<Option<BonsaiChangesetMut>, Error> {
    let diff_stream = find_bonsai_diff(ctx, source_repo, source_parent_cs_id, source_cs_id).await?;
    let diff_changes: Vec<_> = diff_stream
        .map_ok(|diff_result| async move {
            convert_diff_result_into_file_change_for_diamond_merge(ctx, source_repo, diff_result)
                .await
        })
        .try_buffered(100)
        .try_collect()
        .await?;

    let rewritten_changes = diff_changes
        .into_iter()
        .map(|(path, change)| {
            let new_paths = mover(&path)?;
            Ok(new_paths
                .into_iter()
                .map(|new_path| (new_path, change.clone()))
                .collect())
        })
        .collect::<Result<Vec<Vec<_>>, Error>>()?;

    let rewritten_changes: SortedVectorMap<_, _> = rewritten_changes
        .into_iter()
        .flat_map(|changes| changes.into_iter())
        .collect();

    cs.file_changes = rewritten_changes;
    // `validate_can_sync_changeset` already ensures
    // that target_parent_cs_id is one of the existing parents
    cs.parents = vec![target_parent_cs_id];
    let old_message = cs.message;
    cs.message = format!(
        "{}{}{}",
        old_message,
        SQUASH_DELIMITER_MESSAGE,
        side_commits_info.join("\n")
    );
    Ok(Some(cs))
}

pub async fn rewrite_stack_no_merges<'a>(
    ctx: &'a CoreContext,
    css: Vec<BonsaiChangeset>,
    mut rewritten_parent: ChangesetId,
    mover: MultiMover,
    source_repo: &'a impl Repo,
    force_first_parent: Option<ChangesetId>,
    mut modify_bonsai_cs: impl FnMut((ChangesetId, BonsaiChangesetMut)) -> BonsaiChangesetMut,
) -> Result<Vec<Option<BonsaiChangeset>>, Error> {
    borrowed!(mover: &Arc<_>, source_repo);

    for cs in &css {
        if cs.is_merge() {
            return Err(anyhow!(
                "cannot remap merges in a stack - {}",
                cs.get_changeset_id()
            ));
        }
    }

    let css = stream::iter(css)
        .map({
            |cs| async move {
                let deleted_file_changes = if cs.file_changes().next().is_some() {
                    let parents = cs.parents();
                    get_implicit_delete_file_changes(
                        ctx,
                        cs.clone().into_mut(),
                        parents,
                        mover.clone(),
                        source_repo,
                    )
                    .await?
                } else {
                    vec![]
                };

                anyhow::Ok((cs, deleted_file_changes))
            }
        })
        .buffered(100)
        .try_collect::<Vec<_>>()
        .await?;

    let mut res = vec![];
    for (from_cs, implicit_deletes_file_changes) in css {
        let from_cs_id = from_cs.get_changeset_id();
        let from_cs = from_cs.into_mut();

        let mut remapped_parents = HashMap::new();
        if let Some(parent) = from_cs.parents.get(0) {
            remapped_parents.insert(*parent, rewritten_parent);
        }

        let maybe_cs = internal_rewrite_commit_with_implicit_deletes(
            from_cs,
            &remapped_parents,
            mover.clone(),
            force_first_parent,
            implicit_deletes_file_changes,
            Default::default(),
        )?;

        let maybe_cs = maybe_cs
            .map(|cs| modify_bonsai_cs((from_cs_id, cs)))
            .map(|bcs| bcs.freeze())
            .transpose()?;
        if let Some(ref cs) = maybe_cs {
            let to_cs_id = cs.get_changeset_id();
            rewritten_parent = to_cs_id;
        }

        res.push(maybe_cs);
    }

    Ok(res)
}

pub fn internal_rewrite_commit_with_implicit_deletes<'a>(
    mut cs: BonsaiChangesetMut,
    remapped_parents: &'a HashMap<ChangesetId, ChangesetId>,
    mover: MultiMover,
    force_first_parent: Option<ChangesetId>,
    implicit_delete_file_changes: Vec<(MPath, FileChange)>,
    rewrite_opts: RewriteOpts,
) -> Result<Option<BonsaiChangesetMut>, Error> {
    let empty_commit = cs.file_changes.is_empty();
    if !empty_commit
        || rewrite_opts.empty_commit_from_large_repo == EmptyCommitFromLargeRepo::Discard
    {
        let path_rewritten_changes: Result<Vec<Vec<_>>, _> = cs
            .file_changes
            .into_iter()
            .map(|(path, change)| {
                // Just rewrite copy_from information, when we have it
                fn rewrite_copy_from(
                    copy_from: &(MPath, ChangesetId),
                    remapped_parents: &HashMap<ChangesetId, ChangesetId>,
                    mover: MultiMover,
                ) -> Result<Option<(MPath, ChangesetId)>, Error> {
                    let (path, copy_from_commit) = copy_from;
                    let new_paths = mover(path)?;
                    let copy_from_commit =
                        remapped_parents.get(copy_from_commit).ok_or_else(|| {
                            Error::from(ErrorKind::MissingRemappedCommit(*copy_from_commit))
                        })?;

                    // If the source path doesn't remap, drop this copy info.

                    // TODO(stash): a path can be remapped to multiple other paths,
                    // but for copy_from path we pick only the first one. Instead of
                    // picking only the first one, it's a better to have a dedicated
                    // field in a thrift struct which says which path should be picked
                    // as copy from
                    Ok(new_paths
                        .get(0)
                        .cloned()
                        .map(|new_path| (new_path, *copy_from_commit)))
                }

                // Extract any copy_from information, and use rewrite_copy_from on it
                fn rewrite_file_change(
                    change: TrackedFileChange,
                    remapped_parents: &HashMap<ChangesetId, ChangesetId>,
                    mover: MultiMover,
                ) -> Result<FileChange, Error> {
                    let new_copy_from = change
                        .copy_from()
                        .and_then(|copy_from| {
                            rewrite_copy_from(copy_from, remapped_parents, mover).transpose()
                        })
                        .transpose()?;

                    Ok(FileChange::Change(change.with_new_copy_from(new_copy_from)))
                }

                // Rewrite both path and changes
                fn do_rewrite(
                    path: MPath,
                    change: FileChange,
                    remapped_parents: &HashMap<ChangesetId, ChangesetId>,
                    mover: MultiMover,
                ) -> Result<Vec<(MPath, FileChange)>, Error> {
                    let new_paths = mover(&path)?;
                    let change = match change {
                        FileChange::Change(tc) => {
                            rewrite_file_change(tc, remapped_parents, mover.clone())?
                        }
                        FileChange::Deletion => FileChange::Deletion,
                        FileChange::UntrackedDeletion | FileChange::UntrackedChange(_) => {
                            bail!("Can't rewrite untracked changes")
                        }
                    };
                    Ok(new_paths
                        .into_iter()
                        .map(|new_path| (new_path, change.clone()))
                        .collect())
                }
                do_rewrite(path, change, remapped_parents, mover.clone())
            })
            .collect();

        let mut path_rewritten_changes: SortedVectorMap<_, _> = path_rewritten_changes?
            .into_iter()
            .flat_map(|changes| changes.into_iter())
            .collect();

        path_rewritten_changes.extend(implicit_delete_file_changes.into_iter());
        let path_rewritten_changes = minimize_file_change_set(path_rewritten_changes.into_iter());
        let is_merge = cs.parents.len() >= 2;

        // If all parent has < 2 commits then it's not a merge, and it was completely rewritten
        // out. In that case we can just discard it because there are not changes to the working copy.
        // However if it's a merge then we can't discard it, because even
        // though bonsai merge commit might not have file changes inside it can still change
        // a working copy. E.g. if p1 has fileA, p2 has fileB, then empty merge(p1, p2)
        // contains both fileA and fileB.
        if !is_merge
            && ((path_rewritten_changes.is_empty()
                && rewrite_opts.commit_rewritten_to_empty == CommitRewrittenToEmpty::Discard)
                || (empty_commit
                    && rewrite_opts.empty_commit_from_large_repo
                        == EmptyCommitFromLargeRepo::Discard))
        {
            return Ok(None);
        } else {
            cs.file_changes = path_rewritten_changes;
        }
    }

    // Update hashes
    for commit in cs.parents.iter_mut() {
        let remapped = remapped_parents
            .get(commit)
            .ok_or_else(|| Error::from(ErrorKind::MissingRemappedCommit(*commit)))?;

        *commit = *remapped;
    }
    if let Some(first_parent) = force_first_parent {
        if !cs.parents.contains(&first_parent) {
            return Err(Error::from(ErrorKind::MissingForcedParent(first_parent)));
        }
        let mut new_parents = vec![first_parent];
        new_parents.extend(cs.parents.into_iter().filter(|cs| *cs != first_parent));
        cs.parents = new_parents
    }

    Ok(Some(cs))
}

pub async fn upload_commits<'a>(
    ctx: &'a CoreContext,
    rewritten_list: Vec<BonsaiChangeset>,
    source_repo: &'a (impl RepoBlobstoreRef + ChangesetsRef),
    target_repo: &'a (impl RepoBlobstoreRef + ChangesetsRef + FilestoreConfigRef),
) -> Result<(), Error> {
    let mut files_to_sync = vec![];
    for rewritten in &rewritten_list {
        let rewritten_mut = rewritten.clone().into_mut();
        let new_files_to_sync =
            rewritten_mut
                .file_changes
                .values()
                .filter_map(|change| match change {
                    FileChange::Change(tc) => Some(tc.content_id()),
                    FileChange::UntrackedChange(uc) => Some(uc.content_id()),
                    FileChange::Deletion | FileChange::UntrackedDeletion => None,
                });
        files_to_sync.extend(new_files_to_sync);
    }
    copy_file_contents(ctx, source_repo, target_repo, files_to_sync, |_| {}).await?;
    save_bonsai_changesets(rewritten_list.clone(), ctx.clone(), target_repo).await?;
    Ok(())
}

pub async fn copy_file_contents<'a>(
    ctx: &'a CoreContext,
    source_repo: &'a impl RepoBlobstoreRef,
    target_repo: &'a (impl RepoBlobstoreRef + FilestoreConfigRef),
    content_ids: impl IntoIterator<Item = ContentId>,
    progress_reporter: impl Fn(usize),
) -> Result<(), Error> {
    let source_blobstore = source_repo.repo_blobstore();
    let target_blobstore = target_repo.repo_blobstore();
    let target_filestore_config = target_repo.filestore_config();

    let mut i = 0;
    stream::iter(content_ids.into_iter().map({
        |content_id| {
            copy_content(
                ctx,
                source_blobstore,
                target_blobstore,
                target_filestore_config.clone(),
                content_id,
            )
        }
    }))
    .buffer_unordered(100)
    .try_for_each(|_| {
        i += 1;
        progress_reporter(i);
        async { Ok(()) }
    })
    .await
}

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

    use anyhow::bail;
    use blobrepo::save_bonsai_changesets;
    use fbinit::FacebookInit;
    use maplit::btreemap;
    use maplit::hashmap;
    use mononoke_types::ContentId;
    use mononoke_types::FileType;
    use test_repo_factory::TestRepoFactory;
    use tests_utils::list_working_copy_utf8;
    use tests_utils::CreateCommitContext;

    use super::*;

    #[test]
    fn test_multi_mover_simple() -> Result<(), Error> {
        let mapping_rules = SourceMappingRules {
            default_prefix: "".to_string(),
            ..Default::default()
        };
        let multi_mover = create_source_to_target_multi_mover(mapping_rules)?;
        assert_eq!(
            multi_mover(&MPath::new("path")?)?,
            vec![MPath::new("path")?]
        );
        Ok(())
    }

    #[test]
    fn test_multi_mover_prefixed() -> Result<(), Error> {
        let mapping_rules = SourceMappingRules {
            default_prefix: "prefix".to_string(),
            ..Default::default()
        };
        let multi_mover = create_source_to_target_multi_mover(mapping_rules)?;
        assert_eq!(
            multi_mover(&MPath::new("path")?)?,
            vec![MPath::new("prefix/path")?]
        );
        Ok(())
    }

    #[test]
    fn test_multi_mover_prefixed_with_exceptions() -> Result<(), Error> {
        let mapping_rules = SourceMappingRules {
            default_prefix: "prefix".to_string(),
            overrides: btreemap! {
                "override".to_string() => vec![
                    "overriden_1".to_string(),
                    "overriden_2".to_string(),
                ]
            },
            ..Default::default()
        };
        let multi_mover = create_source_to_target_multi_mover(mapping_rules)?;
        assert_eq!(
            multi_mover(&MPath::new("path")?)?,
            vec![MPath::new("prefix/path")?]
        );

        assert_eq!(
            multi_mover(&MPath::new("override/path")?)?,
            vec![
                MPath::new("overriden_1/path")?,
                MPath::new("overriden_2/path")?,
            ]
        );
        Ok(())
    }

    #[test]
    fn test_multi_mover_longest_prefix_first() -> Result<(), Error> {
        let mapping_rules = SourceMappingRules {
            default_prefix: "prefix".to_string(),
            overrides: btreemap! {
                "prefix".to_string() => vec![
                    "prefix_1".to_string(),
                ],
                "prefix/sub".to_string() => vec![
                    "prefix/sub_1".to_string(),
                ]
            },
            ..Default::default()
        };
        let multi_mover = create_source_to_target_multi_mover(mapping_rules)?;
        assert_eq!(
            multi_mover(&MPath::new("prefix/path")?)?,
            vec![MPath::new("prefix_1/path")?]
        );

        assert_eq!(
            multi_mover(&MPath::new("prefix/sub/path")?)?,
            vec![MPath::new("prefix/sub_1/path")?]
        );

        Ok(())
    }

    fn path(p: &str) -> MPath {
        MPath::new(p).unwrap()
    }

    fn verify_minimized(changes: Vec<(&str, Option<()>)>, expected: BTreeMap<&str, Option<()>>) {
        fn to_file_change(o: Option<()>) -> FileChange {
            match o {
                Some(_) => FileChange::tracked(
                    ContentId::from_bytes([1; 32]).unwrap(),
                    FileType::Regular,
                    0,
                    None,
                ),
                None => FileChange::Deletion,
            }
        }
        let changes: Vec<_> = changes
            .into_iter()
            .map(|(p, c)| (path(p), to_file_change(c)))
            .collect();
        let minimized = minimize_file_change_set(changes);
        let expected: SortedVectorMap<MPath, FileChange> = expected
            .into_iter()
            .map(|(p, c)| (path(p), to_file_change(c)))
            .collect();
        assert_eq!(expected, minimized);
    }

    #[fbinit::test]
    fn test_minimize_file_change_set(_fb: FacebookInit) {
        verify_minimized(
            vec![("a", Some(())), ("a", None)],
            btreemap! { "a" => Some(())},
        );
        verify_minimized(vec![("a", Some(()))], btreemap! { "a" => Some(())});
        verify_minimized(vec![("a", None)], btreemap! { "a" => None});
        // directories are deleted implicitly, so explicit deletes are
        // minimized away
        verify_minimized(
            vec![("a/b", None), ("a/c", None), ("a", Some(()))],
            btreemap! { "a" => Some(()) },
        );
        // files, replaced with a directy at a longer path are not
        // deleted implicitly, so they aren't minimized away
        verify_minimized(
            vec![("a", None), ("a/b", Some(()))],
            btreemap! { "a" => None, "a/b" => Some(()) },
        );
    }

    #[fbinit::test]
    async fn test_rewrite_commit(fb: FacebookInit) -> Result<(), Error> {
        let repo: blobrepo::BlobRepo = TestRepoFactory::new(fb)?.build().await?;
        let ctx = CoreContext::test_mock(fb);
        let first = CreateCommitContext::new_root(&ctx, &repo)
            .add_file("path", "path")
            .commit()
            .await?;
        let second = CreateCommitContext::new(&ctx, &repo, vec![first])
            .add_file_with_copy_info("pathsecondcommit", "pathsecondcommit", (first, "path"))
            .commit()
            .await?;
        let third = CreateCommitContext::new(&ctx, &repo, vec![first, second])
            .add_file("path", "pathmodified")
            .commit()
            .await?;

        let mapping_rules = SourceMappingRules {
            default_prefix: "prefix".to_string(),
            overrides: btreemap! {
                "path".to_string() => vec![
                    "path_1".to_string(),
                    "path_2".to_string(),
                ]
            },
            ..Default::default()
        };
        let multi_mover = create_source_to_target_multi_mover(mapping_rules)?;

        let first_rewritten_bcs_id = test_rewrite_commit_cs_id(
            &ctx,
            &repo,
            first,
            HashMap::new(),
            multi_mover.clone(),
            None,
        )
        .await?;

        let first_rewritten_wc =
            list_working_copy_utf8(&ctx, &repo, first_rewritten_bcs_id).await?;
        assert_eq!(
            first_rewritten_wc,
            hashmap! {
                MPath::new("path_1")? => "path".to_string(),
                MPath::new("path_2")? => "path".to_string(),
            }
        );

        let second_rewritten_bcs_id = test_rewrite_commit_cs_id(
            &ctx,
            &repo,
            second,
            hashmap! {
                first => first_rewritten_bcs_id
            },
            multi_mover.clone(),
            None,
        )
        .await?;

        let second_bcs = second_rewritten_bcs_id
            .load(&ctx, &repo.repo_blobstore())
            .await?;
        let maybe_copy_from = match second_bcs
            .file_changes_map()
            .get(&MPath::new("prefix/pathsecondcommit")?)
            .ok_or_else(|| anyhow!("path not found"))?
        {
            FileChange::Change(tc) => tc.copy_from().cloned(),
            _ => bail!("path_is_deleted"),
        };

        assert_eq!(
            maybe_copy_from,
            Some((MPath::new("path_1")?, first_rewritten_bcs_id))
        );

        let second_rewritten_wc =
            list_working_copy_utf8(&ctx, &repo, second_rewritten_bcs_id).await?;
        assert_eq!(
            second_rewritten_wc,
            hashmap! {
                MPath::new("path_1")? => "path".to_string(),
                MPath::new("path_2")? => "path".to_string(),
                MPath::new("prefix/pathsecondcommit")? => "pathsecondcommit".to_string(),
            }
        );

        // Diamond merge test with error during parent reordering
        assert!(
            test_rewrite_commit_cs_id(
                &ctx,
                &repo,
                third,
                hashmap! {
                    first => first_rewritten_bcs_id,
                    second => second_rewritten_bcs_id
                },
                multi_mover.clone(),
                Some(second), // wrong, should be after-rewrite id
            )
            .await
            .is_err()
        );

        // Diamond merge test with success
        let third_rewritten_bcs_id = test_rewrite_commit_cs_id(
            &ctx,
            &repo,
            third,
            hashmap! {
                first => first_rewritten_bcs_id,
                second => second_rewritten_bcs_id
            },
            multi_mover,
            Some(second_rewritten_bcs_id),
        )
        .await?;

        let third_bcs = third_rewritten_bcs_id
            .load(&ctx, &repo.repo_blobstore().clone())
            .await?;

        assert_eq!(
            third_bcs.parents().collect::<Vec<_>>(),
            vec![second_rewritten_bcs_id, first_rewritten_bcs_id],
        );

        Ok(())
    }

    async fn test_rewrite_commit_cs_id<'a>(
        ctx: &'a CoreContext,
        repo: &'a impl Repo,
        bcs_id: ChangesetId,
        parents: HashMap<ChangesetId, ChangesetId>,
        multi_mover: MultiMover,
        force_first_parent: Option<ChangesetId>,
    ) -> Result<ChangesetId, Error> {
        let bcs = bcs_id.load(ctx, &repo.repo_blobstore()).await?;
        let bcs = bcs.into_mut();

        let maybe_rewritten = rewrite_commit(
            ctx,
            bcs,
            &parents,
            multi_mover,
            repo,
            force_first_parent,
            Default::default(),
        )
        .await?;
        let rewritten =
            maybe_rewritten.ok_or_else(|| anyhow!("can't rewrite commit {}", bcs_id))?;
        let rewritten = rewritten.freeze()?;

        save_bonsai_changesets(vec![rewritten.clone()], ctx.clone(), repo).await?;

        Ok(rewritten.get_changeset_id())
    }

    #[test]
    fn test_directory_multi_mover() -> Result<(), Error> {
        let mapping_rules = SourceMappingRules {
            default_prefix: "prefix".to_string(),
            ..Default::default()
        };
        let multi_mover = create_directory_source_to_target_multi_mover(mapping_rules)?;
        assert_eq!(
            multi_mover(&Some(MPath::new("path")?))?,
            vec![Some(MPath::new("prefix/path")?)]
        );

        assert_eq!(multi_mover(&None)?, vec![Some(MPath::new("prefix")?)]);
        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.
 */

#![cfg_attr(not(fbcode_build), allow(unused_crate_dependencies))]

#[cfg(fbcode_build)]
mod facebook;
#[cfg(fbcode_build)]
pub use facebook::ConfigObject;
#[cfg(fbcode_build)]
pub use facebook::MononokeConfigoClient;

// There is no way to implement this for non-fbcode builds
// and it's worth having this crate's users know this, so
// let's not even export any stubs.


// @generated by autocargo
use std::env;
use std::fs;
use std::path::Path;

use thrift_compiler::Config;
use thrift_compiler::GenContext;

#[rustfmt::skip]
fn main() {
    // Rerun if this gets rewritten.
    println!("cargo:rerun-if-changed=thrift_build.rs");

    let out_dir = env::var_os("OUT_DIR").expect("OUT_DIR env not provided");
    let out_dir: &Path = out_dir.as_ref();
    fs::write(
        out_dir.join("cratemap"),
        "megarepo_types_thrift crate
mononoke_types_thrift mononoke_types_thrift
source_control source_control",
    ).expect("Failed to write cratemap");

    let conf = {
        let mut conf = Config::from_env(GenContext::Lib).expect("Failed to instantiate thrift_compiler::Config");

        let path_from_manifest_to_base: &Path = "../../../..".as_ref();
        let cargo_manifest_dir =
            env::var_os("CARGO_MANIFEST_DIR").expect("CARGO_MANIFEST_DIR not provided");
        let cargo_manifest_dir: &Path = cargo_manifest_dir.as_ref();
        let base_path = cargo_manifest_dir
            .join(path_from_manifest_to_base)
            .canonicalize()
            .expect("Failed to canonicalize base_path");
        // TODO: replace canonicalize() with std::path::absolute() when
        // https://github.com/rust-lang/rust/pull/91673 is available (~Rust 1.60)
        // and remove this block.
        #[cfg(windows)]
        let base_path = Path::new(
            base_path
                .as_path()
                .to_string_lossy()
                .trim_start_matches(r"\\?\"),
            )
            .to_path_buf();

        conf.base_path(base_path);

        let options = "types_crate=megarepo_types-thrift__types";
        if !options.is_empty() {
            conf.options(options);
        }

        let include_srcs = vec![

        ];
        conf.include_srcs(include_srcs);

        conf
    };

    conf
        .run(&[
            "megarepo_types_thrift.thrift"
        ])
        .expect("Failed while running thrift compilation");
}


// @generated by autocargo
::codegen_includer_proc_macro::include!();


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

#![feature(trait_alias)]

use std::collections::BTreeMap;
use std::fmt;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use context::CoreContext;
use context::PerfCounterType;
use derived_data::BonsaiDerived;
use filestore::FilestoreConfigRef;
use fsnodes::RootFsnodeId;
use manifest::Entry;
use manifest::ManifestOps;
pub use megarepo_configs::types::Source;
pub use megarepo_configs::types::SourceMappingRules;
pub use megarepo_configs::types::SourceRevision;
pub use megarepo_configs::types::SyncConfigVersion;
pub use megarepo_configs::types::SyncTargetConfig;
pub use megarepo_configs::types::Target;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use mononoke_types::ContentId;
use mononoke_types::FileChange;
use mononoke_types::FileType;
use mononoke_types::MPath;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedDataRef;
use serde::Deserialize;
use serde::Serialize;
use sql::Connection;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;

pub trait Repo = RepoDerivedDataRef + RepoBlobstoreRef + FilestoreConfigRef + Sync + Send;

mononoke_queries! {
    read GetMappingEntry(
        target_repo_id: i64,
        target_bookmark: String,
        target_bcs_id: ChangesetId,
    ) -> (String, ChangesetId, SyncConfigVersion) {
        "SELECT source_name, source_bcs_id, sync_config_version
          FROM megarepo_changeset_mapping
          WHERE target_repo_id = {target_repo_id}
          AND target_bookmark = {target_bookmark}
          AND target_bcs_id = {target_bcs_id}
          "
    }

    write InsertMapping(values: (
        source_name: str,
        target_repo_id: i64,
        target_bookmark: String,
        source_bcs_id: ChangesetId,
        target_bcs_id: ChangesetId,
        sync_config_version: SyncConfigVersion,
    )) {
        insert_or_ignore,
        "{insert_or_ignore} INTO megarepo_changeset_mapping
        (source_name, target_repo_id, target_bookmark, source_bcs_id, target_bcs_id, sync_config_version)
        VALUES {values}"
    }

    read GetReverseMappingEntry(
        target_repo_id: i64,
        target_bookmark: String,
        source_bcs_id: ChangesetId,
    ) -> (String, ChangesetId, SyncConfigVersion) {
        "SELECT source_name, target_bcs_id, sync_config_version
        FROM megarepo_changeset_mapping
        WHERE target_repo_id = {target_repo_id}
        AND target_bookmark = {target_bookmark}
        AND source_bcs_id = {source_bcs_id}"
    }
}

pub struct MegarepoMapping {
    pub(crate) connections: SqlConnections,
}

pub const REMAPPING_STATE_FILE: &str = ".megarepo/remapping_state";

#[derive(
    Clone,
    Debug,
    Hash,
    Eq,
    Ord,
    PartialOrd,
    PartialEq,
    Serialize,
    Deserialize
)]
#[serde(transparent)]
pub struct SourceName(pub String);
impl SourceName {
    pub fn new<T: ToString>(name: T) -> Self {
        SourceName(name.to_string())
    }

    pub fn as_str(&self) -> &str {
        &self.0
    }
}

impl fmt::Display for SourceName {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        self.0.fmt(f)
    }
}

#[derive(Clone, Serialize, Deserialize)]
pub struct CommitRemappingState {
    /// Mapping from source to a changeset id
    pub latest_synced_changesets: BTreeMap<SourceName, ChangesetId>,
    /// Config version that was used to create this commit
    sync_config_version: SyncConfigVersion,
}

impl CommitRemappingState {
    pub fn new(
        latest_synced_changesets: BTreeMap<SourceName, ChangesetId>,
        sync_config_version: SyncConfigVersion,
    ) -> Self {
        Self {
            latest_synced_changesets,
            sync_config_version,
        }
    }

    pub async fn read_state_from_commit(
        ctx: &CoreContext,
        repo: &impl Repo,
        cs_id: ChangesetId,
    ) -> Result<Self, Error> {
        let maybe_state = Self::read_state_from_commit_opt(ctx, repo, cs_id).await?;

        maybe_state.ok_or_else(|| anyhow!("file {} not found", REMAPPING_STATE_FILE))
    }

    pub async fn read_state_from_commit_opt(
        ctx: &CoreContext,
        repo: &impl Repo,
        cs_id: ChangesetId,
    ) -> Result<Option<Self>, Error> {
        let root_fsnode_id = RootFsnodeId::derive(ctx, repo, cs_id).await?;

        let path = MPath::new(REMAPPING_STATE_FILE)?;
        let maybe_entry = root_fsnode_id
            .fsnode_id()
            .find_entry(ctx.clone(), repo.repo_blobstore().clone(), Some(path))
            .await?;

        let entry = match maybe_entry {
            Some(entry) => entry,
            None => {
                return Ok(None);
            }
        };

        let file = match entry {
            Entry::Tree(_) => {
                return Ok(None);
            }
            Entry::Leaf(file) => file,
        };

        let bytes = filestore::fetch_concat(repo.repo_blobstore(), ctx, *file.content_id()).await?;
        let content = String::from_utf8(bytes.to_vec())
            .with_context(|| format!("{} is not utf8", REMAPPING_STATE_FILE))?;
        let state: CommitRemappingState = serde_json::from_str(&content)?;
        Ok(Some(state))
    }

    pub async fn save_in_changeset(
        &self,
        ctx: &CoreContext,
        repo: &impl Repo,
        bcs: &mut BonsaiChangesetMut,
    ) -> Result<(), Error> {
        let (content_id, size) = self.save(ctx, repo).await?;
        let path = MPath::new(REMAPPING_STATE_FILE)?;

        let fc = FileChange::tracked(content_id, FileType::Regular, size, None);
        if bcs.file_changes.insert(path, fc).is_some() {
            return Err(anyhow!(
                "New bonsai changeset already has {} file",
                REMAPPING_STATE_FILE,
            ));
        }

        Ok(())
    }

    pub fn set_source_changeset(&mut self, source: SourceName, cs_id: ChangesetId) {
        self.latest_synced_changesets.insert(source, cs_id);
    }

    pub fn get_latest_synced_changeset(&self, source: &SourceName) -> Option<&ChangesetId> {
        self.latest_synced_changesets.get(source)
    }

    pub fn get_all_latest_synced_changesets(&self) -> &BTreeMap<SourceName, ChangesetId> {
        &self.latest_synced_changesets
    }

    pub fn sync_config_version(&self) -> &SyncConfigVersion {
        &self.sync_config_version
    }

    async fn save(&self, ctx: &CoreContext, repo: &impl Repo) -> Result<(ContentId, u64), Error> {
        let bytes = self.serialize()?;

        let ((content_id, size), fut) = filestore::store_bytes(
            repo.repo_blobstore(),
            *repo.filestore_config(),
            ctx,
            bytes.into(),
        );

        fut.await?;

        Ok((content_id, size))
    }

    fn serialize(&self) -> Result<Vec<u8>, Error> {
        serde_json::to_vec_pretty(&self).map_err(Error::from)
    }
}

pub struct MegarepoMappingEntry {
    pub source_name: SourceName,
    pub target: Target,
    pub source_cs_id: ChangesetId,
    pub target_cs_id: ChangesetId,
    pub sync_config_version: SyncConfigVersion,
}

impl MegarepoMapping {
    /// For a given (target, cs_id) pair return the version that was used
    /// to create target changeset id.
    /// Usually this method is used to find what version do we need to use
    /// for rewriting a commit
    pub async fn get_mapping_entry(
        &self,
        ctx: &CoreContext,
        target: &Target,
        target_cs_id: ChangesetId,
    ) -> Result<Option<MegarepoMappingEntry>, Error> {
        let maybe_version = self
            .get_mapping_entry_impl(
                ctx,
                target,
                target_cs_id,
                PerfCounterType::SqlReadsReplica,
                &self.connections.read_connection,
            )
            .await?;

        if let Some(version) = maybe_version {
            return Ok(Some(version));
        }

        self.get_mapping_entry_impl(
            ctx,
            target,
            target_cs_id,
            PerfCounterType::SqlReadsMaster,
            &self.connections.read_master_connection,
        )
        .await
    }

    async fn get_mapping_entry_impl(
        &self,
        ctx: &CoreContext,
        target: &Target,
        target_cs_id: ChangesetId,
        sql_perf_counter: PerfCounterType,
        connection: &Connection,
    ) -> Result<Option<MegarepoMappingEntry>, Error> {
        ctx.perf_counters().increment_counter(sql_perf_counter);
        let mut rows =
            GetMappingEntry::query(connection, &target.repo_id, &target.bookmark, &target_cs_id)
                .await?;

        if rows.len() > 1 {
            return Err(anyhow!(
                "Programming error - more than 1 row returned for get target config version"
            ));
        }

        Ok(rows.pop().map(|x| MegarepoMappingEntry {
            source_name: SourceName::new(x.0),
            source_cs_id: x.1,
            sync_config_version: x.2,
            target: target.clone(),
            target_cs_id: target_cs_id.clone(),
        }))
    }

    // Reverse lookup of the previous query
    // It is possible to have same source_cs_id mapped to the different targets
    // but this method may return stale(not complete) set of mappings.
    pub async fn get_reverse_mapping_entry(
        &self,
        ctx: &CoreContext,
        target: &Target,
        source_cs_id: ChangesetId,
    ) -> Result<Vec<MegarepoMappingEntry>, Error> {
        let entries = self
            .get_reverse_mapping_entry_impl(
                ctx,
                target,
                source_cs_id,
                PerfCounterType::SqlReadsReplica,
                &self.connections.read_connection,
            )
            .await?;

        if !entries.is_empty() {
            return Ok(entries);
        }

        self.get_reverse_mapping_entry_impl(
            ctx,
            target,
            source_cs_id,
            PerfCounterType::SqlReadsMaster,
            &self.connections.read_master_connection,
        )
        .await
    }

    async fn get_reverse_mapping_entry_impl(
        &self,
        ctx: &CoreContext,
        target: &Target,
        source_cs_id: ChangesetId,
        sql_perf_counter: PerfCounterType,
        connection: &Connection,
    ) -> Result<Vec<MegarepoMappingEntry>, Error> {
        ctx.perf_counters().increment_counter(sql_perf_counter);
        let rows = GetReverseMappingEntry::query(
            connection,
            &target.repo_id,
            &target.bookmark,
            &source_cs_id,
        )
        .await?;

        Ok(rows
            .into_iter()
            .map(|x| MegarepoMappingEntry {
                source_name: SourceName::new(x.0),
                source_cs_id: source_cs_id.clone(),
                sync_config_version: x.2,
                target: target.clone(),
                target_cs_id: x.1,
            })
            .collect())
    }

    /// Add a mapping from a source commit to a target commit
    #[allow(clippy::ptr_arg)]
    pub async fn insert_source_target_cs_mapping(
        &self,
        ctx: &CoreContext,
        source_name: &SourceName,
        target: &Target,
        source_cs_id: ChangesetId,
        target_cs_id: ChangesetId,
        version: &SyncConfigVersion,
    ) -> Result<(), Error> {
        ctx.perf_counters()
            .increment_counter(PerfCounterType::SqlWrites);

        let res = InsertMapping::query(
            &self.connections.write_connection,
            &[(
                source_name.as_str(),
                &target.repo_id,
                &target.bookmark,
                &source_cs_id,
                &target_cs_id,
                version,
            )],
        )
        .await?;
        if res.affected_rows() == 0 {
            // Becase we insert to mapping before moving bookmark (which is fallible)
            // the mapping might be already inserted at that point. If it's the same
            // as what we wanted to insert we can ignore the failure to insert.
            if let Ok(Some(entry)) = self.get_mapping_entry(ctx, target, target_cs_id).await {
                if &entry.source_name != source_name
                    || entry.source_cs_id != source_cs_id
                    || &entry.sync_config_version != version
                {
                    return Err(anyhow!(
                        "trying to insert mapping whille one already exists and it's different!"
                    ));
                }
            } else {
                return Err(anyhow!(
                    "unknown error while inserting mapping (affected_rows=0)"
                ));
            }
        }

        Ok(())
    }
}

impl SqlConstruct for MegarepoMapping {
    const LABEL: &'static str = "megarepo_mapping";

    const CREATION_QUERY: &'static str = include_str!("../schemas/sqlite-megarepo-mapping.sql");

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for MegarepoMapping {}

#[cfg(test)]
mod test {
    use fbinit::FacebookInit;
    use maplit::btreemap;
    use mononoke_types_mocks::changesetid::ONES_CSID;
    use mononoke_types_mocks::changesetid::THREES_CSID;
    use mononoke_types_mocks::changesetid::TWOS_CSID;

    use super::*;

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

        let target = Target {
            repo_id: 0,
            bookmark: "book".to_string(),
            ..Default::default()
        };

        let source_csid = ONES_CSID;
        let target_csid = TWOS_CSID;
        let version = "version".to_string();

        mapping
            .insert_source_target_cs_mapping(
                &ctx,
                &SourceName::new("source_name"),
                &target,
                source_csid,
                target_csid,
                &version,
            )
            .await?;

        // Test to check if insertion is resilient against
        // the mapping being already there.
        mapping
            .insert_source_target_cs_mapping(
                &ctx,
                &SourceName::new("source_name"),
                &target,
                source_csid,
                target_csid,
                &version,
            )
            .await?;

        let res = mapping
            .get_mapping_entry(&ctx, &target, target_csid)
            .await?;

        assert_eq!(res.unwrap().sync_config_version, version);

        Ok(())
    }

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

        let target = Target {
            repo_id: 0,
            bookmark: "book".to_string(),
            ..Default::default()
        };

        let source_csid = ONES_CSID;
        let target_csid = TWOS_CSID;
        let version = "version".to_string();

        mapping
            .insert_source_target_cs_mapping(
                &ctx,
                &SourceName::new("source_name"),
                &target,
                source_csid,
                target_csid,
                &version,
            )
            .await?;

        let mut res = mapping
            .get_reverse_mapping_entry(&ctx, &target, source_csid)
            .await?;

        assert_eq!(res.len(), 1);

        assert_eq!(res.pop().unwrap().target_cs_id, target_csid);

        // query non-existent source_cs
        let res = mapping
            .get_reverse_mapping_entry(&ctx, &target, THREES_CSID)
            .await?;

        assert_eq!(res.len(), 0);

        Ok(())
    }

    #[fbinit::test]
    async fn test_serialize(_fb: FacebookInit) -> Result<(), Error> {
        let state = CommitRemappingState::new(
            btreemap! {
                SourceName::new("source_1") => ONES_CSID,
                SourceName::new("source_2") => TWOS_CSID,
            },
            "version_1".to_string(),
        );
        let res = String::from_utf8(state.serialize()?)?;
        assert_eq!(
            res,
            r#"{
  "latest_synced_changesets": {
    "source_1": "1111111111111111111111111111111111111111111111111111111111111111",
    "source_2": "2222222222222222222222222222222222222222222222222222222222222222"
  },
  "sync_config_version": "version_1"
}"#
        );

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

#![cfg_attr(not(fbcode_build), allow(unused_crate_dependencies))]

use std::path::Path;
use std::path::PathBuf;

use async_trait::async_trait;
use clap::ArgAction;
use clap::Args;
use context::CoreContext;
pub use megarepo_configs::types::MergeMode;
pub use megarepo_configs::types::Source;
pub use megarepo_configs::types::SourceMappingRules;
pub use megarepo_configs::types::SourceRevision;
pub use megarepo_configs::types::Squashed;
pub use megarepo_configs::types::SyncConfigVersion;
pub use megarepo_configs::types::SyncTargetConfig;
pub use megarepo_configs::types::Target;
pub use megarepo_configs::types::WithExtraMoveCommit;
use megarepo_error::MegarepoError;
#[cfg(fbcode_build)]
mod facebook;
#[cfg(not(fbcode_build))]
mod oss;
mod test_impl;
mod verification;

#[cfg(fbcode_build)]
pub use facebook::CfgrMononokeMegarepoConfigs;
#[cfg(not(fbcode_build))]
pub use oss::CfgrMononokeMegarepoConfigs;
pub use test_impl::TestMononokeMegarepoConfigs;
pub use verification::verify_config;

/// Options for instantiating MononokeMegarepoConfigs
#[derive(Clone, PartialEq, Eq)]
pub enum MononokeMegarepoConfigsOptions {
    /// Create prod-style `MononokeMegarepoConfigs` implementation
    /// (requires fb infra to function correctly, although will
    /// successfully instantiate with `unimplemented!` methods
    /// when built outside of fbcode)
    Prod,
    /// Create a config implementation that writes JSON to disk at the
    /// given path instead of calling FB infra.
    /// Used with a testing config store, this gives you a good basis
    /// for integration tests
    IntegrationTest(PathBuf),
    /// Create test-style `MononokeMegarepoConfigs` implementation
    UnitTest,
}

/// Command line arguments for controlling Megarepo configs
#[derive(Args, Debug)]
pub struct MegarepoConfigsArgs {
    /// Whether to instantiate test-style MononokeMegarepoConfigs
    ///
    /// Prod-style instance reads/writes from/to configerator and
    /// requires the FB environment to work properly.
    // For compatibility with existing usage, this arg takes value
    // for example `--with-test-megarepo-configs-client=true`.
    #[clap(long, default_value_t = false, value_name = "BOOL", action = ArgAction::Set)]
    pub with_test_megarepo_configs_client: bool,
}

impl MononokeMegarepoConfigsOptions {
    pub fn from_args(
        local_configerator_path: Option<&Path>,
        megarepo_configs_args: &MegarepoConfigsArgs,
    ) -> Self {
        if megarepo_configs_args.with_test_megarepo_configs_client {
            if let Some(path) = local_configerator_path {
                MononokeMegarepoConfigsOptions::IntegrationTest(path.to_path_buf())
            } else {
                MononokeMegarepoConfigsOptions::UnitTest
            }
        } else {
            MononokeMegarepoConfigsOptions::Prod
        }
    }
}

/// An API for Megarepo Configs
#[async_trait]
pub trait MononokeMegarepoConfigs: Send + Sync {
    /// Get all the versions for a given Target
    fn get_target_config_versions(
        &self,
        ctx: CoreContext,
        target: Target,
    ) -> Result<Vec<SyncConfigVersion>, MegarepoError>;

    /// Get a SyncTargetConfig by its version
    fn get_config_by_version(
        &self,
        ctx: CoreContext,
        target: Target,
        version: SyncConfigVersion,
    ) -> Result<SyncTargetConfig, MegarepoError>;

    /// Add a new unused SyncTargetConfig for an existing Target
    async fn add_config_version(
        &self,
        ctx: CoreContext,
        config: SyncTargetConfig,
    ) -> Result<(), MegarepoError>;
}


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

use async_trait::async_trait;
use cached_config::ConfigStore;
use context::CoreContext;
use fbinit::FacebookInit;
use megarepo_configs::types::SyncConfigVersion;
use megarepo_configs::types::SyncTargetConfig;
use megarepo_configs::types::Target;
use megarepo_error::MegarepoError;
use slog::warn;
use slog::Logger;

use crate::MononokeMegarepoConfigs;

/// This struct is a no-op stub, to allow for successful builds
/// outside of fbcode. While it is allowed to instantiate this struct,
/// (to allow creating of dependent structs that instantiate
/// MononokeMegarepoConfigs as part of their setup) it is illegal
/// to call any of the methods.
pub struct CfgrMononokeMegarepoConfigs;

impl CfgrMononokeMegarepoConfigs {
    pub fn new(
        _fb: FacebookInit,
        logger: &Logger,
        _config_store: ConfigStore,
        _test_write_path: Option<PathBuf>,
    ) -> Result<Self, MegarepoError> {
        warn!(
            logger,
            "CfgrMononokeMegarepoConfigs is not implemented for non-fbcode builds"
        );
        // While this struct should never be used in practice outside of fbcode, I think
        // it shouldn't be an error to simply instantiate it in non-fbcode
        // builds, so let's return `Ok` here.
        Ok(Self)
    }
}

#[async_trait]
impl MononokeMegarepoConfigs for CfgrMononokeMegarepoConfigs {
    fn get_target_config_versions(
        &self,
        _ctx: CoreContext,
        _target: Target,
    ) -> Result<Vec<SyncConfigVersion>, MegarepoError> {
        unimplemented!("OSS CfgrMononokeMegarepoConfigs::get_target_config_versions")
    }

    fn get_config_by_version(
        &self,
        _ctx: CoreContext,
        _target: Target,
        _version: SyncConfigVersion,
    ) -> Result<SyncTargetConfig, MegarepoError> {
        unimplemented!("OSS CfgrMononokeMegarepoConfigs::get_config_by_version")
    }

    async fn add_config_version(
        &self,
        _ctx: CoreContext,
        _config: SyncTargetConfig,
    ) -> Result<(), MegarepoError> {
        unimplemented!("OSS CfgrMononokeMegarepoConfigs::add_config_version")
    }
}


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

use anyhow::anyhow;
use async_trait::async_trait;
use context::CoreContext;
use megarepo_configs::types::SyncConfigVersion;
use megarepo_configs::types::SyncTargetConfig;
use megarepo_configs::types::Target;
use megarepo_error::MegarepoError;
use slog::info;
use slog::Logger;

use crate::verification::verify_config;
use crate::MononokeMegarepoConfigs;

#[derive(Clone)]
pub struct TestMononokeMegarepoConfigs {
    config_versions: Arc<Mutex<HashMap<(Target, SyncConfigVersion), SyncTargetConfig>>>,
}

impl TestMononokeMegarepoConfigs {
    pub fn new(logger: &Logger) -> Self {
        info!(logger, "Creating a new TestMononokeMegarepoConfigs");
        Self {
            config_versions: Arc::new(Mutex::new(HashMap::new())),
        }
    }

    pub fn add(&mut self, key: (Target, SyncConfigVersion), target: SyncTargetConfig) {
        let mut config_versions = self.config_versions.lock().unwrap();
        config_versions.insert(key, target);
    }
}

#[async_trait]
impl MononokeMegarepoConfigs for TestMononokeMegarepoConfigs {
    fn get_target_config_versions(
        &self,
        _ctx: CoreContext,
        _target: Target,
    ) -> Result<Vec<SyncConfigVersion>, MegarepoError> {
        unimplemented!("TestMononokeMegarepoConfigs::get_target_config_versions")
    }

    fn get_config_by_version(
        &self,
        _ctx: CoreContext,
        target: Target,
        version: SyncConfigVersion,
    ) -> Result<SyncTargetConfig, MegarepoError> {
        let config_versions = self.config_versions.lock().unwrap();
        config_versions
            .get(&(target.clone(), version.clone()))
            .cloned()
            .ok_or_else(|| anyhow!("{:?} not found", (target, version)))
            .map_err(MegarepoError::internal)
    }

    async fn add_config_version(
        &self,
        ctx: CoreContext,
        config: SyncTargetConfig,
    ) -> Result<(), MegarepoError> {
        verify_config(&ctx, &config).map_err(MegarepoError::request)?;
        let mut config_versions = self.config_versions.lock().unwrap();
        let key = (config.target.clone(), config.version.clone());
        config_versions.insert(key, config);
        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::HashSet;

use anyhow::anyhow;
use anyhow::Result;
use context::CoreContext;
use megarepo_configs::types::SyncTargetConfig;
use slog::warn;

/// Verify the config
pub fn verify_config(ctx: &CoreContext, config: &SyncTargetConfig) -> Result<()> {
    verify_unique_source_names(ctx, config)
}

fn verify_unique_source_names(ctx: &CoreContext, config: &SyncTargetConfig) -> Result<()> {
    let mut seen = HashSet::new();
    let mut seen_more_than_once = HashSet::new();
    config.sources.iter().for_each(|src| {
        if !seen.insert(&src.source_name) {
            seen_more_than_once.insert(&src.source_name);
        }
    });

    if !seen_more_than_once.is_empty() {
        warn!(
            ctx.logger(),
            "SyncTargetConfig validation error: non-unique source names: {:?}", seen_more_than_once
        );

        Err(anyhow!(
            "Non-unique source names: {:?}",
            seen_more_than_once
        ))
    } else {
        Ok(())
    }
}

#[cfg(test)]
mod verification_tests {
    use fbinit::FacebookInit;
    use maplit::btreemap;
    use megarepo_configs::types::MergeMode;
    use megarepo_configs::types::Source;
    use megarepo_configs::types::SourceMappingRules;
    use megarepo_configs::types::SourceRevision;
    use megarepo_configs::types::Target;
    use megarepo_configs::types::WithExtraMoveCommit;

    use super::*;

    fn s(v: &str) -> String {
        v.to_owned()
    }

    fn get_good_cfg() -> SyncTargetConfig {
        SyncTargetConfig {
            target: Target {
                repo_id: 1,
                bookmark: s("target"),
            },
            version: s("version"),
            sources: vec![
                Source {
                    name: s("name1"),
                    source_name: s("source1"),
                    revision: SourceRevision::bookmark(s("hello")),
                    repo_id: 1,
                    mapping: SourceMappingRules {
                        default_prefix: s("pre/fix1"),
                        linkfiles: btreemap! {
                            s("link/source_1") => s("link_target_1")
                        },
                        overrides: btreemap! {
                            s("deleted") => vec![],
                            s("multiplied1") => vec![s("multiplied1_1"), s("multiplied1_2")],
                            s("copyfile") => vec![
                                s("pre/fix1/copyfile"),
                                s("copiedfile"),
                            ]
                        },
                    },
                    merge_mode: Some(MergeMode::with_move_commit(WithExtraMoveCommit {
                        ..Default::default()
                    })),
                },
                Source {
                    name: s("name2"),
                    source_name: s("source2"),
                    revision: SourceRevision::bookmark(s("hello")),
                    repo_id: 1,
                    mapping: SourceMappingRules {
                        default_prefix: s("pre/fix2"),
                        linkfiles: btreemap! {
                            s("link/source_2") => s("link_target_2")
                        },
                        overrides: btreemap! {
                            s("deleted") => vec![],
                            s("multiplied2") => vec![s("multiplied2_1"), s("multiplied2_2")],
                        },
                    },
                    merge_mode: None,
                },
            ],
        }
    }

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

        let cfg = get_good_cfg();
        assert!(verify_unique_source_names(&ctx, &cfg).is_ok());

        Ok(())
    }

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

        let mut cfg = get_good_cfg();
        cfg.sources[1].source_name = s("source1");
        assert!(verify_unique_source_names(&ctx, &cfg).is_err());

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

#![feature(error_generic_member_access)]
#![feature(provide_any)]

use std::backtrace::BacktraceStatus;
use std::convert::Infallible;
use std::error::Error as StdError;
use std::sync::Arc;

use anyhow::Error;
use blobstore::LoadableError;
use source_control as scs_thrift;
use thiserror::Error;

pub mod macro_reexport {
    pub use anyhow::anyhow;
}

#[macro_export]
macro_rules! cloneable_error {
    ($name: ident) => {
        #[derive(Clone, Debug)]
        pub struct $name(pub ::std::sync::Arc<anyhow::Error>);

        impl $name {
            pub fn backtrace(&self) -> &::std::backtrace::Backtrace {
                self.0.backtrace()
            }
        }

        impl ::std::fmt::Display for $name {
            fn fmt(&self, f: &mut ::std::fmt::Formatter<'_>) -> ::std::fmt::Result {
                self.0.fmt(f)
            }
        }

        impl From<anyhow::Error> for $name {
            fn from(error: anyhow::Error) -> Self {
                Self(::std::sync::Arc::new(error))
            }
        }

        impl ::std::error::Error for $name {
            fn source(&self) -> Option<&(dyn ::std::error::Error + 'static)> {
                Some(&**self.0)
            }

            fn provide<'a>(&'a self, demand: &mut ::std::any::Demand<'a>) {
                demand.provide_ref::<::std::backtrace::Backtrace>(self.backtrace());
            }
        }
    };
}

cloneable_error!(InternalError);
cloneable_error!(RequestError);

#[derive(Clone, Debug, Error)]
pub enum MegarepoError {
    #[error("{0}")]
    RequestError(#[source] RequestError),
    #[error("{0}")]
    InternalError(#[source] InternalError),
}

impl MegarepoError {
    pub fn internal(e: impl Into<Error>) -> Self {
        Self::InternalError(InternalError::from(e.into()))
    }

    pub fn request(e: impl Into<Error>) -> Self {
        Self::RequestError(RequestError::from(e.into()))
    }
}

/// By default, let's treat errors as internal
impl From<Error> for MegarepoError {
    fn from(e: Error) -> Self {
        match e.downcast::<MegarepoError>() {
            Ok(megarepo_error) => match megarepo_error {
                Self::RequestError(e) => Self::RequestError(e),
                Self::InternalError(e) => Self::InternalError(e),
            },
            Err(orig) => Self::internal(orig),
        }
    }
}

impl From<Infallible> for MegarepoError {
    fn from(_i: Infallible) -> Self {
        unreachable!()
    }
}

impl From<LoadableError> for MegarepoError {
    fn from(e: LoadableError) -> Self {
        MegarepoError::InternalError(InternalError(Arc::new(e.into())))
    }
}

#[macro_export]
macro_rules! bail_request {
    ($msg:literal $(,)?) => {
        return Err($crate::MegarepoError::RequestError($crate::RequestError(::std::sync::Arc::new($crate::macro_reexport::anyhow!($msg)))))
    };
    ($fmt:expr, $($arg:tt)*) => {
        return Err($crate::MegarepoError::RequestError($crate::RequestError(::std::sync::Arc::new($crate::macro_reexport::anyhow!($fmt, $($arg)*)))))
    };
}

#[macro_export]
macro_rules! bail_internal {
    ($msg:literal $(,)?) => {
        return Err($crate::MegarepoError::InternalError($crate::InternalError(::std::sync::Arc::new($crate::macro_reexport::anyhow!($msg)))))
    };
    ($fmt:expr, $($arg:tt)*) => {
        return Err($crate::MegarepoError::InternalError($crate::InternalError(::std::sync::Arc::new($crate::macro_reexport::anyhow!($fmt, $($arg)*)))))
    };
}

impl From<MegarepoError> for scs_thrift::MegarepoAsynchronousRequestError {
    fn from(e: MegarepoError) -> Self {
        match e {
            MegarepoError::RequestError(e) => Self::request_error(scs_thrift::RequestErrorStruct {
                kind: scs_thrift::RequestErrorKind::INVALID_REQUEST,
                reason: format!("{}", e),
                ..Default::default()
            }),
            MegarepoError::InternalError(error) => {
                let reason = error.to_string();
                let backtrace = match error.backtrace().status() {
                    BacktraceStatus::Captured => Some(error.backtrace().to_string()),
                    _ => None,
                };
                let mut source_chain = Vec::new();
                let mut error: &dyn StdError = &error;
                while let Some(source) = error.source() {
                    source_chain.push(source.to_string());
                    error = source;
                }

                Self::internal_error(scs_thrift::InternalErrorStruct {
                    reason,
                    backtrace,
                    source_chain,
                    ..Default::default()
                })
            }
        }
    }
}


/*
 * 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::Result;
use async_trait::async_trait;
use bookmarks::BookmarkKey;
use context::CoreContext;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;

mod store;
mod types;

pub use crate::store::SqlLongRunningRequestsQueue;
pub use crate::types::BlobstoreKey;
pub use crate::types::ClaimedBy;
pub use crate::types::LongRunningRequestEntry;
pub use crate::types::RequestId;
pub use crate::types::RequestStatus;
pub use crate::types::RequestType;
pub use crate::types::RowId;

/// A queue of long-running requests
/// This is designed to support the use case of
/// asynchronous request processing, when a client
/// service schedules a request to be processed
/// and later checks if the result is ready.
/// Another service handles the processing and
/// state updates for individual requests.
#[facet::facet]
#[async_trait]
pub trait LongRunningRequestsQueue: Send + Sync {
    /// Schedule an execution of a request, given the request type, the blobstore
    /// key of serialized request parameters and (repo, bookmark) pair
    /// representing the target of the request.
    async fn add_request(
        &self,
        ctx: &CoreContext,
        request_type: &RequestType,
        repo_id: &RepositoryId,
        bookmark: &BookmarkKey,
        args_blobstore_key: &BlobstoreKey,
    ) -> Result<RowId>;

    /// Claim one of new requests. Mark it as in-progress and return it.
    async fn claim_and_get_new_request(
        &self,
        ctx: &CoreContext,
        claimed_by: &ClaimedBy,
        supported_repos: &[RepositoryId],
    ) -> Result<Option<LongRunningRequestEntry>>;

    /// Get the full request object entry by id
    /// Since this does not take `request_type`, it is
    /// mainly intended to be used in tests (`request_type`
    /// is a type-safety feature)
    async fn test_get_request_entry_by_id(
        &self,
        ctx: &CoreContext,
        id: &RowId,
    ) -> Result<Option<LongRunningRequestEntry>>;

    /// Mark request as in-progress
    async fn mark_in_progress(
        &self,
        ctx: &CoreContext,
        req_id: &RequestId,
        claimed_by: &ClaimedBy,
    ) -> Result<bool>;

    /// Update the inprogress_last_updated_at timestamp
    /// This is used to mark that request is still executing
    async fn update_in_progress_timestamp(
        &self,
        ctx: &CoreContext,
        req_id: &RequestId,
    ) -> Result<bool>;

    /// Find requests that have "inprogress" status but which timestamp
    /// hasn't been updated after `abandoned_timestamp`.
    async fn find_abandoned_requests(
        &self,
        ctx: &CoreContext,
        repo_ids: &[RepositoryId],
        abandoned_timestamp: Timestamp,
    ) -> Result<Vec<RequestId>>;

    /// If `request_id` is still abandoned, then mark it as new so that
    /// somebody else can pick it up
    async fn mark_abandoned_request_as_new(
        &self,
        ctx: &CoreContext,
        request_id: RequestId,
        abandoned_timestamp: Timestamp,
    ) -> Result<bool>;

    /// Mark request as ready
    async fn mark_ready(
        &self,
        ctx: &CoreContext,
        req_id: &RequestId,
        blobstore_result_key: BlobstoreKey,
    ) -> Result<bool>;

    /// Mark request as new (used for requeuing requests from CLI)
    async fn mark_new(&self, ctx: &CoreContext, req_id: &RequestId) -> Result<bool>;

    /// Mark request as polled by a client
    /// To be used in tests only
    async fn test_mark(
        &self,
        ctx: &CoreContext,
        row_id: &RowId,
        status: RequestStatus,
    ) -> Result<bool>;

    /// Query request and change it's state to `polled`
    /// if it is ready. This fn will return a corresponding
    /// LongRunningRequestEntry if the request is in `ready`
    /// or `polled` state at the time of the call
    /// Otherwise, it will return None. It will error out
    /// for unknown requests
    async fn poll(
        &self,
        ctx: &CoreContext,
        req_id: &RequestId,
    ) -> Result<Option<(bool, LongRunningRequestEntry)>>;

    async fn list_requests(
        &self,
        ctx: &CoreContext,
        repo_ids: &[RepositoryId],
        last_update_newer_than: Option<&Timestamp>,
    ) -> Result<Vec<LongRunningRequestEntry>>;
}


/*
 * 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::bail;
use anyhow::Result;
use async_trait::async_trait;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkName;
use context::CoreContext;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::mononoke_queries;
use sql_ext::SqlConnections;

use crate::BlobstoreKey;
use crate::ClaimedBy;
use crate::LongRunningRequestEntry;
use crate::LongRunningRequestsQueue;
use crate::RequestId;
use crate::RequestStatus;
use crate::RequestType;
use crate::RowId;

mononoke_queries! {
    read TestGetRequest(id: RowId) -> (
        RowId,
        RequestType,
        RepositoryId,
        BookmarkName,
        BlobstoreKey,
        Option<BlobstoreKey>,
        Timestamp,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        RequestStatus,
        Option<ClaimedBy>,
    ) {
        "SELECT id,
            request_type,
            repo_id,
            bookmark,
            args_blobstore_key,
            result_blobstore_key,
            created_at,
            started_processing_at,
            inprogress_last_updated_at,
            ready_at,
            polled_at,
            status,
            claimed_by
        FROM long_running_request_queue
        WHERE id = {id}"
    }

    read GetRequest(id: RowId, request_type: RequestType) -> (
        RowId,
        RequestType,
        RepositoryId,
        BookmarkName,
        BlobstoreKey,
        Option<BlobstoreKey>,
        Timestamp,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        RequestStatus,
        Option<ClaimedBy>,
    ) {
        "SELECT id,
            request_type,
            repo_id,
            bookmark,
            args_blobstore_key,
            result_blobstore_key,
            created_at,
            started_processing_at,
            inprogress_last_updated_at,
            ready_at,
            polled_at,
            status,
            claimed_by
        FROM long_running_request_queue
        WHERE id = {id} AND request_type = {request_type}"
    }

    read GetOneNewRequestForRepos(>list supported_repo_ids: RepositoryId) -> (
        RowId,
        RequestType,
        RepositoryId,
        BookmarkName,
        BlobstoreKey,
        Option<BlobstoreKey>,
        Timestamp,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        RequestStatus,
        Option<ClaimedBy>,
    ) {
        "SELECT id,
            request_type,
            repo_id,
            bookmark,
            args_blobstore_key,
            result_blobstore_key,
            created_at,
            started_processing_at,
            inprogress_last_updated_at,
            ready_at,
            polled_at,
            status,
            claimed_by
        FROM long_running_request_queue
        WHERE status = 'new' AND repo_id IN {supported_repo_ids}
        ORDER BY created_at ASC
        LIMIT 1
        "
    }

    write AddRequest(request_type: RequestType, repo_id: RepositoryId, bookmark: BookmarkName, args_blobstore_key: BlobstoreKey, created_at: Timestamp) {
        none,
        "INSERT INTO long_running_request_queue
         (request_type, repo_id, bookmark, args_blobstore_key, status, created_at)
         VALUES ({request_type}, {repo_id}, {bookmark}, {args_blobstore_key}, 'new', {created_at})
        "
    }

    read FindAbandonedRequests(
        abandoned_timestamp: Timestamp,
        >list repo_ids: RepositoryId
    ) -> (RowId, RequestType) {
        "
        SELECT id, request_type
        FROM long_running_request_queue
        WHERE repo_id IN {repo_ids} AND status = 'inprogress' AND inprogress_last_updated_at <= {abandoned_timestamp}
        "
    }

    write MarkRequestAsNewAgainIfAbandoned(
        id: RowId,
        request_type: RequestType,
        abandoned_timestamp: Timestamp,
    )
    {
        none,
        "UPDATE long_running_request_queue
         SET status = 'new', claimed_by = NULL, inprogress_last_updated_at = NULL
         WHERE id = {id} AND request_type = {request_type} AND status = 'inprogress' AND inprogress_last_updated_at <= {abandoned_timestamp}
        "
    }

    write MarkRequestInProgress(
        id: RowId,
        request_type: RequestType,
        started_processing_at: Timestamp,
        claimed_by: ClaimedBy,
    ) {
        none,
        "UPDATE long_running_request_queue
         SET started_processing_at = {started_processing_at},
             inprogress_last_updated_at = {started_processing_at},
             status = 'inprogress',
             claimed_by = {claimed_by}
         WHERE id = {id} AND request_type = {request_type} AND status = 'new'
        "
    }

    write UpdateInProgressTimestamp(
        id: RowId,
        request_type: RequestType,
        inprogress_last_updated_at: Timestamp,
    ) {
        none,
        "UPDATE long_running_request_queue
         SET inprogress_last_updated_at = {inprogress_last_updated_at}
         WHERE id = {id} AND request_type = {request_type} AND status = 'inprogress'
        "
    }

    write MarkRequestReady(id: RowId, request_type: RequestType, ready_at: Timestamp, result_blobstore_key: BlobstoreKey) {
        none,
        "UPDATE long_running_request_queue
         SET ready_at = {ready_at}, status = 'ready', result_blobstore_key = {result_blobstore_key}
         WHERE id = {id} AND request_type = {request_type} AND status = 'inprogress'
        "
    }

    write MarkRequestAsNew(id: RowId, request_type: RequestType) {
        none,
        "UPDATE long_running_request_queue
         SET status = 'new'
         WHERE id = {id} AND request_type = {request_type}
        "
    }

    write MarkRequestPolled(id: RowId, request_type: RequestType, polled_at: Timestamp) {
        none,
        "UPDATE long_running_request_queue
         SET polled_at = {polled_at}, status = 'polled'
         WHERE id = {id} AND request_type = {request_type} AND status = 'ready'
        "
    }

    write TestMark(id: RowId, status: RequestStatus) {
        none,
        "UPDATE long_running_request_queue
         SET status = {status}
         WHERE id = {id}
        "
    }

    read ListRequests(last_udate_newer_than: Timestamp, >list repo_ids: RepositoryId) -> (
        RowId,
        RequestType,
        RepositoryId,
        BookmarkName,
        BlobstoreKey,
        Option<BlobstoreKey>,
        Timestamp,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        RequestStatus,
        Option<ClaimedBy>,
    ) {
        "SELECT id,
            request_type,
            repo_id,
            bookmark,
            args_blobstore_key,
            result_blobstore_key,
            created_at,
            started_processing_at,
            inprogress_last_updated_at,
            ready_at,
            polled_at,
            status,
            claimed_by
        FROM long_running_request_queue
        WHERE repo_id IN {repo_ids} AND inprogress_last_updated_at > {last_udate_newer_than}"
    }
}

fn row_to_entry(
    row: (
        RowId,
        RequestType,
        RepositoryId,
        BookmarkName,
        BlobstoreKey,
        Option<BlobstoreKey>,
        Timestamp,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        Option<Timestamp>,
        RequestStatus,
        Option<ClaimedBy>,
    ),
) -> LongRunningRequestEntry {
    let (
        id,
        request_type,
        repo_id,
        bookmark,
        args_blobstore_key,
        result_blobstore_key,
        created_at,
        started_processing_at,
        inprogress_last_updated_at,
        ready_at,
        polled_at,
        status,
        claimed_by,
    ) = row;
    LongRunningRequestEntry {
        id,
        repo_id,
        bookmark,
        request_type,
        args_blobstore_key,
        result_blobstore_key,
        created_at,
        started_processing_at,
        inprogress_last_updated_at,
        ready_at,
        polled_at,
        status,
        claimed_by,
    }
}

#[derive(Clone)]
pub struct SqlLongRunningRequestsQueue {
    pub(crate) connections: SqlConnections,
}

#[async_trait]
impl LongRunningRequestsQueue for SqlLongRunningRequestsQueue {
    async fn add_request(
        &self,
        _ctx: &CoreContext,
        request_type: &RequestType,
        repo_id: &RepositoryId,
        bookmark: &BookmarkKey,
        args_blobstore_key: &BlobstoreKey,
    ) -> Result<RowId> {
        let res = AddRequest::query(
            &self.connections.write_connection,
            request_type,
            repo_id,
            bookmark.name(),
            args_blobstore_key,
            &Timestamp::now(),
        )
        .await?;

        match res.last_insert_id() {
            Some(last_insert_id) if res.affected_rows() == 1 => Ok(RowId(last_insert_id)),
            _ => bail!("Failed to insert a new request of type {}", request_type),
        }
    }

    /// Claim one of new requests. Mark it as in-progress and return it.
    async fn claim_and_get_new_request(
        &self,
        ctx: &CoreContext,
        claimed_by: &ClaimedBy,
        supported_repos: &[RepositoryId],
    ) -> Result<Option<LongRunningRequestEntry>> {
        // Spin until we win the race or there's nothing to do.
        loop {
            let rows = GetOneNewRequestForRepos::query(
                &self.connections.read_master_connection, // reaching DB master improves our chances.
                supported_repos,
            )
            .await?;
            let mut entry = match rows.into_iter().next() {
                None => {
                    return Ok(None);
                }
                Some(row) => row_to_entry(row),
            };
            if self
                .mark_in_progress(
                    ctx,
                    &RequestId(entry.id, entry.request_type.clone()),
                    claimed_by,
                )
                .await?
            {
                // Success, we won the race!
                entry.status = RequestStatus::InProgress;
                return Ok(Some(entry));
            }
            // Failure, let's try again.
        }
    }

    async fn test_get_request_entry_by_id(
        &self,
        _ctx: &CoreContext,
        id: &RowId,
    ) -> Result<Option<LongRunningRequestEntry>> {
        let rows = TestGetRequest::query(&self.connections.read_connection, id).await?;
        match rows.into_iter().next() {
            None => Ok(None),
            Some(row) => Ok(Some(row_to_entry(row))),
        }
    }

    async fn mark_in_progress(
        &self,
        _ctx: &CoreContext,
        req_id: &RequestId,
        claimed_by: &ClaimedBy,
    ) -> Result<bool> {
        let res = MarkRequestInProgress::query(
            &self.connections.write_connection,
            &req_id.0,
            &req_id.1,
            &Timestamp::now(),
            claimed_by,
        )
        .await?;
        Ok(res.affected_rows() > 0)
    }

    async fn update_in_progress_timestamp(
        &self,
        _ctx: &CoreContext,
        req_id: &RequestId,
    ) -> Result<bool> {
        let res = UpdateInProgressTimestamp::query(
            &self.connections.write_connection,
            &req_id.0,
            &req_id.1,
            &Timestamp::now(),
        )
        .await?;
        Ok(res.affected_rows() > 0)
    }

    async fn find_abandoned_requests(
        &self,
        _ctx: &CoreContext,
        repo_ids: &[RepositoryId],
        abandoned_timestamp: Timestamp,
    ) -> Result<Vec<RequestId>> {
        let rows = FindAbandonedRequests::query(
            &self.connections.write_connection,
            &abandoned_timestamp,
            repo_ids,
        )
        .await?;
        Ok(rows.into_iter().map(|(id, ty)| RequestId(id, ty)).collect())
    }

    async fn mark_abandoned_request_as_new(
        &self,
        _ctx: &CoreContext,
        request_id: RequestId,
        abandoned_timestamp: Timestamp,
    ) -> Result<bool> {
        let res = MarkRequestAsNewAgainIfAbandoned::query(
            &self.connections.write_connection,
            &request_id.0,
            &request_id.1,
            &abandoned_timestamp,
        )
        .await?;

        Ok(res.affected_rows() > 0)
    }

    async fn mark_ready(
        &self,
        _ctx: &CoreContext,
        req_id: &RequestId,
        blobstore_result_key: BlobstoreKey,
    ) -> Result<bool> {
        let res = MarkRequestReady::query(
            &self.connections.write_connection,
            &req_id.0,
            &req_id.1,
            &Timestamp::now(),
            &blobstore_result_key,
        )
        .await?;

        Ok(res.affected_rows() > 0)
    }

    async fn mark_new(&self, _ctx: &CoreContext, req_id: &RequestId) -> Result<bool> {
        let res = MarkRequestAsNew::query(&self.connections.write_connection, &req_id.0, &req_id.1)
            .await?;

        Ok(res.affected_rows() > 0)
    }

    async fn test_mark(
        &self,
        _ctx: &CoreContext,
        row_id: &RowId,
        status: RequestStatus,
    ) -> Result<bool> {
        let res = TestMark::query(&self.connections.write_connection, row_id, &status).await?;
        Ok(res.affected_rows() > 0)
    }

    async fn poll(
        &self,
        _ctx: &CoreContext,
        req_id: &RequestId,
    ) -> Result<Option<(bool, LongRunningRequestEntry)>> {
        let txn = self
            .connections
            .write_connection
            .start_transaction()
            .await?;
        let (mut txn, rows) = GetRequest::query_with_transaction(txn, &req_id.0, &req_id.1).await?;
        let entry = match rows.into_iter().next() {
            None => bail!("unknown request polled: {:?}", req_id),
            Some(row) => {
                let mut entry = row_to_entry(row);

                match &entry.status {
                    RequestStatus::Ready | RequestStatus::Polled
                        if entry.result_blobstore_key.is_none() =>
                    {
                        bail!(
                            "no result stored for {:?} request {:?}",
                            entry.status,
                            req_id
                        );
                    }
                    RequestStatus::Ready => {
                        txn = MarkRequestPolled::query_with_transaction(
                            txn,
                            &req_id.0,
                            &req_id.1,
                            &Timestamp::now(),
                        )
                        .await?
                        .0;

                        entry.status = RequestStatus::Polled;
                        Some((true, entry))
                    }
                    RequestStatus::Polled => Some((false, entry)),
                    _ => None,
                }
            }
        };
        txn.commit().await?;
        Ok(entry)
    }

    async fn list_requests(
        &self,
        _ctx: &CoreContext,
        repo_ids: &[RepositoryId],
        last_update_newer_than: Option<&Timestamp>,
    ) -> Result<Vec<LongRunningRequestEntry>> {
        let entries = ListRequests::query(
            &self.connections.read_connection,
            last_update_newer_than.unwrap_or(&Timestamp::from_timestamp_nanos(0)),
            repo_ids,
        )
        .await?
        .into_iter()
        .map(row_to_entry)
        .collect();
        Ok(entries)
    }
}

impl SqlConstruct for SqlLongRunningRequestsQueue {
    const LABEL: &'static str = "long_running_requests_queue";

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

    fn from_sql_connections(connections: SqlConnections) -> Self {
        Self { connections }
    }
}

impl SqlConstructFromMetadataDatabaseConfig for SqlLongRunningRequestsQueue {}

#[cfg(test)]
mod test {
    use std::time::Duration;

    use fbinit::FacebookInit;

    use super::*;

    #[fbinit::test]
    async fn test_mark_inprogress(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let queue = SqlLongRunningRequestsQueue::with_sqlite_in_memory()?;
        let id = queue
            .add_request(
                &ctx,
                &RequestType("type".to_string()),
                &RepositoryId::new(0),
                &BookmarkKey::new("book")?,
                &BlobstoreKey("key".to_string()),
            )
            .await?;

        let request = queue.test_get_request_entry_by_id(&ctx, &id).await?;
        assert!(request.is_some());
        let request = request.unwrap();
        assert!(request.inprogress_last_updated_at.is_none());

        queue
            .claim_and_get_new_request(&ctx, &ClaimedBy("me".to_string()), &[RepositoryId::new(0)])
            .await?;

        let request = queue.test_get_request_entry_by_id(&ctx, &id).await?;
        assert!(request.is_some());
        let request = request.unwrap();
        assert!(request.inprogress_last_updated_at.is_some());

        let timestamp = request.inprogress_last_updated_at.unwrap();

        tokio::time::sleep(Duration::from_secs(3)).await;

        let updated = queue
            .update_in_progress_timestamp(&ctx, &RequestId(request.id, request.request_type))
            .await?;
        assert!(updated);
        let request = queue.test_get_request_entry_by_id(&ctx, &id).await?;
        // Check that timestamp was updated
        assert!(request.unwrap().inprogress_last_updated_at.unwrap() > timestamp);

        Ok(())
    }

    #[fbinit::test]
    async fn test_find_abandoned_requests(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let queue = SqlLongRunningRequestsQueue::with_sqlite_in_memory()?;
        let repo_id = RepositoryId::new(0);
        let id = queue
            .add_request(
                &ctx,
                &RequestType("type".to_string()),
                &repo_id,
                &BookmarkKey::new("book")?,
                &BlobstoreKey("key".to_string()),
            )
            .await?;

        // This claims new request from queue and makes it inprogress
        let req = queue
            .claim_and_get_new_request(&ctx, &ClaimedBy("me".to_string()), &[RepositoryId::new(0)])
            .await?;
        assert!(req.is_some());

        tokio::time::sleep(Duration::from_secs(3)).await;
        let now = Timestamp::now();
        let abandoned_timestamp = Timestamp::from_timestamp_secs(now.timestamp_seconds() - 1);
        let abandoned = queue
            .find_abandoned_requests(&ctx, &[repo_id], abandoned_timestamp)
            .await?;
        assert_eq!(abandoned.len(), 1);
        assert_eq!(abandoned[0].0, id);

        // Now update timestamp of the request, and check that it's not considered
        // abandoned anymore
        let updated = queue
            .update_in_progress_timestamp(&ctx, &abandoned[0])
            .await?;
        assert!(updated);
        assert_eq!(
            queue
                .find_abandoned_requests(&ctx, &[repo_id], abandoned_timestamp)
                .await?,
            vec![]
        );

        // Now mark ready first, then make sure that update_in_progress_timestamp does nothing
        assert!(
            queue
                .mark_ready(&ctx, &abandoned[0], BlobstoreKey("key".to_string()))
                .await?
        );
        assert!(
            !queue
                .update_in_progress_timestamp(&ctx, &abandoned[0])
                .await?
        );

        // And also check that marking as new does nothing on a completed request
        tokio::time::sleep(Duration::from_secs(3)).await;
        let now = Timestamp::now();
        let abandoned_timestamp = Timestamp::from_timestamp_secs(now.timestamp_seconds() - 1);
        assert!(
            !queue
                .mark_abandoned_request_as_new(&ctx, abandoned[0].clone(), abandoned_timestamp)
                .await?
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_mark_as_new(fb: FacebookInit) -> Result<()> {
        let ctx = CoreContext::test_mock(fb);
        let queue = SqlLongRunningRequestsQueue::with_sqlite_in_memory()?;
        let repo_id = RepositoryId::new(0);
        let id = queue
            .add_request(
                &ctx,
                &RequestType("type".to_string()),
                &repo_id,
                &BookmarkKey::new("book")?,
                &BlobstoreKey("key".to_string()),
            )
            .await?;

        // This claims new request from queue and makes it inprogress
        let req = queue
            .claim_and_get_new_request(&ctx, &ClaimedBy("me".to_string()), &[RepositoryId::new(0)])
            .await?;
        assert!(req.is_some());

        tokio::time::sleep(Duration::from_secs(3)).await;
        let now = Timestamp::now();
        let abandoned_timestamp = Timestamp::from_timestamp_secs(now.timestamp_seconds() - 1);
        let abandoned = queue
            .find_abandoned_requests(&ctx, &[repo_id], abandoned_timestamp)
            .await?;
        assert_eq!(abandoned.len(), 1);
        assert_eq!(abandoned[0].0, id);

        let res = queue
            .mark_abandoned_request_as_new(&ctx, abandoned[0].clone(), abandoned_timestamp)
            .await?;
        assert!(res);

        let request = queue.test_get_request_entry_by_id(&ctx, &id).await?;
        assert!(request.is_some());
        let request = request.unwrap();
        assert_eq!(request.status, RequestStatus::New);

        Ok(())
    }
}