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

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

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use blobrepo::save_bonsai_changesets;
use blobstore::Loadable;
use changeset_fetcher::ChangesetFetcherRef;
use changesets::ChangesetsRef;
use commit_transformation::create_directory_source_to_target_multi_mover;
use commit_transformation::create_source_to_target_multi_mover;
use commit_transformation::rewrite_as_squashed_commit;
use commit_transformation::rewrite_commit;
use commit_transformation::upload_commits;
use context::CoreContext;
use futures::stream;
use futures::StreamExt;
use futures::TryStreamExt;
use megarepo_config::MononokeMegarepoConfigs;
use megarepo_config::Source;
use megarepo_config::SourceMappingRules;
use megarepo_config::SourceRevision;
use megarepo_config::Target;
use megarepo_error::MegarepoError;
use megarepo_mapping::CommitRemappingState;
use megarepo_mapping::MegarepoMapping;
use megarepo_mapping::SourceName;
use mononoke_api::ChangesetContext;
use mononoke_api::Mononoke;
use mononoke_api::RepoContext;
use mononoke_types::BonsaiChangeset;
use mononoke_types::ChangesetId;
use mutable_renames::MutableRenames;
use repo_blobstore::RepoBlobstoreRef;

use crate::common::find_source_config;
use crate::common::find_target_bookmark_and_value;
use crate::common::find_target_sync_config;
use crate::common::MegarepoOp;
use crate::common::SourceAndMovedChangesets;
use crate::Repo;

pub(crate) struct SyncChangeset<'a> {
    megarepo_configs: &'a Arc<dyn MononokeMegarepoConfigs>,
    mononoke: &'a Arc<Mononoke>,
    target_megarepo_mapping: &'a Arc<MegarepoMapping>,
    mutable_renames: &'a Arc<MutableRenames>,
}

#[async_trait]
impl<'a> MegarepoOp for SyncChangeset<'a> {
    fn mononoke(&self) -> &Arc<Mononoke> {
        self.mononoke
    }
}

pub enum MergeMode {
    Squashed {
        side_commits: Vec<ChangesetContext>,
    },
    ExtraMoveCommits {
        side_parents_move_commits: Vec<SourceAndMovedChangesets>,
    },
}

fn get_squashing_overrides(repo_name: &str, target_bookmark: &str) -> (Option<i64>, Option<bool>) {
    let targets = tunables::tunables()
        .by_repo_megarepo_squashing_config_override_targets(repo_name)
        .unwrap_or_default();
    if targets
        .iter()
        .any(|target| target.as_str() == target_bookmark)
    {
        (
            tunables::tunables().by_repo_megarepo_override_squashing_limit(repo_name),
            tunables::tunables().by_repo_megarepo_override_author_check(repo_name),
        )
    } else {
        (None, None)
    }
}

pub struct SquashingConfig {
    squash_limit: usize,
    check_author: bool,
}

const MERGE_COMMIT_MOVES_CONCURRENCY: usize = 10;

impl<'a> SyncChangeset<'a> {
    pub(crate) fn new(
        megarepo_configs: &'a Arc<dyn MononokeMegarepoConfigs>,
        mononoke: &'a Arc<Mononoke>,
        target_megarepo_mapping: &'a Arc<MegarepoMapping>,
        mutable_renames: &'a Arc<MutableRenames>,
    ) -> Self {
        Self {
            megarepo_configs,
            mononoke,
            target_megarepo_mapping,
            mutable_renames,
        }
    }

    pub(crate) async fn sync(
        &self,
        ctx: &CoreContext,
        source_cs_id: ChangesetId,
        source_name: &SourceName,
        target: &Target,
        target_location: ChangesetId,
    ) -> Result<ChangesetId, MegarepoError> {
        let target_repo = self.find_repo_by_id(ctx, target.repo_id).await?;

        // Now we need to find the target config version that was used to create the latest
        // target commit. This config version will be used to sync the new changeset
        let (_, actual_target_location) =
            find_target_bookmark_and_value(ctx, &target_repo, target).await?;

        if target_location != actual_target_location {
            // Check if previous call was successful and return result if so
            return self
                .check_if_this_method_has_already_succeeded(
                    ctx,
                    source_cs_id,
                    source_name,
                    (target_location, actual_target_location),
                    &target_repo,
                )
                .await;
        }

        let (commit_remapping_state, target_config) = find_target_sync_config(
            ctx,
            target_repo.inner_repo(),
            target_location,
            target,
            self.megarepo_configs,
        )
        .await?;

        // Given the SyncTargetConfig, let's find config for the source
        // we are going to sync from
        let source_config = find_source_config(source_name, &target_config)?;

        // Find source repo and changeset that we need to sync
        let source_repo = self.find_repo_by_id(ctx, source_config.repo_id).await?;
        let source_cs = source_cs_id
            .load(ctx, source_repo.inner_repo().repo_blobstore())
            .await?;

        validate_can_sync_changeset(
            ctx,
            target,
            &source_cs,
            &commit_remapping_state,
            &source_repo,
            source_config,
        )
        .await?;

        // In case of merge commits we need to add move commits on top of the
        // merged-in commits or squash side-branch.
        let maybe_squashing_config = match &source_config.merge_mode {
            Some(megarepo_config::MergeMode::squashed(sq)) => {
                let (maybe_squash_limit, maybe_check_author) =
                    get_squashing_overrides(target_repo.name(), &target.bookmark);
                Some(SquashingConfig {
                    squash_limit: maybe_squash_limit
                        .unwrap_or(sq.squash_limit)
                        .try_into()
                        .context("couldn't convert squash commits limit")?,
                    check_author: maybe_check_author.unwrap_or(true),
                })
            }
            None | Some(megarepo_config::MergeMode::with_move_commit(_)) => None,
            Some(megarepo_config::MergeMode::UnknownField(_)) => {
                return Err(anyhow!("Unknown MergeMode").into());
            }
        };

        let (is_squashable, side_commits) = if let Some(sq) = maybe_squashing_config {
            self.is_commit_squashable(
                target,
                source_cs_id,
                &commit_remapping_state,
                source_name,
                &source_repo,
                &sq,
            )
            .await?
        } else {
            (false, vec![])
        };

        let merge_mode = if is_squashable {
            MergeMode::Squashed { side_commits }
        } else {
            MergeMode::ExtraMoveCommits {
                side_parents_move_commits: self
                    .create_move_commits(
                        ctx,
                        target,
                        &source_cs,
                        &commit_remapping_state,
                        &source_repo,
                        source_name,
                        source_config,
                    )
                    .await?,
            }
        };

        // Finally create a commit in the target and update the mapping.
        let source_cs_id = source_cs.get_changeset_id();
        let new_target_cs_id = sync_changeset_to_target(
            ctx,
            &source_config.mapping,
            source_name,
            source_repo.inner_repo(),
            source_cs,
            target_repo.inner_repo(),
            target_location,
            target,
            commit_remapping_state,
            merge_mode,
        )
        .await?;

        self.target_megarepo_mapping
            .insert_source_target_cs_mapping(
                ctx,
                source_name,
                target,
                source_cs_id,
                new_target_cs_id,
                &target_config.version,
            )
            .await?;

        // Move the bookmark and record latest synced source changeset
        self.move_bookmark_conditionally(
            ctx,
            target_repo.inner_repo(),
            target.bookmark.clone(),
            (target_location, new_target_cs_id),
        )
        .await?;

        Ok(new_target_cs_id)
    }

    async fn is_commit_squashable(
        &self,
        target: &Target,
        source_cs_id: ChangesetId,
        commit_remapping_state: &CommitRemappingState,
        source_name: &SourceName,
        source_repo: &RepoContext,
        squashing_config: &SquashingConfig,
    ) -> Result<(bool, Vec<ChangesetContext>)> {
        if squashing_config.squash_limit == 0 {
            return Ok((false, vec![]));
        }

        let latest_synced_cs_id =
            find_latest_synced_cs_id(commit_remapping_state, source_name, target)?;

        // Take all ancestors of source_cs_id which are not ancestors of
        // latest_synced_cs_id up to limit+1 commits.
        let side_commits: Vec<_> = source_repo
            .difference_of_unions_of_ancestors(vec![source_cs_id], vec![latest_synced_cs_id])
            .await?
            .skip(1) // The source_cs_id is always the first returned commit.
            .take(squashing_config.squash_limit + 1) // Get one above limit so we know if there's more than limit.
            .try_collect()
            .await?;

        // Bail if we got nothing to do
        if side_commits.is_empty() {
            return Ok((false, vec![]));
        }

        // Bail if we got over limit of commits
        if side_commits.len() > squashing_config.squash_limit {
            return Ok((false, vec![]));
        }

        // The author we'll be comparing to ensure all commits are the same author.
        if squashing_config.check_author {
            let author = match side_commits.first() {
                Some(p) => p.author().await?,
                None => {
                    return Ok((false, vec![]));
                }
            };

            for side_commit in side_commits.iter() {
                // Author is different, bail
                if author != side_commit.author().await? {
                    return Ok((false, vec![]));
                }
            }
        }

        // All checks passed, squashing is possible
        Ok((true, side_commits))
    }

    // Creates move commits on top of the merge parents in the source that
    // hasn't already been synced to targets (all but one). These move commits
    // put all source files into a correct places in a target so the file
    // history is correct.
    async fn create_move_commits(
        &self,
        ctx: &CoreContext,
        target: &Target,
        source_cs: &BonsaiChangeset,
        commit_remapping_state: &CommitRemappingState,
        target_repo: &RepoContext,
        source_name: &SourceName,
        source: &Source,
    ) -> Result<Vec<SourceAndMovedChangesets>, MegarepoError> {
        let latest_synced_cs_id =
            find_latest_synced_cs_id(commit_remapping_state, source_name, target)?;

        // All parents except the one that's already synced to the target
        let side_parents = source_cs.parents().filter(|p| *p != latest_synced_cs_id);
        let mover = create_source_to_target_multi_mover(source.mapping.clone())
            .map_err(MegarepoError::request)?;
        let directory_mover = create_directory_source_to_target_multi_mover(source.mapping.clone())
            .map_err(MegarepoError::request)?;
        let moved_commits = stream::iter(side_parents)
            .map(|parent| {
                self.create_single_move_commit(
                    ctx,
                    target_repo.inner_repo(),
                    parent.clone(),
                    &mover,
                    &directory_mover,
                    Default::default(),
                    source_name,
                )
            })
            .buffer_unordered(MERGE_COMMIT_MOVES_CONCURRENCY)
            .try_collect::<Vec<_>>()
            .await?;

        save_bonsai_changesets(
            moved_commits.iter().map(|css| css.moved.clone()).collect(),
            ctx.clone(),
            target_repo.inner_repo(),
        )
        .await?;

        let mutable_renames_count: usize = moved_commits
            .iter()
            .map(|css| css.mutable_renames.len())
            .sum();
        let mut scuba = ctx.scuba().clone();
        scuba.add("mutable_renames_count", mutable_renames_count);
        scuba.log_with_msg("Started saving mutable renames", None);
        self.save_mutable_renames(
            ctx,
            target_repo.inner_repo().changesets(),
            self.mutable_renames,
            moved_commits.iter().map(|css| &css.mutable_renames),
        )
        .await?;
        scuba.log_with_msg("Saved mutable renames", None);

        Ok(moved_commits)
    }

    // If that sync_changeset() call was successful, but failed to send
    // successful result to the client (e.g. network issues) then
    // client will retry a request. We need to detect this situation and
    // send a successful response to the client.
    async fn check_if_this_method_has_already_succeeded(
        &self,
        ctx: &CoreContext,
        source_cs_id: ChangesetId,
        source_name: &SourceName,
        (expected_target_location, actual_target_location): (ChangesetId, ChangesetId),
        repo: &RepoContext,
    ) -> Result<ChangesetId, MegarepoError> {
        // Bookmark points a non-expected commit - let's see if changeset it points to was created
        // by a previous sync_changeset call

        // Check that first parent is a target location
        let parents = repo
            .inner_repo()
            .changeset_fetcher()
            .get_parents(ctx, actual_target_location)
            .await?;
        if parents.get(0) != Some(&expected_target_location) {
            return Err(MegarepoError::request(anyhow!(
                "Neither {} nor its first parent {:?} point to a target location {}",
                actual_target_location,
                parents.get(0),
                expected_target_location,
            )));
        }

        let state = self
            .read_remapping_state_file(ctx, repo, actual_target_location)
            .await?;

        let latest_synced = state.latest_synced_changesets.get(source_name);
        if Some(&source_cs_id) != latest_synced {
            return Err(MegarepoError::request(anyhow!(
                "In target commit {} latest synced source commit is {:?}, but expected {}",
                actual_target_location,
                latest_synced,
                source_cs_id,
            )));
        }

        Ok(actual_target_location)
    }
}

// We allow syncing changeset from a source if one of its parents was the latest synced changeset
// from this source into this target.
async fn validate_can_sync_changeset(
    _ctx: &CoreContext,
    target: &Target,
    source_cs: &BonsaiChangeset,
    commit_remapping_state: &CommitRemappingState,
    _source_repo: &RepoContext,
    source: &Source,
) -> Result<(), MegarepoError> {
    match &source.revision {
        SourceRevision::hash(_) => {
            /* If the revision is hardcoded hash it should be changed using remerge_source */
            return Err(MegarepoError::request(anyhow!(
                "can't sync changeset from source {} because this source points to a changeset",
                source.source_name,
            )));
        }
        SourceRevision::bookmark(_bookmark) => {
            /* If the source is following a git repo branch we can't verify much as the bookmark
            doesn't have to exist in the megarepo */
        }
        SourceRevision::UnknownField(_) => {
            return Err(MegarepoError::internal(anyhow!(
                "unexpected source revision!"
            )));
        }
    };

    let latest_synced_cs_id = find_latest_synced_cs_id(
        commit_remapping_state,
        &SourceName::new(&source.source_name),
        target,
    )?;

    let found = source_cs.parents().find(|p| *p == latest_synced_cs_id);
    if found.is_none() {
        return Err(MegarepoError::request(anyhow!(
            "Can't sync {}, because latest synced commit is not a parent of this commit. \
                    Latest synced source changeset is {}",
            source_cs.get_changeset_id(),
            latest_synced_cs_id,
        )));
    }
    Ok(())
}

async fn sync_changeset_to_target(
    ctx: &CoreContext,
    mapping: &SourceMappingRules,
    source: &SourceName,
    source_repo: &impl Repo,
    source_cs: BonsaiChangeset,
    target_repo: &impl Repo,
    target_cs_id: ChangesetId,
    target: &Target,
    mut state: CommitRemappingState,
    merge_mode: MergeMode,
) -> Result<ChangesetId, MegarepoError> {
    let mover =
        create_source_to_target_multi_mover(mapping.clone()).map_err(MegarepoError::internal)?;

    let source_cs_id = source_cs.get_changeset_id();
    // Create a new commit using a mover
    let source_cs_mut = source_cs.into_mut();
    let latest_synced_cs_id = find_latest_synced_cs_id(&state, source, target)?;

    let mut rewritten_commit = match merge_mode {
        MergeMode::ExtraMoveCommits {
            side_parents_move_commits,
        } => {
            let mut remapped_parents = HashMap::new();

            remapped_parents.insert(latest_synced_cs_id, target_cs_id);
            for css in side_parents_move_commits.iter() {
                remapped_parents.insert(css.source, css.moved.get_changeset_id());
            }
            rewrite_commit(
                ctx, // this is already a reference
                source_cs_mut,
                &remapped_parents,
                mover,
                source_repo,
                // In case of octopus merges only first two parent get preserved during
                // hg derivation. This ensures that mainline is within those two so is
                // represented in the commit graph and the sync is a fast-forward move.
                Some(target_cs_id),
                Default::default(),
            )
            .await
            .map_err(MegarepoError::internal)?
            .ok_or_else(|| {
                MegarepoError::internal(anyhow!(
                    "failed to rewrite commit {}, target: {:?}",
                    source_cs_id,
                    target
                ))
            })
        }
        MergeMode::Squashed { side_commits } => {
            let side_commits_info = stream::iter(side_commits.into_iter())
                .map(|cs_ctx| async move {
                    let hash = match cs_ctx.git_sha1().await? {
                        None => format!("HG hash: {}", cs_ctx.id()),
                        Some(hash) => format!("Git hash: {}", hash),
                    };
                    let author_date = cs_ctx.author_date().await?;
                    let message = cs_ctx.message().await?;
                    let title = message.trim_start().lines().next().unwrap_or("");
                    Ok::<_, MegarepoError>(format!(" * {}\t{}\t{}", hash, author_date, title))
                })
                .buffer_unordered(100)
                .try_collect()
                .await?;
            rewrite_as_squashed_commit(
                ctx,
                source_repo,
                source_cs_id,
                (latest_synced_cs_id, target_cs_id),
                source_cs_mut,
                mover,
                side_commits_info,
            )
            .await
            .map_err(MegarepoError::internal)?
            .ok_or_else(|| {
                MegarepoError::internal(anyhow!(
                    "failed to rewrite as squashed commit {}, target: {:?}",
                    source_cs_id,
                    target
                ))
            })
        }
    }?;

    state.set_source_changeset(source.clone(), source_cs_id);
    state
        .save_in_changeset(ctx, target_repo, &mut rewritten_commit)
        .await?;

    let rewritten_commit = rewritten_commit.freeze().map_err(MegarepoError::internal)?;
    let target_cs_id = rewritten_commit.get_changeset_id();
    upload_commits(ctx, vec![rewritten_commit], source_repo, target_repo)
        .await
        .map_err(MegarepoError::internal)?;

    Ok(target_cs_id)
}

fn find_latest_synced_cs_id(
    commit_remapping_state: &CommitRemappingState,
    source_name: &SourceName,
    target: &Target,
) -> Result<ChangesetId, MegarepoError> {
    let maybe_latest_synced_cs_id = commit_remapping_state.get_latest_synced_changeset(source_name);
    if let Some(latest_synced_cs_id) = maybe_latest_synced_cs_id {
        Ok(latest_synced_cs_id.clone())
    } else {
        Err(MegarepoError::internal(anyhow!(
            "Source {:?} was not synced into target {:?}",
            source_name,
            target
        )))
    }
}

#[cfg(test)]
mod test {
    use anyhow::Error;
    use fbinit::FacebookInit;
    use maplit::hashmap;
    use megarepo_mapping::REMAPPING_STATE_FILE;
    use mononoke_types::FileChange;
    use mononoke_types::MPath;
    use tests_utils::bookmark;
    use tests_utils::list_working_copy_utf8;
    use tests_utils::resolve_cs_id;
    use tests_utils::CreateCommitContext;

    use super::*;
    use crate::megarepo_test_utils::MegarepoTest;
    use crate::megarepo_test_utils::SyncTargetConfigBuilder;

    #[fbinit::test]
    async fn test_sync_changeset_simple(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let mut test = MegarepoTest::new(&ctx).await?;
        let target: Target = test.target("target".to_string());

        let source_name = SourceName::new("source_1");
        let version = "version_1".to_string();
        SyncTargetConfigBuilder::new(test.repo_id(), target.clone(), version.clone())
            .source_builder(source_name.clone())
            .set_prefix_bookmark_to_source_name()
            .build_source()?
            .build(&mut test.configs_storage);

        println!("Create initial source commit and bookmark");
        let init_source_cs_id = CreateCommitContext::new_root(&ctx, &test.blobrepo)
            .add_file("file", "content")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(init_source_cs_id)
            .await?;

        let latest_target_cs_id = test
            .prepare_initial_commit_in_target(&ctx, &version, &target)
            .await?;

        let configs_storage: Arc<dyn MononokeMegarepoConfigs> = Arc::new(test.configs_storage);
        let sync_changeset = SyncChangeset::new(
            &configs_storage,
            &test.mononoke,
            &test.megarepo_mapping,
            &test.mutable_renames,
        );
        println!("Trying to sync already synced commit again");
        let res = sync_changeset
            .sync(
                &ctx,
                init_source_cs_id,
                &source_name,
                &target,
                latest_target_cs_id,
            )
            .await;
        assert!(res.is_err());

        let source_cs_id = CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source_cs_id])
            .add_file("anotherfile", "anothercontent")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(source_cs_id)
            .await?;

        println!("Syncing new commit");
        sync_changeset
            .sync(
                &ctx,
                source_cs_id,
                &source_name,
                &target,
                latest_target_cs_id,
            )
            .await?;

        let cs_id = resolve_cs_id(&ctx, &test.blobrepo, "target").await?;
        let mut wc = list_working_copy_utf8(&ctx, &test.blobrepo, cs_id).await?;

        // Remove file with commit remapping state because it's never present in source
        wc.remove(&MPath::new(REMAPPING_STATE_FILE)?);

        assert_eq!(
            wc,
            hashmap! {
                MPath::new("source_1/file")? => "content".to_string(),
                MPath::new("source_1/anotherfile")? => "anothercontent".to_string(),
            }
        );

        Ok(())
    }

    #[fbinit::test]
    async fn test_sync_changeset_octopus_merge(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let mut test = MegarepoTest::new(&ctx).await?;
        let target: Target = test.target("target".to_string());

        let source_name = SourceName::new("source_1");
        let version = "version_1".to_string();
        SyncTargetConfigBuilder::new(test.repo_id(), target.clone(), version.clone())
            .source_builder(source_name.clone())
            .set_prefix_bookmark_to_source_name()
            .copyfile("file", "copyfile")
            .build_source()?
            .build(&mut test.configs_storage);

        println!("Create initial source commit and bookmark");
        let init_source_cs_id = CreateCommitContext::new_root(&ctx, &test.blobrepo)
            .add_file("file", "content")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(init_source_cs_id)
            .await?;

        let latest_target_cs_id = test
            .prepare_initial_commit_in_target(&ctx, &version, &target)
            .await?;

        let configs_storage: Arc<dyn MononokeMegarepoConfigs> = Arc::new(test.configs_storage);
        let sync_changeset = SyncChangeset::new(
            &configs_storage,
            &test.mononoke,
            &test.megarepo_mapping,
            &test.mutable_renames,
        );

        let merge_parent_1_source =
            CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source_cs_id])
                .add_file("file", "anothercontent")
                .add_file("file_from_parent_1", "parent_1")
                .commit()
                .await?;

        let merge_parent_2_source =
            CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source_cs_id])
                .add_file("file", "totallydifferentcontent")
                .add_file("file_from_parent_2", "parent_2")
                .commit()
                .await?;

        let merge_parent_3_source =
            CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source_cs_id])
                .add_file("file", "contentfromthirdparent")
                .add_file("file_from_parent_3", "parent_3")
                .commit()
                .await?;

        let merge_source = CreateCommitContext::new(
            &ctx,
            &test.blobrepo,
            vec![
                merge_parent_2_source,
                merge_parent_3_source,
                // Commit parent comming from the target last to ensure that
                // parent reordering works as expected.
                merge_parent_1_source,
            ],
        )
        .add_file("file", "mergeresolution")
        .add_file_with_copy_info(
            "copy_of_file",
            "totallydifferentcontent",
            (merge_parent_2_source, "file"),
        )
        .commit()
        .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(merge_parent_1_source)
            .await?;
        println!("Syncing first merge parent");
        let merge_parent_1_target = sync_changeset
            .sync(
                &ctx,
                merge_parent_1_source,
                &source_name,
                &target,
                latest_target_cs_id,
            )
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(merge_source)
            .await?;
        println!("Syncing merge commit parent");
        let merge_target = sync_changeset
            .sync(
                &ctx,
                merge_source,
                &source_name,
                &target,
                merge_parent_1_target,
            )
            .await?;

        let mut wc = list_working_copy_utf8(&ctx, &test.blobrepo, merge_target).await?;

        // Remove file with commit remapping state because it's never present in source
        wc.remove(&MPath::new(REMAPPING_STATE_FILE)?);

        assert_eq!(
            wc,
            hashmap! {
                MPath::new("source_1/file")? => "mergeresolution".to_string(),
                MPath::new("source_1/file_from_parent_1")? => "parent_1".to_string(),
                MPath::new("source_1/file_from_parent_2")? => "parent_2".to_string(),
                MPath::new("source_1/file_from_parent_3")? => "parent_3".to_string(),
                MPath::new("source_1/copy_of_file")? => "totallydifferentcontent".to_string(),
                MPath::new("copyfile")? => "mergeresolution".to_string(),
            }
        );

        let merge_target_cs = merge_target
            .load(&ctx, &test.blobrepo.repo_blobstore())
            .await?;

        let copied_file_change_from_bonsai = match merge_target_cs
            .file_changes()
            .find(|(p, _)| p == &&MPath::new("source_1/copy_of_file").unwrap())
            .unwrap()
            .1
        {
            FileChange::Change(tc) => tc,
            _ => panic!(),
        };
        assert_eq!(
            copied_file_change_from_bonsai.copy_from().unwrap().0,
            MPath::new("source_1/file")?
        );

        // All parents are preserved.
        assert_eq!(merge_target_cs.parents().count(), 3);

        // The parent from target comes first.
        assert_eq!(
            merge_target_cs.parents().next().unwrap(),
            merge_parent_1_target,
        );
        Ok(())
    }

    #[fbinit::test]
    async fn test_sync_changeset_two_sources_one_with_diamond_merge(
        fb: FacebookInit,
    ) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let mut test = MegarepoTest::new(&ctx).await?;
        let target: Target = test.target("target".to_string());

        let source1_name = SourceName::new("source_1");
        let source2_name = SourceName::new("source_2");
        let version = "version_1".to_string();
        SyncTargetConfigBuilder::new(test.repo_id(), target.clone(), version.clone())
            .source_builder(source1_name.clone())
            .set_prefix_bookmark_to_source_name()
            .build_source()?
            .source_builder(source2_name.clone())
            .set_prefix_bookmark_to_source_name()
            .build_source()?
            .build(&mut test.configs_storage);

        println!("Create initial first source commit and bookmark");
        let init_source1_cs_id = CreateCommitContext::new_root(&ctx, &test.blobrepo)
            .add_file("file1", "content1")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source1_name.to_string())
            .set_to(init_source1_cs_id)
            .await?;

        println!("Create initial second source commit and bookmark");
        let init_source2_cs_id = CreateCommitContext::new_root(&ctx, &test.blobrepo)
            .add_file("file2", "content2")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source2_name.to_string())
            .set_to(init_source2_cs_id)
            .await?;

        let mut latest_target_cs_id = test
            .prepare_initial_commit_in_target(&ctx, &version, &target)
            .await?;

        print!("Syncing one commit to each of sources... 1");
        let configs_storage: Arc<dyn MononokeMegarepoConfigs> = Arc::new(test.configs_storage);
        let sync_changeset = SyncChangeset::new(
            &configs_storage,
            &test.mononoke,
            &test.megarepo_mapping,
            &test.mutable_renames,
        );
        let source1_cs_id =
            CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source1_cs_id])
                .add_file("anotherfile1", "anothercontent")
                .commit()
                .await?;
        bookmark(&ctx, &test.blobrepo, source1_name.to_string())
            .set_to(source1_cs_id)
            .await?;
        latest_target_cs_id = sync_changeset
            .sync(
                &ctx,
                source1_cs_id,
                &source1_name,
                &target,
                latest_target_cs_id,
            )
            .await?;
        println!(", 2");

        let source2_cs_id =
            CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source2_cs_id])
                .add_file("anotherfile2", "anothercontent")
                .commit()
                .await?;
        bookmark(&ctx, &test.blobrepo, source2_name.to_string())
            .set_to(source2_cs_id)
            .await?;
        latest_target_cs_id = sync_changeset
            .sync(
                &ctx,
                source2_cs_id,
                &source2_name,
                &target,
                latest_target_cs_id,
            )
            .await?;

        println!("Trying to sync already synced commit again");
        let res = sync_changeset
            .sync(
                &ctx,
                source1_cs_id,
                &source1_name,
                &target,
                latest_target_cs_id,
            )
            .await;
        assert!(res.is_err());
        println!("Trying to sync a diamond merge commit");

        let source1_diamond_merge_cs_id = CreateCommitContext::new(
            &ctx,
            &test.blobrepo,
            vec![source1_cs_id, init_source1_cs_id],
        )
        .add_file("anotherfile1", "content_from_diamond_merge")
        .commit()
        .await?;
        bookmark(&ctx, &test.blobrepo, source1_name.to_string())
            .set_to(source1_diamond_merge_cs_id)
            .await?;
        let _diamond_merge_synced = sync_changeset
            .sync(
                &ctx,
                source1_diamond_merge_cs_id,
                &source1_name,
                &target,
                latest_target_cs_id,
            )
            .await?;

        let target_cs_id = resolve_cs_id(&ctx, &test.blobrepo, "target").await?;
        let mut wc = list_working_copy_utf8(&ctx, &test.blobrepo, target_cs_id).await?;

        // Remove file with commit remapping state because it's never present in source
        wc.remove(&MPath::new(REMAPPING_STATE_FILE)?);

        assert_eq!(
            wc,
            hashmap! {
                MPath::new("source_1/file1")? => "content1".to_string(),
                MPath::new("source_1/anotherfile1")? => "content_from_diamond_merge".to_string(),
                MPath::new("source_2/file2")? => "content2".to_string(),
                MPath::new("source_2/anotherfile2")? => "anothercontent".to_string(),
            }
        );

        let target_cs = target_cs_id
            .load(&ctx, &test.blobrepo.repo_blobstore())
            .await?;
        // All parents are preserved.
        assert_eq!(target_cs.parents().count(), 2);

        Ok(())
    }

    #[fbinit::test]
    async fn test_sync_changeset_repeat_same_request(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let mut test = MegarepoTest::new(&ctx).await?;
        let target: Target = test.target("target".to_string());

        let source_name = SourceName::new("source_1");
        let version = "version_1".to_string();
        SyncTargetConfigBuilder::new(test.repo_id(), target.clone(), version.clone())
            .source_builder(source_name.clone())
            .set_prefix_bookmark_to_source_name()
            .build_source()?
            .build(&mut test.configs_storage);

        println!("Create initial source commit and bookmark");
        let init_source_cs_id = CreateCommitContext::new_root(&ctx, &test.blobrepo)
            .add_file("file", "content")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(init_source_cs_id)
            .await?;

        let latest_target_cs_id = test
            .prepare_initial_commit_in_target(&ctx, &version, &target)
            .await?;

        let configs_storage: Arc<dyn MononokeMegarepoConfigs> = Arc::new(test.configs_storage);
        let sync_changeset = SyncChangeset::new(
            &configs_storage,
            &test.mononoke,
            &test.megarepo_mapping,
            &test.mutable_renames,
        );

        let source_cs_id = CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source_cs_id])
            .add_file("anotherfile", "anothercontent")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(source_cs_id)
            .await?;

        println!("Syncing new commit");
        let res1 = sync_changeset
            .sync(
                &ctx,
                source_cs_id,
                &source_name,
                &target,
                latest_target_cs_id,
            )
            .await?;

        println!("Now syncing the same commit again - should succeed");
        let res2 = sync_changeset
            .sync(
                &ctx,
                source_cs_id,
                &source_name,
                &target,
                latest_target_cs_id,
            )
            .await?;

        assert_eq!(res1, res2);

        Ok(())
    }

    #[fbinit::test]
    async fn test_sync_changeset_squash_commit(fb: FacebookInit) -> Result<(), Error> {
        let ctx = CoreContext::test_mock(fb);
        let mut test = MegarepoTest::new(&ctx).await?;
        let target: Target = test.target("target".to_string());

        let source_name = SourceName::new("source_1");
        let version = "version_1".to_string();
        SyncTargetConfigBuilder::new(test.repo_id(), target.clone(), version.clone())
            .source_builder(source_name.clone())
            .set_prefix_bookmark_to_source_name()
            .merge_mode(megarepo_config::MergeMode::squashed(
                megarepo_config::Squashed { squash_limit: 3 },
            ))
            .build_source()?
            .build(&mut test.configs_storage);

        println!("Create initial source commit and bookmark");
        let init_source_cs_id = CreateCommitContext::new_root(&ctx, &test.blobrepo)
            .add_file("file", "content")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(init_source_cs_id)
            .await?;

        let latest_target_cs_id = test
            .prepare_initial_commit_in_target(&ctx, &version, &target)
            .await?;

        let configs_storage: Arc<dyn MononokeMegarepoConfigs> = Arc::new(test.configs_storage);
        let sync_changeset = SyncChangeset::new(
            &configs_storage,
            &test.mononoke,
            &test.megarepo_mapping,
            &test.mutable_renames,
        );

        let main_line = CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source_cs_id])
            .add_file("file_in_mainline", "mainline1")
            .commit()
            .await?;

        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(main_line)
            .await?;
        let main_line_target = sync_changeset
            .sync(&ctx, main_line, &source_name, &target, latest_target_cs_id)
            .await?;

        let side_branch_1 = CreateCommitContext::new(&ctx, &test.blobrepo, vec![init_source_cs_id])
            .add_file("file", "totallydifferentcontent")
            .add_file("file_in_sidebranch_1", "sidebranch1")
            .commit()
            .await?;

        let side_branch_2 = CreateCommitContext::new(&ctx, &test.blobrepo, vec![side_branch_1])
            .add_file("file", "amended")
            .add_file("file_in_sidebranch_2", "sidebranch2")
            .commit()
            .await?;

        let merge = CreateCommitContext::new(&ctx, &test.blobrepo, vec![side_branch_2, main_line])
            .add_file("file", "mergeresolution")
            .commit()
            .await?;
        println!("Syncing merge");
        bookmark(&ctx, &test.blobrepo, source_name.to_string())
            .set_to(merge)
            .await?;
        let merge_target = sync_changeset
            .sync(&ctx, merge, &source_name, &target, main_line_target)
            .await?;

        let _mcs = merge.load(&ctx, test.blobrepo.repo_blobstore()).await?;

        // Find source repo and changeset that we need to sync
        let target_repo = sync_changeset.find_repo_by_id(&ctx, target.repo_id).await?;
        let merge_cs = merge_target
            .load(&ctx, target_repo.inner_repo().repo_blobstore())
            .await?;

        let parents: Vec<_> = merge_cs.parents().collect();

        let mut wc = list_working_copy_utf8(&ctx, &test.blobrepo, merge_target).await?;

        // Remove file with commit remapping state because it's never present in source
        wc.remove(&MPath::new(REMAPPING_STATE_FILE)?);

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

        assert_eq!(
            wc,
            hashmap! {
                MPath::new("source_1/file")? => "mergeresolution".to_string(),
                MPath::new("source_1/file_in_sidebranch_1")? => "sidebranch1".to_string(),
                MPath::new("source_1/file_in_sidebranch_2")? => "sidebranch2".to_string(),
                MPath::new("source_1/file_in_mainline")? => "mainline1".to_string(),
            }
        );
        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.
 */

//! Overall coordinator for parsing bundle2 streams.

use std::fmt;
use std::fmt::Debug;
use std::fmt::Display;
use std::fmt::Formatter;
use std::io::BufRead;
use std::io::Chain;
use std::io::Cursor;
use std::io::Read;
use std::mem;

use anyhow::Error;
use bytes_old::BytesMut;
use futures_ext::io::Either;
use futures_ext::BoxFuture;
use futures_old::Async;
use futures_old::Poll;
use futures_old::Stream;
use slog::Logger;
use tokio_codec::Framed;
use tokio_codec::FramedParts;
use tokio_io::AsyncRead;

use crate::errors::ErrorKind;
use crate::part_inner::inner_stream;
use crate::part_outer::outer_stream;
use crate::part_outer::OuterFrame;
use crate::part_outer::OuterStream;
use crate::stream_start::StartDecoder;
use crate::Bundle2Item;
use crate::OldBundle2Item;

pub enum StreamEvent<I, S> {
    Next(I),
    Done(S),
}

impl<I, S> StreamEvent<I, S> {
    pub fn into_next(self) -> Result<I, Self> {
        if let StreamEvent::Next(v) = self {
            Ok(v)
        } else {
            Err(self)
        }
    }
}

impl<I, S> Debug for StreamEvent<I, S> {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match *self {
            StreamEvent::Next(_) => write!(f, "Next(...)"),
            StreamEvent::Done(_) => write!(f, "Done(...)"),
        }
    }
}

pub type Remainder<R> = (BytesMut, R);

#[derive(Debug)]
pub struct Bundle2Stream<R>
where
    R: AsyncRead + BufRead + 'static + Send,
{
    inner: Bundle2StreamInner,
    current_stream: CurrentStream<R>,
}

#[derive(Debug)]
struct Bundle2StreamInner {
    logger: Logger,
    app_errors: Vec<ErrorKind>,
}

enum CurrentStream<R>
where
    R: AsyncRead + BufRead + 'static + Send,
{
    Start(Framed<R, StartDecoder>),
    Outer(OuterStream<Chain<Cursor<BytesMut>, R>>),
    Inner(BoxFuture<OuterStream<Chain<Cursor<BytesMut>, R>>, Error>),
    Invalid,
    End,
}

impl<R> CurrentStream<R>
where
    R: AsyncRead + BufRead + 'static + Send,
{
    pub fn take(&mut self) -> Self {
        mem::replace(self, CurrentStream::Invalid)
    }
}

impl<R> Display for CurrentStream<R>
where
    R: AsyncRead + BufRead + 'static + Send,
{
    fn fmt(&self, fmt: &mut fmt::Formatter<'_>) -> fmt::Result {
        use self::CurrentStream::*;

        let s = match *self {
            Start(_) => "start",
            Outer(_) => "outer",
            Inner(_) => "inner",
            Invalid => "invalid",
            End => "end",
        };
        write!(fmt, "{}", s)
    }
}

impl<R> Debug for CurrentStream<R>
where
    R: AsyncRead + BufRead + Debug + 'static + Send,
{
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        match *self {
            CurrentStream::Start(ref framed_stream) => write!(f, "Start({:?})", framed_stream),
            CurrentStream::Outer(ref outer_stream) => write!(f, "Outer({:?})", outer_stream),
            // InnerStream currently doesn't implement Debug because
            // part_inner::BoolFuture doesn't implement Debug.
            CurrentStream::Inner(_) => write!(f, "Inner(inner_stream)"),
            CurrentStream::Invalid => write!(f, "Invalid"),
            CurrentStream::End => write!(f, "End"),
        }
    }
}

impl<R> Bundle2Stream<R>
where
    R: AsyncRead + BufRead + 'static + Send,
{
    pub fn new(logger: Logger, read: R) -> Bundle2Stream<R> {
        Bundle2Stream {
            inner: Bundle2StreamInner {
                logger,
                app_errors: Vec::new(),
            },
            current_stream: CurrentStream::Start(Framed::from_parts(FramedParts::new(
                read,
                StartDecoder,
            ))),
        }
    }

    pub fn app_errors(&self) -> &[ErrorKind] {
        &self.inner.app_errors
    }
}

impl<R> Stream for Bundle2Stream<R>
where
    R: AsyncRead + BufRead + 'static + Send,
{
    type Item = StreamEvent<Bundle2Item<'static>, Remainder<R>>;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Self::Error> {
        let current_stream = self.current_stream.take();

        let (ret, stream) = self.inner.poll_next(current_stream);
        self.current_stream = stream;
        match ret {
            Ok(Async::Ready(Some(StreamEvent::Next(item)))) => {
                Ok(Async::Ready(Some(StreamEvent::Next(item.into()))))
            }
            Ok(Async::Ready(Some(StreamEvent::Done(rem)))) => {
                Ok(Async::Ready(Some(StreamEvent::Done(rem))))
            }
            Ok(Async::Ready(None)) => Ok(Async::Ready(None)),
            Ok(Async::NotReady) => Ok(Async::NotReady),
            Err(err) => Err(err),
        }
    }
}

impl Bundle2StreamInner {
    fn poll_next<R>(
        &mut self,
        current_stream: CurrentStream<R>,
    ) -> (
        Poll<Option<StreamEvent<OldBundle2Item, Remainder<R>>>, Error>,
        CurrentStream<R>,
    )
    where
        R: AsyncRead + BufRead + 'static + Send,
    {
        match current_stream {
            CurrentStream::Start(mut stream) => {
                match stream.poll() {
                    Err(e) => (Err(e), CurrentStream::Start(stream)),
                    Ok(Async::NotReady) => (Ok(Async::NotReady), CurrentStream::Start(stream)),
                    Ok(Async::Ready(None)) => {
                        (Ok(Async::Ready(None)), CurrentStream::Start(stream))
                    }
                    Ok(Async::Ready(Some(start))) => {
                        let FramedParts {
                            io,
                            read_buf,
                            write_buf,
                            ..
                        } = stream.into_parts();
                        assert!(
                            write_buf.is_empty(),
                            "write_buf must be empty, since io is not AsyncWrite"
                        );

                        match outer_stream(
                            self.logger.clone(),
                            &start,
                            Cursor::new(read_buf).chain(io),
                        ) {
                            Err(e) => {
                                // Can't do much if reading stream level params
                                // failed -- go to the invalid state.
                                (Err(e), CurrentStream::Invalid)
                            }
                            Ok(v) => {
                                let outer = CurrentStream::Outer(v);
                                (
                                    Ok(Async::Ready(Some(StreamEvent::Next(
                                        OldBundle2Item::Start(start),
                                    )))),
                                    outer,
                                )
                            }
                        }
                    }
                }
            }
            CurrentStream::Outer(mut stream) => {
                match stream.poll() {
                    Err(e) => match e.downcast::<ErrorKind>() {
                        Ok(ek) => {
                            if ek.is_app_error() {
                                self.app_errors.push(ek);
                                self.poll_next(CurrentStream::Outer(stream))
                            } else {
                                (Err(Error::from(ek)), CurrentStream::Outer(stream))
                            }
                        }
                        Err(e) => (Err(e), CurrentStream::Outer(stream)),
                    },
                    Ok(Async::NotReady) => (Ok(Async::NotReady), CurrentStream::Outer(stream)),
                    Ok(Async::Ready(None)) => {
                        (Ok(Async::Ready(None)), CurrentStream::Outer(stream))
                    }
                    Ok(Async::Ready(Some(OuterFrame::Header(header)))) => {
                        let (bundle2item, remainder) =
                            inner_stream(self.logger.clone(), header, stream);
                        (
                            Ok(Async::Ready(Some(StreamEvent::Next(bundle2item)))),
                            CurrentStream::Inner(remainder),
                        )
                    }
                    Ok(Async::Ready(Some(OuterFrame::Discard))) => {
                        self.poll_next(CurrentStream::Outer(stream))
                    }
                    Ok(Async::Ready(Some(OuterFrame::StreamEnd))) => {
                        // No more parts to go.
                        let FramedParts {
                            io,
                            mut read_buf,
                            write_buf,
                            ..
                        } = stream.into_parts();
                        assert!(
                            write_buf.is_empty(),
                            "write_buf must be empty, since io is not AsyncWrite"
                        );

                        let io = match io {
                            Either::A(io) => io,
                            Either::B(decompressor) => decompressor.into_inner().into_inner(),
                        };

                        let (cursor, io) = io.into_inner();
                        read_buf
                            .extend_from_slice(&cursor.get_ref()[(cursor.position() as usize)..]);

                        (
                            Ok(Async::Ready(Some(StreamEvent::Done((read_buf, io))))),
                            CurrentStream::End,
                        )
                    }
                    _ => panic!("Expected a header or StreamEnd!"),
                }
            }
            CurrentStream::Inner(mut remainder) => match remainder.poll() {
                Err(e) => (Err(e), CurrentStream::Invalid),
                Ok(Async::NotReady) => (Ok(Async::NotReady), CurrentStream::Inner(remainder)),
                Ok(Async::Ready(remainder)) => self.poll_next(CurrentStream::Outer(remainder)),
            },
            CurrentStream::Invalid => (
                Err(ErrorKind::Bundle2Decode("corrupt byte stream".into()).into()),
                CurrentStream::Invalid,
            ),
            CurrentStream::End => (Ok(Async::Ready(None)), CurrentStream::End),
        }
    }
}


/*
 * 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::io;
use std::io::Cursor;
use std::mem;
use std::vec::IntoIter;

use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use async_compression::Compressor;
use async_compression::CompressorType;
use byteorder::BigEndian;
use byteorder::ByteOrder;
use bytes_old::Buf;
use bytes_old::BufMut;
use bytes_old::Bytes;
use bytes_old::IntoBuf;
use futures_ext::io::Either;
use futures_ext::io::Either::A as UncompressedRead;
use futures_ext::io::Either::B as CompressedRead;
use futures_old::stream::Forward;
use futures_old::try_ready;
use futures_old::Async;
use futures_old::AsyncSink;
use futures_old::Future;
use futures_old::Poll;
use futures_old::Sink;
use futures_old::StartSend;
use futures_old::Stream;
use mercurial_types::percent_encode;
use tokio_codec::FramedWrite;
use tokio_io::AsyncWrite;

use crate::chunk::Chunk;
use crate::chunk::ChunkEncoder;
use crate::errors::ErrorKind;
use crate::part_encode::PartEncode;
use crate::part_encode::PartEncodeBuilder;
use crate::part_header::PartId;
use crate::types::StreamHeader;
use crate::utils::capitalize_first;
use crate::utils::get_compression_param;
use crate::utils::is_mandatory_param;

/// This is a general wrapper around a Sink to prevent closing of the underlying Sink. This is
/// useful when using Sink::send_all, because in addition to writing and flushing the data it also
/// closes the sink, which may result in IO errors if called repeatedly on the same Sink
#[derive(Debug)]
struct NotClosingSink<S> {
    pub inner: S,
}
impl<S: Sink> Sink for NotClosingSink<S> {
    type SinkItem = S::SinkItem;
    type SinkError = S::SinkError;

    fn start_send(&mut self, item: Self::SinkItem) -> StartSend<Self::SinkItem, Self::SinkError> {
        self.inner.start_send(item)
    }

    fn poll_complete(&mut self) -> Poll<(), Self::SinkError> {
        self.inner.poll_complete()
    }

    fn close(&mut self) -> Poll<(), Self::SinkError> {
        Ok(Async::Ready(()))
    }
}

/// Builder to generate a bundle2.
pub struct Bundle2EncodeBuilder<W> {
    writer: W,
    header: StreamHeader,
    compressor_type: Option<CompressorType>,
    parts: Vec<PartEncode>,
}

impl<W> Bundle2EncodeBuilder<W>
where
    W: AsyncWrite + Send,
{
    pub fn new(writer: W) -> Self {
        Bundle2EncodeBuilder {
            writer,
            header: StreamHeader {
                m_stream_params: HashMap::new(),
                a_stream_params: HashMap::new(),
            },
            compressor_type: None,
            parts: Vec::new(),
        }
    }

    pub fn add_stream_param(&mut self, key: String, val: String) -> Result<&mut Self> {
        if &key.to_lowercase() == "compression" {
            let msg = "stream compression should be set through set_compressor_type";
            bail!(ErrorKind::Bundle2Encode(msg.into()));
        }
        if is_mandatory_param(&key)
            .with_context(|| ErrorKind::Bundle2Encode("stream key is invalid".into()))?
        {
            self.header.m_stream_params.insert(key.to_lowercase(), val);
        } else {
            self.header.a_stream_params.insert(key.to_lowercase(), val);
        }
        Ok(self)
    }

    pub fn add_part(&mut self, part: PartEncodeBuilder) -> &mut Self {
        let part_id = self.parts.len() as PartId;
        self.parts.push(part.build(part_id));
        self
    }

    pub fn set_compressor_type<C: Into<Option<CompressorType>>>(&mut self, ct: C) -> &mut Self {
        self.compressor_type = ct.into();
        self
    }

    pub fn build(self) -> Bundle2Encode<W> {
        let mut mparams = self.header.m_stream_params;

        // The compression type becomes a mandatory param.
        mparams.insert(
            "compression".into(),
            get_compression_param(&self.compressor_type).into(),
        );

        // Build the buffer required for the stream header.
        let mut header_buf: Vec<u8> = Vec::new();

        header_buf.put_slice(b"HG20");
        // Reserve 4 bytes for the length.
        header_buf.put_u32_be(0);
        // Now write out the stream header.

        let params = mparams
            .into_iter()
            .map(|x| (capitalize_first(x.0), x.1))
            .chain(self.header.a_stream_params.into_iter());
        Self::build_stream_params(params, &mut header_buf);

        let header_len = (header_buf.len() - 8) as u32;
        BigEndian::write_u32(&mut header_buf[4..], header_len);

        Bundle2Encode {
            state: EncodeState::Start(StartState {
                writer: self.writer,
                compressor_type: self.compressor_type,
                header_buf: Bytes::from(header_buf).into_buf(),
                parts: self.parts,
            }),
        }
    }

    fn build_stream_params<I>(params: I, header_buf: &mut Vec<u8>)
    where
        I: Iterator<Item = (String, String)>,
    {
        let mut first = true;
        for (key, val) in params {
            if !first {
                header_buf.put(b' ');
            }
            first = false;
            header_buf.put(percent_encode(&key).into_bytes().as_slice());
            header_buf.put(b'=');
            header_buf.put(percent_encode(&val).into_bytes().as_slice());
        }
    }
}

/// A sink that chunks generated by PartEncodes goes into.
type PartSink<W> = NotClosingSink<FramedWrite<Either<W, Compressor<W>>, ChunkEncoder>>;

/// A future to drive writing a part to a sink.
type PartFuture<W> = Forward<PartEncode, PartSink<W>>;

#[derive(Debug)]
enum EncodeState<W>
where
    W: AsyncWrite + 'static,
{
    Start(StartState<W>),
    Part(PartFuture<W>, IntoIter<PartEncode>),
    EndOfStream(PartSink<W>, bool),
    Finish(Compressor<W>),
    Done,
    Invalid,
}

impl<W: AsyncWrite> EncodeState<W> {
    fn take(&mut self) -> Self {
        mem::replace(self, EncodeState::Invalid)
    }
}

pub struct Bundle2Encode<W>
where
    W: AsyncWrite + 'static,
{
    state: EncodeState<W>,
}

#[derive(Debug)]
struct StartState<W>
where
    W: AsyncWrite + 'static,
{
    writer: W,
    compressor_type: Option<CompressorType>,
    header_buf: Cursor<Bytes>,
    parts: Vec<PartEncode>,
}

impl<W> Future for StartState<W>
where
    W: AsyncWrite + Send + 'static,
{
    type Item = ();
    type Error = Error;

    fn poll(&mut self) -> Poll<(), Error> {
        try_ready!(self.writer.write_buf(&mut self.header_buf));
        if self.header_buf.has_remaining() {
            Ok(Async::NotReady)
        } else {
            Ok(Async::Ready(()))
        }
    }
}

impl<W> StartState<W>
where
    W: AsyncWrite + Send,
{
    fn finish(self) -> (IntoIter<PartEncode>, PartSink<W>) {
        assert!(!self.header_buf.has_remaining());
        (
            self.parts.into_iter(),
            NotClosingSink {
                inner: FramedWrite::new(
                    match self.compressor_type {
                        None => UncompressedRead(self.writer),
                        Some(compressor_type) => {
                            CompressedRead(Compressor::new(self.writer, compressor_type))
                        }
                    },
                    ChunkEncoder,
                ),
            },
        )
    }
}

impl<W> Future for Bundle2Encode<W>
where
    W: AsyncWrite + Send,
{
    type Item = W;
    type Error = Error;
    fn poll(&mut self) -> Poll<W, Error> {
        let (ret, state) = Self::poll_next(self.state.take());
        self.state = state;
        ret
    }
}

impl<W> Bundle2Encode<W>
where
    W: AsyncWrite + Send,
{
    fn poll_next(state: EncodeState<W>) -> (Poll<W, Error>, EncodeState<W>) {
        match state {
            EncodeState::Start(mut start_state) => {
                match start_state.poll() {
                    Ok(Async::Ready(())) => {
                        // Writing to the stream header is done. Set up the sink
                        // and remaining parts.
                        let (iter, sink) = start_state.finish();
                        Self::poll_next_part(iter, sink)
                    }
                    Ok(Async::NotReady) => (Ok(Async::NotReady), EncodeState::Start(start_state)),
                    Err(err) => {
                        // Somehow writing out the stream header failed. Not
                        // much to do here unfortunately -- we must abort the future.
                        (Err(err), EncodeState::Invalid)
                    }
                }
            }
            EncodeState::Part(part_fut, iter) => Self::poll_part(part_fut, iter),
            EncodeState::EndOfStream(sink, eos_written) => Self::poll_eos(sink, eos_written),
            EncodeState::Finish(compressor) => Self::poll_finish(CompressedRead(compressor)),
            EncodeState::Done => panic!("polled Bundle2Encode future after it is complete"),
            EncodeState::Invalid => {
                panic!("polled Bundle2Encode future after it returned an error")
            }
        }
    }

    fn poll_part(
        mut part_fut: PartFuture<W>,
        iter: IntoIter<PartEncode>,
    ) -> (Poll<W, Error>, EncodeState<W>) {
        match part_fut.poll() {
            Ok(Async::Ready((_part_encoder, sink))) => {
                // This part is done.
                Self::poll_next_part(iter, sink)
            }
            Ok(Async::NotReady) => {
                // This part is still writing.
                (Ok(Async::NotReady), EncodeState::Part(part_fut, iter))
            }
            Err(err) => {
                // TODO: bail on (some forms of) errors?
                (Err(err), EncodeState::Part(part_fut, iter))
            }
        }
    }

    fn poll_next_part(
        mut iter: IntoIter<PartEncode>,
        sink: PartSink<W>,
    ) -> (Poll<W, Error>, EncodeState<W>) {
        match iter.next() {
            Some(part_enc) => Self::poll_part(part_enc.forward(sink), iter),
            None => Self::poll_eos(sink, false),
        }
    }

    fn poll_eos(mut sink: PartSink<W>, eos_written: bool) -> (Poll<W, Error>, EncodeState<W>) {
        if !eos_written {
            match sink.start_send(Chunk::empty()) {
                Ok(AsyncSink::Ready) => {}
                Ok(AsyncSink::NotReady(_)) => {
                    return (Ok(Async::NotReady), EncodeState::EndOfStream(sink, false));
                }
                Err(err) => return (Err(err), EncodeState::Invalid),
            };
        }

        match sink.poll_complete() {
            Ok(Async::Ready(())) => Self::poll_finish(sink.inner.into_inner()),
            Ok(Async::NotReady) => (Ok(Async::NotReady), EncodeState::EndOfStream(sink, true)),
            Err(err) => (Err(err), EncodeState::Invalid),
        }
    }

    fn poll_finish(compressor: Either<W, Compressor<W>>) -> (Poll<W, Error>, EncodeState<W>) {
        let compressor = match compressor {
            UncompressedRead(inner) => return (Ok(Async::Ready(inner)), EncodeState::Done),
            CompressedRead(inner) => inner,
        };

        match compressor.try_finish() {
            Ok(inner) => (Ok(Async::Ready(inner)), EncodeState::Done),
            Err((compressor, err)) => {
                if err.kind() == io::ErrorKind::WouldBlock {
                    (Ok(Async::NotReady), EncodeState::Finish(compressor))
                } else {
                    (
                        Err(Error::from(err).context(ErrorKind::Bundle2Encode(
                            "error while completing write".into(),
                        ))),
                        EncodeState::Invalid,
                    )
                }
            }
        }
    }
}

/// Ensure that Bundle2Encode is Send.
fn _assert_send() {
    fn _assert<T: Send>(_val: &T) {}

    let builder = Bundle2EncodeBuilder::new(Cursor::new(Vec::new()));
    _assert(&builder);
    _assert(&builder.build());
}


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

//! Unpacking capabilities

use std::borrow::Cow;
use std::collections::HashMap;

use anyhow::Error;
use anyhow::Result;
use bytes_old::Bytes;
use bytes_old::BytesMut;
use mercurial_types::utils::percent_encode;
use percent_encoding::percent_decode;
use tokio_io::codec::Decoder;

#[derive(Debug, PartialEq, Eq)]
pub struct Capabilities {
    caps: HashMap<String, Vec<String>>,
}

impl Capabilities {
    pub fn new(caps: HashMap<String, Vec<String>>) -> Self {
        Self { caps }
    }
}

/// This is a tokio_io Decoder for capabilities used f.e. in "replycaps" part of bundle2
///
/// The format is as follows:
/// <capabilities> := '' EOF | <line> ['\n' <line>]* EOF
/// <line> := <key> | <key> '=' <values>
/// <values> := '' | <value> [',' <value>]*
/// <key> := `url encoded key`
/// <value> := `url encoded value`
///
/// Notice that this decoder needs to reach EOF and it decodes everything as a single item
pub struct CapabilitiesUnpacker;

impl Decoder for CapabilitiesUnpacker {
    type Item = Capabilities;
    type Error = Error;

    fn decode(&mut self, _buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        Ok(None) // This unpacker unpacks single element, wait for EOF
    }

    fn decode_eof(&mut self, buf: &mut BytesMut) -> Result<Option<Self::Item>> {
        let mut caps = HashMap::new();
        for kv in buf.split(|b| b == &b'\n') {
            let mut kv = kv.splitn(2, |b| b == &b'=');
            let key = percent_decode(kv.next().expect("must have at least 1 element"))
                .decode_utf8()?
                .into_owned();
            let values = {
                match kv.next() {
                    None => Vec::new(),
                    Some(values) => {
                        let res: Result<Vec<_>, _> = values
                            .split(|b| b == &b',')
                            .filter(|v| !v.is_empty())
                            .map(|v| percent_decode(v).decode_utf8().map(Cow::into_owned))
                            .collect();
                        res?
                    }
                }
            };
            caps.insert(key, values);
        }

        buf.clear(); // all buf was consumed

        Ok(Some(Capabilities { caps }))
    }
}

pub fn encode_capabilities(caps: Capabilities) -> Bytes {
    let mut res = vec![];
    for (key, values) in caps.caps {
        let key = percent_encode(&key);
        let values = itertools::join(values.into_iter().map(|v| percent_encode(&v)), ",");
        res.push(format!("{}={}", key, values));
    }
    Bytes::from(itertools::join(res, "\n").as_str())
}

#[cfg(test)]
mod test {
    use maplit::hashmap;

    use super::*;

    #[test]
    fn caps_roundtrip() {
        let caps = hashmap! {
            "key1".to_string() => vec!["value11".to_string(), "value12".to_string()],
            "key_empty".to_string() => vec![],
            "key2".to_string() => vec!["value22".to_string()],
            "weirdkey,=,=".to_string() => vec!["weirdvalue,==,".to_string(), "value".to_string()],
        };

        let encoded = encode_capabilities(Capabilities::new(caps.clone()));
        let mut unpacker = CapabilitiesUnpacker;

        let decoded = unpacker
            .decode_eof(&mut BytesMut::from(encoded))
            .unwrap()
            .unwrap();
        assert_eq!(decoded.caps, caps);
    }
}


/*
 * 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::Error;
use anyhow::Result;
use bytes_old::BufMut;
use bytes_old::Bytes;
use bytes_old::BytesMut;
use tokio_codec::Decoder;
use tokio_codec::Encoder;

use crate::errors::ErrorKind;
use crate::utils::BytesExt;

/// A bundle2 chunk.
///
/// Chunks underlie the bundle2 protocol. A chunk is a series of bytes and is
/// encoded as:
///
/// i32 chunk_size
/// [u8] data (chunk_size bytes)
///
/// Normally chunk_size > 0.
///
/// There are two special kinds of chunks:
///
/// 1. An "empty chunk", which is simply a chunk of size 0. This is represented
///    as a Normal chunk below with an empty Bytes.
/// 2. An "error chunk", which is a chunk with size -1 and no data. Error chunks
///    interrupt a chunk stream and are followed by a new part.
#[derive(Clone, Debug, PartialEq)]
pub struct Chunk(ChunkInner);

// This is separate to prevent constructing chunks with unexpected Bytes objects.
#[derive(Clone, Debug, PartialEq)]
enum ChunkInner {
    Normal(Bytes),
    Error,
}

impl Chunk {
    pub fn new<T: Into<Bytes>>(val: T) -> Result<Self> {
        let bytes: Bytes = val.into();
        if bytes.len() > i32::max_value() as usize {
            bail!(ErrorKind::Bundle2Chunk(format!(
                "chunk of length {} exceeds maximum {}",
                bytes.len(),
                i32::max_value()
            )));
        }
        Ok(Chunk(ChunkInner::Normal(bytes)))
    }

    pub fn empty() -> Self {
        Chunk(ChunkInner::Normal(Bytes::new()))
    }

    pub fn error() -> Self {
        Chunk(ChunkInner::Error)
    }

    pub fn is_normal(&self) -> bool {
        match self.0 {
            ChunkInner::Normal(_) => true,
            ChunkInner::Error => false,
        }
    }

    pub fn is_empty(&self) -> bool {
        match self.0 {
            ChunkInner::Normal(ref bytes) => bytes.is_empty(),
            ChunkInner::Error => false,
        }
    }

    pub fn is_error(&self) -> bool {
        match self.0 {
            ChunkInner::Normal(_) => false,
            ChunkInner::Error => true,
        }
    }

    /// Convert a chunk into its inner bytes.
    ///
    /// Returns an error if this chunk was an error chunk, since those do not
    /// have any bytes associated with them.
    pub fn into_bytes(self) -> Result<Bytes> {
        match self.0 {
            ChunkInner::Normal(bytes) => Ok(bytes),
            ChunkInner::Error => bail!("error chunk, no associated bytes"),
        }
    }
}

/// Encode a bundle2 chunk into a bytestream.
#[derive(Debug)]
pub struct ChunkEncoder;

impl Encoder for ChunkEncoder {
    type Item = Chunk;
    type Error = Error;

    fn encode(&mut self, item: Chunk, dst: &mut BytesMut) -> Result<()> {
        match item.0 {
            ChunkInner::Normal(bytes) => {
                dst.reserve(4 + bytes.len());
                dst.put_i32_be(bytes.len() as i32);
                dst.put_slice(&bytes);
            }
            ChunkInner::Error => {
                dst.reserve(4);
                dst.put_i32_be(-1);
            }
        }
        Ok(())
    }
}

/// Decode a bytestream into bundle2 chunks.
#[allow(dead_code)]
pub struct ChunkDecoder;

impl Decoder for ChunkDecoder {
    type Item = Chunk;
    type Error = Error;

    fn decode(&mut self, src: &mut BytesMut) -> Result<Option<Chunk>> {
        if src.len() < 4 {
            return Ok(None);
        }

        let len = src.peek_i32();
        if len == -1 {
            src.drain_i32();
            return Ok(Some(Chunk::error()));
        }
        if len < 0 {
            bail!(ErrorKind::Bundle2Chunk(format!(
                "chunk length must be >= -1, found {}",
                len
            ),));
        }

        let len = len as usize;
        if src.len() < 4 + len {
            return Ok(None);
        }

        src.drain_i32();
        let chunk = Chunk::new(src.split_to(len))?;
        Ok(Some(chunk))
    }
}

#[cfg(test)]
mod test {
    use std::io::Cursor;

    use assert_matches::assert_matches;
    use futures::compat::Future01CompatExt;
    use futures_old::stream;
    use futures_old::Future;
    use futures_old::Sink;
    use futures_old::Stream;
    use quickcheck::quickcheck;
    use quickcheck::TestResult;
    use tokio_codec::FramedRead;
    use tokio_codec::FramedWrite;

    use super::*;

    #[test]
    fn test_empty_chunk() {
        let mut buf = BytesMut::with_capacity(4);
        buf.put_i32_be(0);

        let mut decoder = ChunkDecoder;
        let chunk = decoder.decode(&mut buf).unwrap().unwrap();

        assert_eq!(chunk, Chunk::empty());
        assert!(chunk.is_normal());
        assert!(chunk.is_empty());
        assert!(!chunk.is_error());
    }

    #[test]
    fn test_error_chunk() {
        let mut buf = BytesMut::with_capacity(4);
        buf.put_i32_be(-1);

        let mut decoder = ChunkDecoder;
        let chunk = decoder.decode(&mut buf).unwrap().unwrap();

        assert_eq!(chunk, Chunk::error());
        assert!(!chunk.is_normal());
        assert!(!chunk.is_empty());
        assert!(chunk.is_error());
    }

    #[test]
    fn test_invalid_chunk() {
        let mut buf = BytesMut::with_capacity(4);
        buf.put_i32_be(-2);

        let mut decoder = ChunkDecoder;
        let chunk_err = decoder.decode(&mut buf);

        let err = chunk_err.unwrap_err();
        assert_matches!(
            err.downcast::<ErrorKind>().unwrap(),
            ErrorKind::Bundle2Chunk(_)
        );
    }

    #[test]
    fn test_roundtrip() {
        // Avoid using the quickcheck! macro because it eats up line numbers in
        // stack traces.
        quickcheck(roundtrip as fn(Vec<Option<Vec<u8>>>) -> TestResult);
    }

    fn roundtrip(data: Vec<Option<Vec<u8>>>) -> TestResult {
        let count = data.len();
        // Treat Some(bytes) as a normal chunk, None as an error chunk.
        let chunks: Vec<Chunk> = data
            .into_iter()
            .map(|x| match x {
                Some(b) => Chunk::new(b).unwrap(),
                None => Chunk::error(),
            })
            .collect();
        // Make a clone so we can use chunks later.
        let chunks_res: Vec<Result<Chunk>> = chunks.clone().into_iter().map(Ok).collect();

        let cursor = Cursor::new(Vec::with_capacity(32 * 1024));
        let sink = FramedWrite::new(cursor, ChunkEncoder);

        let encode_fut = sink
            .send_all(stream::iter_ok(chunks_res).and_then(|x| x))
            .map_err(|err| panic!("{:#?}", err))
            .and_then(move |(sink, _)| {
                let mut cursor = sink.into_inner();
                cursor.set_position(0);

                // cursor will now have the encoded byte stream. Run it through the decoder.
                let stream = FramedRead::new(cursor, ChunkDecoder);

                let collector: Vec<Chunk> = Vec::with_capacity(count);
                collector.send_all(stream.map_err(|err| {
                    panic!("Unexpected error: {}", err);
                }))
            })
            .map(move |(collector, _)| {
                assert_eq!(collector, chunks);
            })
            .map_err(|err| panic!("{:#?}", err));

        let rt = tokio::runtime::Runtime::new().unwrap();
        rt.block_on(encode_fut.compat()).unwrap();

        TestResult::passed()
    }
}


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

//! Code to deal with deltas received or sent over the wire.

use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use bufsize::SizeCounter;
use bytes_old::BufMut;
use bytes_old::BytesMut;
use mercurial_types::delta::Delta;
use mercurial_types::delta::Fragment;

use crate::errors::ErrorKind;
use crate::utils::BytesExt;

const DELTA_HEADER_LEN: usize = 12;

/// Decodes this delta. Consumes the entire buffer, so accepts a BytesMut.
pub fn decode_delta(buf: BytesMut) -> Result<Delta> {
    let mut buf = buf;
    let mut frags = vec![];
    let mut remaining = buf.len();

    while remaining >= DELTA_HEADER_LEN {
        // Each delta fragment has:
        // ---
        // start offset: i32
        // end offset: i32
        // new length: i32
        // content (new length bytes)
        // ---
        let start = buf.drain_i32();
        let end = buf.drain_i32();
        let new_len = buf.drain_i32();
        // TODO: handle negative values for all the above

        let delta_len = (new_len as usize) + DELTA_HEADER_LEN;
        if remaining < delta_len {
            bail!(ErrorKind::InvalidDelta(format!(
                "expected {} bytes, {} remaining",
                delta_len, remaining
            )));
        }

        frags.push(Fragment {
            start: start as usize,
            end: end as usize,
            // TODO: avoid copies here by switching this to Bytes
            content: buf.split_to(new_len as usize).to_vec(),
        });

        remaining -= delta_len;
    }

    if remaining != 0 {
        bail!(ErrorKind::InvalidDelta(format!(
            "{} trailing bytes in encoded delta",
            remaining
        ),));
    }

    Delta::new(frags).with_context(|| ErrorKind::InvalidDelta("invalid fragment list".into()))
}

#[inline]
pub fn encoded_len(delta: &Delta) -> usize {
    let mut size_counter = SizeCounter::new();
    encode_delta(delta, &mut size_counter);
    size_counter.size()
}

pub fn encode_delta<B: BufMut>(delta: &Delta, out: &mut B) {
    for fragment in delta.fragments() {
        out.put_i32_be(fragment.start as i32);
        out.put_i32_be(fragment.end as i32);
        out.put_i32_be(fragment.content.len() as i32);
        out.put_slice(&fragment.content[..]);
    }
}

#[cfg(test)]
mod test {
    use assert_matches::assert_matches;
    use failure_ext::err_downcast;
    use failure_ext::err_downcast_ref;
    use quickcheck::quickcheck;

    use super::*;

    #[test]
    fn invalid_deltas() {
        let short_delta = BytesMut::from(&b"\0\0\0\0\0\0\0\0\0\0\0\x20"[..]);
        assert_matches!(
            err_downcast!(decode_delta(short_delta).unwrap_err(), err: ErrorKind => err),
            Ok(ErrorKind::InvalidDelta(ref msg))
            if msg == "expected 44 bytes, 12 remaining"
        );

        let short_header = BytesMut::from(&b"\0\0\0\0\0\0"[..]);
        assert_matches!(
            err_downcast!(decode_delta(short_header).unwrap_err(), err: ErrorKind => err),
            Ok(ErrorKind::InvalidDelta(ref msg))
            if msg == "6 trailing bytes in encoded delta"
        );

        // start = 2, end = 0
        let start_after_end = BytesMut::from(&b"\0\0\0\x02\0\0\0\0\0\0\0\0"[..]);
        match decode_delta(start_after_end) {
            Ok(bad) => panic!("unexpected success {:?}", bad),
            Err(err) => match err_downcast_ref!(err, err: ErrorKind => err) {
                Some(&ErrorKind::InvalidDelta(..)) => {}
                Some(bad) => panic!("Bad ErrorKind {:?}", bad),
                None => panic!("Unexpected error {:?}", err),
            },
        }
    }

    quickcheck! {
        fn roundtrip(delta: Delta) -> bool {
            let mut out = vec![];
            encode_delta(&delta, &mut out);
            assert_eq!(encoded_len(&delta), out.len());
            delta == decode_delta(out.into()).unwrap()
        }
    }
}


/*
 * 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 thiserror::Error;

use crate::part_header::PartHeader;
use crate::part_header::PartHeaderType;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("bundle2 decode error: {0}")]
    Bundle2Decode(String),
    #[error("changegroup decode error: {0}")]
    CgDecode(String),
    #[error("changegroup2 encode error: {0}")]
    Cg2Encode(String),
    #[error("wirepack decode error: {0}")]
    WirePackDecode(String),
    #[error("wirepack encode error: {0}")]
    WirePackEncode(String),
    #[error("bundle2 encode error: {0}")]
    Bundle2Encode(String),
    #[error("bundle2 chunk error: {0}")]
    Bundle2Chunk(String),
    #[error("invalid delta: {0}")]
    InvalidDelta(String),
    #[error("invalid wire pack entry: {0}")]
    InvalidWirePackEntry(String),
    #[error("unknown part type: {0:?}")]
    BundleUnknownPart(PartHeader),
    #[error("unknown params for bundle2 part '{0:?}': {1:?}")]
    BundleUnknownPartParams(PartHeaderType, Vec<String>),
    #[error("error while generating listkey part")]
    ListkeyGeneration,
    #[error("error while generating phase-heads part")]
    PhaseHeadsGeneration,
}

impl ErrorKind {
    pub fn is_app_error(&self) -> bool {
        match self {
            &ErrorKind::BundleUnknownPart(_) | &ErrorKind::BundleUnknownPartParams(..) => true,
            _ => false,
        }
    }
}


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

pub mod bundle2;
pub mod bundle2_encode;
pub mod capabilities;
pub mod changegroup;
mod chunk;
mod delta;
pub mod infinitepush;
pub mod obsmarkers;
pub mod part_encode;
mod part_header;
mod part_inner;
mod part_outer;
pub mod parts;
mod pushrebase;
mod quickcheck_types;
pub mod stream_start;
#[cfg(test)]
mod test;
mod types;
pub mod wirepack;

mod errors;
pub use crate::errors::ErrorKind;
mod utils;

use std::fmt;

use anyhow::bail;
use anyhow::Error;
use anyhow::Result;
use bytes_old::Bytes;
use futures::compat::Future01CompatExt;
use futures::compat::Stream01CompatExt;
use futures::future::BoxFuture;
use futures::future::FutureExt;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures_ext::BoxFuture as OldBoxFuture;
use futures_ext::BoxStream as OldBoxStream;
use futures_ext::SinkToAsyncWrite;
use futures_old::sync::mpsc;
use futures_old::sync::oneshot;
use futures_old::Future as OldFuture;
use futures_old::Stream as OldStream;

pub use crate::bundle2_encode::Bundle2EncodeBuilder;
pub use crate::part_header::PartHeader;
pub use crate::part_header::PartHeaderInner;
pub use crate::part_header::PartHeaderType;
pub use crate::part_header::PartId;
pub use crate::types::StreamHeader;

pub enum Bundle2Item<'a> {
    Start(StreamHeader),
    Changegroup(PartHeader, BoxStream<'a, Result<changegroup::Part>>),
    B2xCommonHeads(
        PartHeader,
        BoxStream<'a, Result<mercurial_types::HgChangesetId>>,
    ),
    B2xInfinitepush(PartHeader, BoxStream<'a, Result<changegroup::Part>>),
    B2xTreegroup2(PartHeader, BoxStream<'a, Result<wirepack::Part>>),
    // B2xInfinitepushBookmarks returns Bytes because this part is not going to be used.
    B2xInfinitepushBookmarks(PartHeader, BoxStream<'a, Result<bytes_old::Bytes>>),
    B2xInfinitepushMutation(
        PartHeader,
        BoxStream<'a, Result<Vec<mercurial_mutation::HgMutationEntry>>>,
    ),
    B2xRebasePack(PartHeader, BoxStream<'a, Result<wirepack::Part>>),
    B2xRebase(PartHeader, BoxStream<'a, Result<changegroup::Part>>),
    Replycaps(
        PartHeader,
        BoxFuture<'a, Result<capabilities::Capabilities>>,
    ),
    Pushkey(PartHeader, BoxFuture<'a, Result<()>>),
    Pushvars(PartHeader, BoxFuture<'a, Result<()>>),
}

impl<'a> Bundle2Item<'a> {
    #[cfg(test)]
    pub(crate) fn unwrap_start(self) -> StreamHeader {
        match self {
            Bundle2Item::Start(stream_header) => stream_header,
            other => panic!("expected item to be Start, was {:?}", other),
        }
    }
}

impl<'a> fmt::Debug for Bundle2Item<'a> {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        use crate::Bundle2Item::*;
        match *self {
            Start(ref header) => write!(f, "Bundle2Item::Start({:?})", header),
            Changegroup(ref header, _) => write!(f, "Bundle2Item::Changegroup({:?}, ...)", header),
            B2xCommonHeads(ref header, _) => {
                write!(f, "Bundle2Item::B2xCommonHeads({:?}, ...)", header)
            }
            B2xInfinitepush(ref header, _) => {
                write!(f, "Bundle2Item::B2xInfinitepush({:?}, ...)", header)
            }
            B2xInfinitepushBookmarks(ref header, _) => write!(
                f,
                "Bundle2Item::B2xInfinitepushBookmarks({:?}, ...)",
                header
            ),
            B2xInfinitepushMutation(ref header, _) => {
                write!(f, "Bundle2Item::B2xInfinitepushMutation({:?}, ...)", header)
            }
            B2xTreegroup2(ref header, _) => {
                write!(f, "Bundle2Item::B2xTreegroup2({:?}, ...)", header)
            }
            B2xRebasePack(ref header, _) => {
                write!(f, "Bundle2Item::B2xRebasePack({:?}, ...)", header)
            }
            B2xRebase(ref header, _) => write!(f, "Bundle2Item::B2xRebase({:?}, ...)", header),
            Replycaps(ref header, _) => write!(f, "Bundle2Item::Replycaps({:?}, ...)", header),
            Pushkey(ref header, _) => write!(f, "Bundle2Item::Pushkey({:?}, ...)", header),
            Pushvars(ref header, _) => write!(f, "Bundle2Item::Pushvars({:?}, ...)", header),
        }
    }
}

impl<'a> From<OldBundle2Item> for Bundle2Item<'a> {
    fn from(this: OldBundle2Item) -> Self {
        use crate::OldBundle2Item::*;
        match this {
            Start(header) => Bundle2Item::Start(header),
            Changegroup(header, stream) => {
                Bundle2Item::Changegroup(header, stream.compat().boxed())
            }
            B2xCommonHeads(header, stream) => {
                Bundle2Item::B2xCommonHeads(header, stream.compat().boxed())
            }
            B2xInfinitepush(header, stream) => {
                Bundle2Item::B2xInfinitepush(header, stream.compat().boxed())
            }
            B2xTreegroup2(header, stream) => {
                Bundle2Item::B2xTreegroup2(header, stream.compat().boxed())
            }
            // B2xInfinitepushBookmarks returns Bytes because this part is not going to be used.
            B2xInfinitepushBookmarks(header, stream) => {
                Bundle2Item::B2xInfinitepushBookmarks(header, stream.compat().boxed())
            }
            B2xInfinitepushMutation(header, stream) => {
                Bundle2Item::B2xInfinitepushMutation(header, stream.compat().boxed())
            }
            B2xRebasePack(header, stream) => {
                Bundle2Item::B2xRebasePack(header, stream.compat().boxed())
            }
            B2xRebase(header, stream) => Bundle2Item::B2xRebase(header, stream.compat().boxed()),
            Replycaps(header, future) => Bundle2Item::Replycaps(header, future.compat().boxed()),
            Pushkey(header, future) => Bundle2Item::Pushkey(header, future.compat().boxed()),
            Pushvars(header, future) => Bundle2Item::Pushvars(header, future.compat().boxed()),
        }
    }
}

pub(crate) enum OldBundle2Item {
    Start(StreamHeader),
    Changegroup(PartHeader, OldBoxStream<changegroup::Part, Error>),
    B2xCommonHeads(
        PartHeader,
        OldBoxStream<mercurial_types::HgChangesetId, Error>,
    ),
    B2xInfinitepush(PartHeader, OldBoxStream<changegroup::Part, Error>),
    B2xTreegroup2(PartHeader, OldBoxStream<wirepack::Part, Error>),
    // B2xInfinitepushBookmarks returns Bytes because this part is not going to be used.
    B2xInfinitepushBookmarks(PartHeader, OldBoxStream<bytes_old::Bytes, Error>),
    B2xInfinitepushMutation(
        PartHeader,
        OldBoxStream<Vec<mercurial_mutation::HgMutationEntry>, Error>,
    ),
    B2xRebasePack(PartHeader, OldBoxStream<wirepack::Part, Error>),
    B2xRebase(PartHeader, OldBoxStream<changegroup::Part, Error>),
    Replycaps(PartHeader, OldBoxFuture<capabilities::Capabilities, Error>),
    Pushkey(PartHeader, OldBoxFuture<(), Error>),
    Pushvars(PartHeader, OldBoxFuture<(), Error>),
}

/// Given bundle parts, returns a stream of Bytes that represent an encoded bundle with these parts
pub fn create_bundle_stream<C: Into<Option<async_compression::CompressorType>>>(
    parts: Vec<part_encode::PartEncodeBuilder>,
    ct: C,
) -> impl OldStream<Item = bytes_old::Bytes, Error = Error> {
    let (sender, receiver) = mpsc::channel::<Bytes>(1);
    // Sends either and empty Bytes if bundle generation was successful or an error.
    // Empty Bytes are used just to make chaining of streams below easier.
    let (result_sender, result_receiver) = oneshot::channel::<Result<Bytes>>();
    // Bundle2EncodeBuilder accepts writer which implements AsyncWrite. To workaround that we
    // use SinkToAsyncWrite. It implements AsyncWrite trait and sends everything that was written
    // into the Sender
    let mut bundle = Bundle2EncodeBuilder::new(SinkToAsyncWrite::new(sender));
    bundle.set_compressor_type(ct);
    for part in parts {
        bundle.add_part(part);
    }

    tokio::spawn(
        bundle
            .build()
            .then(move |val| {
                // Ignore send errors, because they can only happen if receiver was deallocated already
                match val {
                    Ok(_) => {
                        // Bundle was successfully generated, so there is nothing add.
                        // So just add empty bytes.
                        let _ = result_sender.send(Ok(Bytes::new()));
                    }
                    Err(err) => {
                        let _ = result_sender.send(Err(err));
                    }
                };
                Result::<_, Error>::Ok(())
            })
            .compat(),
    );

    receiver
        .map(Ok)
        .chain(result_receiver.into_stream().map_err(|_err| ()))
        .then(|entry| match entry {
            Ok(res) => res,
            Err(()) => bail!("error while receiving gettreepack response from the channel"),
        })
}


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

//! Scaffolding for encoding bundle2 parts.

use std::fmt;
use std::fmt::Debug;
use std::fmt::Formatter;
use std::mem;

use anyhow::Error;
use anyhow::Result;
use bytes::Bytes as BytesNew;
use bytes_old::Bytes;
use futures_ext::BoxStream;
use futures_ext::StreamExt;
use futures_old::Async;
use futures_old::Future;
use futures_old::Poll;
use futures_old::Stream;

use crate::chunk::Chunk;
use crate::part_header::PartHeader;
use crate::part_header::PartHeaderBuilder;
use crate::part_header::PartHeaderType;
use crate::part_header::PartId;

/// Represents a stream of chunks produced by the individual part handler.
pub struct ChunkStream(BoxStream<Chunk, Error>);

impl Debug for ChunkStream {
    fn fmt(&self, f: &mut Formatter<'_>) -> fmt::Result {
        f.debug_tuple("ChunkStream").finish()
    }
}

#[derive(Debug)]
pub enum PartEncodeData {
    None,
    Fixed(Chunk),
    Generated(ChunkStream),
}

pub struct PartEncodeBuilder {
    headerb: PartHeaderBuilder,
    data: PartEncodeData,
}

#[derive(Debug)]
pub struct PartEncode {
    state: GenerationState,
}

#[derive(Debug)]
enum GenerationState {
    NotStarted(PartHeader, PartEncodeData),
    Fixed(Chunk),
    Generating(ChunkStream),
    EmptyChunk,
    Done,
    Invalid,
}

impl GenerationState {
    fn take(&mut self) -> Self {
        mem::replace(self, GenerationState::Invalid)
    }
}

impl PartEncodeBuilder {
    pub fn mandatory(part_type: PartHeaderType) -> Result<Self> {
        Ok(PartEncodeBuilder {
            headerb: PartHeaderBuilder::new(part_type, true)?,
            data: PartEncodeData::None,
        })
    }

    pub fn advisory(part_type: PartHeaderType) -> Result<Self> {
        Ok(PartEncodeBuilder {
            headerb: PartHeaderBuilder::new(part_type, false)?,
            data: PartEncodeData::None,
        })
    }

    #[inline]
    pub fn add_mparam<S, B>(&mut self, key: S, val: B) -> Result<&mut Self>
    where
        S: Into<String>,
        B: Into<BytesNew>,
    {
        self.headerb.add_mparam(key, val)?;
        Ok(self)
    }

    #[inline]
    pub fn add_aparam<S, B>(&mut self, key: S, val: B) -> Result<&mut Self>
    where
        S: Into<String>,
        B: Into<BytesNew>,
    {
        self.headerb.add_aparam(key, val)?;
        Ok(self)
    }

    pub fn set_data_fixed<T: Into<Chunk>>(&mut self, data: T) -> &mut Self {
        self.data = PartEncodeData::Fixed(data.into());
        self
    }

    pub fn set_data_bytes<T: Into<Bytes>>(&mut self, data: T) -> Result<&mut Self> {
        self.data = PartEncodeData::Fixed(Chunk::new(data.into())?);
        Ok(self)
    }

    pub fn set_data_future<T>(&mut self, future: T) -> &mut Self
    where
        T: Future<Error = Error> + Send + 'static,
        T::Item: Into<Bytes>,
    {
        let stream = future.and_then(Chunk::new).into_stream();
        self.set_data_generated(stream)
    }

    pub fn set_data_generated<S>(&mut self, stream: S) -> &mut Self
    where
        S: 'static + Stream<Item = Chunk, Error = Error> + Send,
    {
        let stream = stream.filter(|chunk| !chunk.is_empty());
        self.data = PartEncodeData::Generated(ChunkStream(stream.boxify()));
        self
    }

    pub fn build(self, part_id: PartId) -> PartEncode {
        PartEncode {
            state: GenerationState::NotStarted(self.headerb.build(part_id), self.data),
        }
    }
}

impl Stream for PartEncode {
    type Item = Chunk;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Chunk>, Error> {
        let (ret, next_state) = Self::poll_next(self.state.take());
        self.state = next_state;
        ret
    }
}

impl PartEncode {
    fn poll_next(state: GenerationState) -> (Poll<Option<Chunk>, Error>, GenerationState) {
        // An individual part has three sections:
        // (1) a header (1 chunk)
        // (2) the payload (0-many chunks)
        // (3) an empty chunk, indicating the end of the stream.
        //
        // The state machine captures the generation as:
        // NotStarted = header not output yet
        // Generating = payload currently being generated by inner stream
        // Fixed = fixed-length payload (no generation, just one chunk)
        // EmptyChunk = end of payload (or no payload)
        // Done = chunk completed
        // Invalid = some sort of error occured
        use self::GenerationState::*;

        match state {
            NotStarted(header, data) => {
                let header_chunk = header.encode();
                let next_state = match data {
                    PartEncodeData::Fixed(b) => Fixed(b),
                    PartEncodeData::None => EmptyChunk,
                    PartEncodeData::Generated(ChunkStream(stream)) => {
                        Generating(ChunkStream(stream))
                    }
                };
                (Ok(Async::Ready(Some(header_chunk))), next_state)
            }
            Generating(ChunkStream(mut stream)) => {
                match stream.poll() {
                    Ok(Async::Ready(Some(v))) => {
                        // TODO: don't send too large or too small chunks to clients
                        (Ok(Async::Ready(Some(v))), Generating(ChunkStream(stream)))
                    }
                    Ok(Async::Ready(None)) => (Ok(Async::Ready(Some(Chunk::empty()))), Done),
                    Ok(Async::NotReady) => (Ok(Async::NotReady), Generating(ChunkStream(stream))),
                    // TODO: produce an error part for (some kinds of?) errors
                    Err(e) => (Err(e), Generating(ChunkStream(stream))),
                }
            }
            Fixed(chunk) => (Ok(Async::Ready(Some(chunk))), EmptyChunk),
            EmptyChunk => (Ok(Async::Ready(Some(Chunk::empty()))), Done),
            Done => (Ok(Async::Ready(None)), Done),
            Invalid => panic!("invalid state"),
        }
    }
}