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

use anyhow::bail;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobrepo::BlobRepo;
use blobstore::Loadable;
use bytes::Bytes;
use cloned::cloned;
use context::CoreContext;
use filestore::FetchKey;
use filestore::FilestoreConfigRef;
use futures::future::BoxFuture;
use futures::Future;
use futures::FutureExt;
use futures::Stream;
use futures::TryFutureExt;
use futures::TryStreamExt;
use futures_ext::future::TryShared;
use futures_ext::FbTryFutureExt;
use mercurial_bundles::changegroup::CgDeltaChunk;
use mercurial_types::blobs::ContentBlobMeta;
use mercurial_types::blobs::File;
use mercurial_types::blobs::UploadHgFileContents;
use mercurial_types::blobs::UploadHgFileEntry;
use mercurial_types::blobs::UploadHgNodeHash;
use mercurial_types::delta;
use mercurial_types::Delta;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgNodeHash;
use mercurial_types::HgNodeKey;
use mercurial_types::MPath;
use mercurial_types::RepoPath;
use mercurial_types::RevFlags;
use mercurial_types::NULL_HASH;
use quickcheck::Arbitrary;
use quickcheck::Gen;
use remotefilelog::create_raw_filenode_blob;
use repo_blobstore::RepoBlobstoreArc;
use repo_blobstore::RepoBlobstoreRef;

use crate::stats::*;
use crate::upload_blobs::UploadableHgBlob;

#[derive(Debug, Eq, PartialEq)]
pub struct FilelogDeltaed {
    pub path: MPath,
    pub chunk: CgDeltaChunk,
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub enum FilelogData {
    RawBytes(Bytes),
    LfsMetaData(ContentBlobMeta),
}

#[derive(Debug, Clone, Eq, PartialEq)]
pub struct Filelog {
    pub node_key: HgNodeKey,
    pub p1: Option<HgNodeHash>,
    pub p2: Option<HgNodeHash>,
    pub linknode: HgNodeHash,
    pub data: FilelogData,
    pub flags: RevFlags,
}

impl UploadableHgBlob for Filelog {
    // * Shared is required here because a single file node can be referred to by more than
    //   one changeset, and all of those will want to refer to the corresponding future.
    type Value = TryShared<BoxFuture<'static, Result<(HgFileNodeId, RepoPath)>>>;

    fn upload(self, ctx: &CoreContext, repo: &BlobRepo) -> Result<(HgNodeKey, Self::Value)> {
        let node_key = self.node_key;
        let path = match &node_key.path {
            RepoPath::FilePath(path) => path.clone(),
            other => bail!("internal error: expected file path, got {}", other),
        };

        // If LFSMetaData
        let contents = match self.data {
            FilelogData::RawBytes(bytes) => {
                UploadHgFileContents::RawBytes(bytes, *repo.filestore_config())
            }
            FilelogData::LfsMetaData(meta) => UploadHgFileContents::ContentUploaded(meta),
        };

        let upload = UploadHgFileEntry {
            upload_node_id: UploadHgNodeHash::Checked(node_key.hash),
            contents,
            p1: self.p1.map(HgFileNodeId::new),
            p2: self.p2.map(HgFileNodeId::new),
        };

        let fut = upload.upload_with_path(ctx.clone(), repo.repo_blobstore_arc(), path);
        Ok((node_key, fut.boxed().try_shared()))
    }
}

pub(crate) fn convert_to_revlog_filelog(
    ctx: CoreContext,
    repo: BlobRepo,
    deltaed: impl Stream<Item = Result<FilelogDeltaed>> + Send + 'static,
) -> impl Stream<Item = Result<Filelog>> {
    let mut delta_cache = DeltaCache::new(repo.clone());
    deltaed
        .map_ok(move |FilelogDeltaed { path, chunk }| {
            let CgDeltaChunk {
                node,
                base,
                delta,
                p1,
                p2,
                linknode,
                flags,
            } = chunk;

            delta_cache
                .decode(&ctx, node, base.into_option(), delta)
                .and_then({
                    cloned!(ctx, node, path, repo);
                    move |data| {
                        cloned!(ctx, node, path, repo);
                        let flags = flags.unwrap_or(RevFlags::REVIDX_DEFAULT_FLAGS);
                        async move {
                            let file_log_data = get_filelog_data(&ctx, &repo, data, flags).await?;
                            Ok(Filelog {
                                node_key: HgNodeKey {
                                    path: RepoPath::FilePath(path),
                                    hash: node,
                                },
                                p1: p1.into_option(),
                                p2: p2.into_option(),
                                linknode,
                                data: file_log_data,
                                flags,
                            })
                        }
                    }
                })
                .map(move |res| {
                    res.with_context(move || {
                        format!(
                            "While decoding delta cache for file id {}, path {}",
                            node, path
                        )
                    })
                })
        })
        .try_buffer_unordered(100)
}

async fn generate_lfs_meta_data(
    ctx: &CoreContext,
    repo: &BlobRepo,
    data: Bytes,
) -> Result<ContentBlobMeta, Error> {
    // TODO(anastasiyaz): check size
    let lfs_content = File::data_only(data).get_lfs_content()?;
    let content_id = FetchKey::from(lfs_content.oid())
        .load(ctx, repo.repo_blobstore())
        .await?;
    Ok(ContentBlobMeta {
        id: content_id,
        copy_from: lfs_content.copy_from(),
        size: lfs_content.size(),
    })
}

async fn get_filelog_data(
    ctx: &CoreContext,
    repo: &BlobRepo,
    data: Bytes,
    flags: RevFlags,
) -> Result<FilelogData, Error> {
    if flags.contains(RevFlags::REVIDX_EXTSTORED) {
        let cbmeta = generate_lfs_meta_data(ctx, repo, data).await?;
        Ok(FilelogData::LfsMetaData(cbmeta))
    } else {
        Ok(FilelogData::RawBytes(data))
    }
}

struct DeltaCache {
    repo: BlobRepo,
    bytes_cache: HashMap<HgNodeHash, TryShared<BoxFuture<'static, Result<Bytes>>>>,
}

impl DeltaCache {
    fn new(repo: BlobRepo) -> Self {
        Self {
            repo,
            bytes_cache: HashMap::new(),
        }
    }

    fn decode(
        &mut self,
        ctx: &CoreContext,
        node: HgNodeHash,
        base: Option<HgNodeHash>,
        delta: Delta,
    ) -> impl Future<Output = Result<Bytes>> {
        let bytes = self.bytes_cache.get(&node).cloned().unwrap_or_else(|| {
            let bytes = {
                let vec_u8 = match base {
                    None => async move {
                        delta::apply(b"", &delta)
                            .with_context(|| format!("File content empty, delta: {:?}", delta))
                    }
                    .left_future(),
                    Some(base) => self
                        .apply_delta_on_base(ctx, base, delta)
                        .map_err(move |err| {
                            err.context(format!(
                                "While looking for base {:?} to apply on delta {:?}",
                                base, node
                            ))
                        })
                        .right_future(),
                };
                vec_u8.map_ok(Bytes::from)
            };

            let bytes = bytes.boxed().try_shared();

            if self.bytes_cache.insert(node, bytes.clone()).is_some() {
                panic!("Logic error: byte cache returned Some for HashMap::get with node");
            }
            bytes
        });

        async move {
            let bytes = bytes.await?;

            let fsize = mem::size_of_val(&bytes) as i64;
            STATS::deltacache_fsize.add_value(fsize);
            STATS::deltacache_fsize_large.add_value(fsize);

            Ok(bytes)
        }
    }

    fn apply_delta_on_base(
        &self,
        ctx: &CoreContext,
        base: HgNodeHash,
        delta: Delta,
    ) -> impl Future<Output = Result<Vec<u8>>> {
        let cache_entry = self.bytes_cache.get(&base).cloned();
        cloned!(ctx, self.repo);
        async move {
            let bytes = match cache_entry {
                Some(bytes) => bytes.clone().await?,
                None => {
                    let validate_hash = false;
                    create_raw_filenode_blob(&ctx, &repo, HgFileNodeId::new(base), validate_hash)
                        .await?
                }
            };
            delta::apply(&bytes, &delta)
                .with_context(|| format!("File content: {:?} delta: {:?}", bytes, delta))
        }
    }
}

impl Arbitrary for Filelog {
    fn arbitrary(g: &mut Gen) -> Self {
        Filelog {
            node_key: HgNodeKey {
                path: RepoPath::FilePath(MPath::arbitrary(g)),
                hash: HgNodeHash::arbitrary(g),
            },
            p1: HgNodeHash::arbitrary(g).into_option(),
            p2: HgNodeHash::arbitrary(g).into_option(),
            linknode: HgNodeHash::arbitrary(g),
            data: FilelogData::RawBytes(Bytes::from(Vec::<u8>::arbitrary(g))),
            flags: RevFlags::REVIDX_DEFAULT_FLAGS,
        }
    }

    fn shrink(&self) -> Box<dyn Iterator<Item = Self>> {
        fn append(result: &mut Vec<Filelog>, f: Filelog) {
            result.append(&mut f.shrink().collect());
            result.push(f);
        }

        let mut result = Vec::new();

        if self.node_key.hash != NULL_HASH {
            let mut f = self.clone();
            f.node_key.hash = NULL_HASH;
            append(&mut result, f);
        }

        if self.p1.is_some() {
            let mut f = self.clone();
            f.p1 = None;
            append(&mut result, f);
        }

        if self.p2.is_some() {
            let mut f = self.clone();
            f.p2 = None;
            append(&mut result, f);
        }

        if self.linknode != NULL_HASH {
            let mut f = self.clone();
            f.linknode = NULL_HASH;
            append(&mut result, f);
        }

        if let FilelogData::RawBytes(ref bytes) = self.data {
            if !bytes.is_empty() {
                let mut f = self.clone();
                f.data = FilelogData::RawBytes(Bytes::from(Vec::new()));
                append(&mut result, f);
            }
        }

        Box::new(result.into_iter())
    }
}

#[cfg(test)]
mod tests {
    use std::cmp::min;

    use fbinit::FacebookInit;
    use futures::stream::iter;
    use itertools::assert_equal;
    use itertools::EitherOrBoth;
    use itertools::Itertools;
    use mercurial_types::delta::Fragment;
    use mercurial_types::NULL_HASH;
    use mercurial_types_mocks::nodehash::*;
    use quickcheck_macros::quickcheck;

    use super::*;

    struct NodeHashGen {
        bytes: Vec<u8>,
    }

    impl NodeHashGen {
        fn new() -> Self {
            Self {
                bytes: Vec::from(NULL_HASH.as_ref()),
            }
        }

        fn next(&mut self) -> HgNodeHash {
            for i in 0..self.bytes.len() {
                if self.bytes[i] == 255 {
                    self.bytes[i] = 0;
                } else {
                    self.bytes[i] += 1;
                    return HgNodeHash::from_bytes(self.bytes.as_slice()).unwrap();
                }
            }

            panic!("NodeHashGen overflow");
        }
    }

    async fn check_conversion<I, J>(ctx: CoreContext, inp: I, exp: J)
    where
        I: IntoIterator<Item = FilelogDeltaed>,
        J: IntoIterator<Item = Filelog>,
    {
        let repo_factory = test_repo_factory::build_empty(ctx.fb).await.unwrap();
        let result = convert_to_revlog_filelog(
            ctx,
            repo_factory,
            iter(inp.into_iter().map(Ok).collect::<Vec<_>>()),
        )
        .try_collect::<Vec<_>>()
        .await
        .unwrap();

        assert_equal(result, exp);
    }

    fn filelog_to_deltaed(f: &Filelog) -> FilelogDeltaed {
        match f.data {
            FilelogData::RawBytes(ref bytes) => FilelogDeltaed {
                path: f.node_key.path.mpath().unwrap().clone(),
                chunk: CgDeltaChunk {
                    node: f.node_key.hash.clone(),
                    p1: f.p1.clone().unwrap_or(NULL_HASH),
                    p2: f.p2.clone().unwrap_or(NULL_HASH),
                    base: NULL_HASH,
                    linknode: f.linknode.clone(),
                    delta: Delta::new_fulltext(bytes.as_ref()),
                    flags: None,
                },
            },
            _ => panic!("RawBytes FilelogData is only supported in tests"),
        }
    }

    fn filelog_compute_delta(b1: &FilelogData, b2: &FilelogData) -> Delta {
        match (b1, b2) {
            (FilelogData::RawBytes(b1_data), FilelogData::RawBytes(b2_data)) => {
                compute_delta(b1_data, b2_data)
            }
            _ => panic!("RawBytes FilelogData is only supported in tests"),
        }
    }

    fn compute_delta(b1: &[u8], b2: &[u8]) -> Delta {
        let mut frags = Vec::new();
        let mut start = 0;
        let mut frag = Vec::new();
        for (idx, val) in b1.iter().zip_longest(b2.iter()).enumerate() {
            match val {
                EitherOrBoth::Both(v1, v2) => {
                    if v1 == v2 && !frag.is_empty() {
                        frags.push(Fragment {
                            start,
                            end: start + frag.len(),
                            content: std::mem::take(&mut frag),
                        });
                    } else if v1 != v2 {
                        if frag.is_empty() {
                            start = idx;
                        }
                        frag.push(*v2);
                    }
                }
                EitherOrBoth::Left(_) => continue,
                EitherOrBoth::Right(v) => {
                    if frag.is_empty() {
                        start = idx;
                    }
                    frag.push(*v)
                }
            }
        }
        if !frag.is_empty() {
            frags.push(Fragment {
                start,
                end: min(start + frag.len(), b1.len()),
                content: std::mem::take(&mut frag),
            });
        }
        if b1.len() > b2.len() {
            frags.push(Fragment {
                start: b2.len(),
                end: b1.len(),
                content: Vec::new(),
            });
        }

        Delta::new(frags).unwrap()
    }

    #[fbinit::test]
    async fn two_fulltext_files(fb: FacebookInit) {
        let ctx = CoreContext::test_mock(fb);
        let f1 = Filelog {
            node_key: HgNodeKey {
                path: RepoPath::FilePath(MPath::new(b"test").unwrap()),
                hash: ONES_HASH,
            },
            p1: Some(TWOS_HASH),
            p2: Some(THREES_HASH),
            linknode: FOURS_HASH,
            data: FilelogData::RawBytes(Bytes::from("test file content")),
            flags: RevFlags::REVIDX_DEFAULT_FLAGS,
        };

        let f2 = Filelog {
            node_key: HgNodeKey {
                path: RepoPath::FilePath(MPath::new(b"test2").unwrap()),
                hash: FIVES_HASH,
            },
            p1: Some(SIXES_HASH),
            p2: Some(SEVENS_HASH),
            linknode: EIGHTS_HASH,
            data: FilelogData::RawBytes(Bytes::from("test2 file content")),
            flags: RevFlags::REVIDX_DEFAULT_FLAGS,
        };

        check_conversion(
            ctx,
            vec![filelog_to_deltaed(&f1), filelog_to_deltaed(&f2)],
            vec![f1, f2],
        )
        .await;
    }

    async fn files_check_order(ctx: CoreContext, correct_order: bool) {
        let f1 = Filelog {
            node_key: HgNodeKey {
                path: RepoPath::FilePath(MPath::new(b"test").unwrap()),
                hash: ONES_HASH,
            },
            p1: Some(TWOS_HASH),
            p2: Some(THREES_HASH),
            linknode: FOURS_HASH,
            data: FilelogData::RawBytes(Bytes::from("test file content")),
            flags: RevFlags::REVIDX_DEFAULT_FLAGS,
        };

        let f2 = Filelog {
            node_key: HgNodeKey {
                path: RepoPath::FilePath(MPath::new(b"test2").unwrap()),
                hash: FIVES_HASH,
            },
            p1: Some(SIXES_HASH),
            p2: Some(SEVENS_HASH),
            linknode: EIGHTS_HASH,
            data: FilelogData::RawBytes(Bytes::from("test2 file content")),
            flags: RevFlags::REVIDX_DEFAULT_FLAGS,
        };

        let f1_deltaed = filelog_to_deltaed(&f1);
        let mut f2_deltaed = filelog_to_deltaed(&f2);

        f2_deltaed.chunk.base = f1.node_key.hash.clone();
        f2_deltaed.chunk.delta = filelog_compute_delta(&f1.data, &f2.data);

        let inp = if correct_order {
            vec![f1_deltaed, f2_deltaed]
        } else {
            vec![f2_deltaed, f1_deltaed]
        };

        let repo_factory = test_repo_factory::build_empty(ctx.fb).await.unwrap();
        let result = convert_to_revlog_filelog(ctx, repo_factory, iter(inp.into_iter().map(Ok)))
            .try_collect::<Vec<_>>()
            .await;

        match result {
            Ok(_) => assert!(
                correct_order,
                "Successfuly converted even though order was incorrect"
            ),
            Err(_) => assert!(
                !correct_order,
                "Filed to convert even though order was correct"
            ),
        }
    }

    #[fbinit::test]
    async fn files_order_correct(fb: FacebookInit) {
        files_check_order(CoreContext::test_mock(fb), true).await;
    }

    #[fbinit::test]
    async fn files_order_incorrect(fb: FacebookInit) {
        files_check_order(CoreContext::test_mock(fb), false).await;
    }

    #[quickcheck]
    fn sanitycheck_delta_computation(b1: Vec<u8>, b2: Vec<u8>) -> bool {
        assert_equal(&b2, &delta::apply(&b1, &compute_delta(&b1, &b2)).unwrap());
        true
    }

    #[quickcheck_async::tokio]
    async fn correct_conversion_single(fb: FacebookInit, f: Filelog) -> bool {
        let ctx = CoreContext::test_mock(fb);
        check_conversion(ctx, vec![filelog_to_deltaed(&f)], vec![f]).await;

        true
    }

    #[quickcheck_async::tokio]
    async fn correct_conversion_delta_against_first(
        fb: FacebookInit,
        f: Filelog,
        fs: Vec<Filelog>,
    ) -> bool {
        let ctx = CoreContext::test_mock(fb);
        let mut hash_gen = NodeHashGen::new();

        let mut f = f.clone();
        f.node_key.hash = hash_gen.next();

        let mut fs = fs.clone();
        for el in fs.iter_mut() {
            el.node_key.hash = hash_gen.next();
        }

        let mut deltas = vec![filelog_to_deltaed(&f)];
        for filelog in &fs {
            let mut delta = filelog_to_deltaed(filelog);
            delta.chunk.base = f.node_key.hash.clone();
            delta.chunk.delta = filelog_compute_delta(&f.data, &filelog.data);
            deltas.push(delta);
        }

        check_conversion(ctx, deltas, vec![f].into_iter().chain(fs)).await;

        true
    }

    #[quickcheck_async::tokio]
    async fn correct_conversion_delta_against_next(fb: FacebookInit, fs: Vec<Filelog>) -> bool {
        let ctx = CoreContext::test_mock(fb);
        let mut hash_gen = NodeHashGen::new();

        let mut fs = fs.clone();
        for el in fs.iter_mut() {
            el.node_key.hash = hash_gen.next();
        }

        let deltas = {
            let mut it = fs.iter();
            let mut deltas = match it.next() {
                None => return true, // empty test case
                Some(f) => vec![filelog_to_deltaed(f)],
            };

            for (prev, next) in fs.iter().zip(it) {
                let mut delta = filelog_to_deltaed(next);
                delta.chunk.base = prev.node_key.hash.clone();
                delta.chunk.delta = filelog_compute_delta(&prev.data, &next.data);
                deltas.push(delta);
            }

            deltas
        };

        check_conversion(ctx, deltas, fs).await;

        true
    }
}


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

mod changeset;
mod filelog;
mod split;

pub(crate) use self::changeset::convert_to_revlog_changesets;
pub(crate) use self::filelog::convert_to_revlog_filelog;
pub(crate) use self::filelog::Filelog;
pub(crate) use self::split::split_changegroup;


/*
 * 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::pin::Pin;

use anyhow::anyhow;
use anyhow::Result;
use futures::future;
use futures::ready;
use futures::task::Context;
use futures::task::Poll;
use futures::Stream;
use futures::StreamExt;
use futures::TryFutureExt;
use futures::TryStreamExt;
use futures_ext::FbStreamExt;
use mercurial_bundles::changegroup::Part;
use mercurial_bundles::changegroup::Section;
use pin_project::pin_project;

use crate::changegroup::changeset::ChangesetDeltaed;
use crate::changegroup::filelog::FilelogDeltaed;

pub(crate) fn split_changegroup(
    cg2s: impl Stream<Item = Result<Part>> + Send + 'static,
) -> (
    impl Stream<Item = Result<ChangesetDeltaed>>,
    impl Stream<Item = Result<FilelogDeltaed>>,
) {
    let cg2s = CheckEnd::new(cg2s).boxed();

    let (changesets, remainder) = cg2s
        .try_take_while(|part| {
            future::ready(match part {
                &Part::CgChunk(Section::Changeset, _) => Ok(true),
                &Part::SectionEnd(Section::Changeset) => Ok(false),
                bad => Err(anyhow!("Expected Changeset chunk or end, found: {:?}", bad)),
            })
        })
        .return_remainder();

    let changesets = changesets
        .and_then(|part| {
            future::ready(match part {
                Part::CgChunk(Section::Changeset, chunk) => Ok(ChangesetDeltaed { chunk }),
                bad => Err(anyhow!("Expected Changeset chunk, found: {:?}", bad)),
            })
        })
        .map_err(|err| err.context("While extracting Changesets from Changegroup"));

    let filelogs = remainder
        .err_into()
        .map_ok(|take_while_stream| take_while_stream.into_inner())
        .try_flatten_stream()
        .try_skip_while({
            let mut seen_manifest_end = false;
            move |part| future::ready(match part {
                &Part::SectionEnd(Section::Manifest) if !seen_manifest_end => {
                    seen_manifest_end = true;
                    Ok(true)
                }
                _ if seen_manifest_end => Ok(false),
                bad => Err(anyhow!("Expected Manifest end, found: {:?}", bad)),
            })
        })
        .map_err(|err| {
            err.context("While skipping Manifests in Changegroup")
        })
        .try_filter_map({
            let mut seen_path = None;
            move |part| {
                if let Some(seen_path) = &seen_path {
                    match &part {
                        &Part::CgChunk(Section::Filelog(ref path), _)
                        | &Part::SectionEnd(Section::Filelog(ref path)) => {
                            if seen_path != path {
                                return future::ready(Err(anyhow!(
                                    "Mismatched path found {0} ({0:?}) != {1} ({1:?}), for part: {2:?}",
                                    seen_path,
                                    path,
                                    part
                                )));
                            }
                        }
                        _ => (), // Handled in the next pattern-match
                    }
                }

                future::ready(match part {
                    Part::CgChunk(Section::Filelog(path), chunk) => {
                        seen_path = Some(path.clone());
                        Ok(Some(FilelogDeltaed { path, chunk }))
                    }
                    Part::SectionEnd(Section::Filelog(_)) if seen_path.is_some() => {
                        seen_path = None;
                        Ok(None)
                    }
                    Part::SectionEnd(Section::Treemanifest) => Ok(None),
                    // Checking that there is exactly one Part::end is is covered by CheckEnd
                    // wrapper
                    Part::End if seen_path.is_none() => Ok(None),
                    bad => Err(
                        if seen_path.is_some() {
                            anyhow!(
                                "Expected Filelog chunk or end, seen_path was {:?}, found: {:?}",
                                seen_path,
                                bad
                            )
                        } else {
                            anyhow!(
                                "Expected Filelog chunk or Part::End, seen_path was {:?}, found: {:?}",
                                seen_path,
                                bad
                            )
                        }
                    )
                })
            }
        })
        .map_err(|err| {
            err.context("While extracting Filelogs from Changegroup")
        });

    (changesets, filelogs)
}

/// Wrapper for Stream of Part that is supposed to ensure that there is exactly one Part::End in
/// the stream and that it is at the end of it
#[pin_project]
struct CheckEnd<S> {
    #[pin]
    cg2s: S,
    seen_end: bool,
}

impl<S> CheckEnd<S>
where
    S: Stream<Item = Result<Part>>,
{
    fn new(cg2s: S) -> Self {
        Self {
            cg2s,
            seen_end: false,
        }
    }
}

impl<S> Stream for CheckEnd<S>
where
    S: Stream<Item = Result<Part>>,
{
    type Item = S::Item;

    fn poll_next(self: Pin<&mut Self>, cx: &mut Context) -> Poll<Option<Self::Item>> {
        let mut this = self.project();

        let part = ready!(this.cg2s.as_mut().poll_next(cx));
        Poll::Ready((|| {
            match part {
                Some(Ok(Part::End)) => {
                    if *this.seen_end {
                        return Some(Err(anyhow!("More than one Part::End noticed")));
                    }
                    *this.seen_end = true;
                }
                None => {
                    if !*this.seen_end {
                        return Some(Err(anyhow!(
                            "End of stream reached, but no Part::End noticed"
                        )));
                    }
                }
                _ => {} // all good, proceed
            }
            part
        })())
    }
}

#[cfg(test)]
mod tests {
    use futures::stream::iter;
    use itertools::assert_equal;
    use itertools::equal;
    use mercurial_bundles::changegroup::CgDeltaChunk;
    use mercurial_types::MPath;

    use super::*;

    async fn check_splitting<S, I, J>(cg2s: S, exp_cs: I, exp_fs: J) -> bool
    where
        S: Stream<Item = Result<Part>> + Send + 'static,
        I: IntoIterator<Item = ChangesetDeltaed>,
        J: IntoIterator<Item = FilelogDeltaed>,
    {
        let (cs, fs) = split_changegroup(cg2s);

        let cs = cs
            .try_collect::<Vec<_>>()
            .await
            .expect("error in changesets");
        let fs = fs
            .try_collect::<Vec<_>>()
            .await
            .expect("error in changesets");

        equal(cs, exp_cs) && equal(fs, exp_fs)
    }

    #[tokio::test]
    async fn splitting_empty() {
        assert!(
            check_splitting(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                ),
                vec![],
                vec![],
            )
            .await
        )
    }

    #[quickcheck_async::tokio]
    async fn splitting_minimal(c: CgDeltaChunk, f: CgDeltaChunk, f_p: MPath) -> bool {
        check_splitting(
            iter(
                vec![
                    Part::CgChunk(Section::Changeset, c.clone()),
                    Part::SectionEnd(Section::Changeset),
                    Part::SectionEnd(Section::Manifest),
                    Part::End,
                ]
                .into_iter()
                .map(Ok),
            ),
            vec![ChangesetDeltaed { chunk: c.clone() }],
            vec![],
        )
        .await
            && check_splitting(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::CgChunk(Section::Filelog(f_p.clone()), f.clone()),
                        Part::SectionEnd(Section::Filelog(f_p.clone())),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                ),
                vec![],
                vec![FilelogDeltaed {
                    path: f_p.clone(),
                    chunk: f.clone(),
                }],
            )
            .await
            && check_splitting(
                iter(
                    vec![
                        Part::CgChunk(Section::Changeset, c.clone()),
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::CgChunk(Section::Filelog(f_p.clone()), f.clone()),
                        Part::SectionEnd(Section::Filelog(f_p.clone())),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                ),
                vec![ChangesetDeltaed { chunk: c.clone() }],
                vec![FilelogDeltaed {
                    path: f_p.clone(),
                    chunk: f.clone(),
                }],
            )
            .await
    }

    #[quickcheck_async::tokio]
    async fn splitting_complex(
        c1: CgDeltaChunk,
        c2: CgDeltaChunk,
        f1: CgDeltaChunk,
        f1_bis: CgDeltaChunk,
        f1_p: MPath,
        f2: CgDeltaChunk,
        f2_bis: CgDeltaChunk,
        f2_p: MPath,
    ) -> bool {
        check_splitting(
            iter(
                vec![
                    Part::CgChunk(Section::Changeset, c1.clone()),
                    Part::CgChunk(Section::Changeset, c2.clone()),
                    Part::SectionEnd(Section::Changeset),
                    Part::SectionEnd(Section::Manifest),
                    Part::CgChunk(Section::Filelog(f1_p.clone()), f1.clone()),
                    Part::CgChunk(Section::Filelog(f1_p.clone()), f1_bis.clone()),
                    Part::SectionEnd(Section::Filelog(f1_p.clone())),
                    Part::CgChunk(Section::Filelog(f2_p.clone()), f2.clone()),
                    Part::CgChunk(Section::Filelog(f2_p.clone()), f2_bis.clone()),
                    Part::SectionEnd(Section::Filelog(f2_p.clone())),
                    Part::End,
                ]
                .into_iter()
                .map(Ok),
            ),
            vec![
                ChangesetDeltaed { chunk: c1 },
                ChangesetDeltaed { chunk: c2 },
            ],
            vec![
                FilelogDeltaed {
                    path: f1_p.clone(),
                    chunk: f1,
                },
                FilelogDeltaed {
                    path: f1_p,
                    chunk: f1_bis,
                },
                FilelogDeltaed {
                    path: f2_p.clone(),
                    chunk: f2,
                },
                FilelogDeltaed {
                    path: f2_p,
                    chunk: f2_bis,
                },
            ],
        )
        .await
    }

    #[quickcheck_async::tokio]
    async fn splitting_error_filelog_end(f: CgDeltaChunk, f1_p: MPath, f2_p: MPath) -> bool {
        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::SectionEnd(Section::Filelog(f1_p.clone())),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }

        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::CgChunk(Section::Filelog(f1_p.clone()), f.clone()),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }

        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::CgChunk(Section::Filelog(f1_p.clone()), f.clone()),
                        Part::SectionEnd(Section::Filelog(f2_p.clone())),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(f1_p == f2_p || fs.try_collect::<Vec<_>>().await.is_err());
        }

        true
    }

    #[quickcheck_async::tokio]
    async fn splitting_error_manifest(
        c: CgDeltaChunk,
        m: CgDeltaChunk,
        f: CgDeltaChunk,
        f_p: MPath,
    ) -> bool {
        let (cs, fs) = split_changegroup(
            iter(
                vec![
                    Part::CgChunk(Section::Changeset, c.clone()),
                    Part::SectionEnd(Section::Changeset),
                    Part::CgChunk(Section::Manifest, m.clone()),
                    Part::SectionEnd(Section::Manifest),
                    Part::CgChunk(Section::Filelog(f_p.clone()), f.clone()),
                    Part::SectionEnd(Section::Filelog(f_p.clone())),
                    Part::End,
                ]
                .into_iter()
                .map(Ok),
            )
            .boxed(),
        );

        equal(
            cs.try_collect::<Vec<_>>().await.unwrap(),
            vec![ChangesetDeltaed { chunk: c }],
        ) && fs.try_collect::<Vec<_>>().await.is_err()
    }

    #[tokio::test]
    async fn splitting_error_two_ends() {
        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }

        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::SectionEnd(Section::Manifest),
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }

        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                        Part::End,
                        Part::End,
                    ]
                    .into_iter()
                    .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }
    }

    #[tokio::test]
    async fn splitting_error_missing_end() {
        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![Part::SectionEnd(Section::Manifest), Part::End]
                        .into_iter()
                        .map(Ok),
                )
                .boxed(),
            );

            assert!(cs.try_collect::<Vec<_>>().await.is_err());
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }

        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![Part::SectionEnd(Section::Changeset), Part::End]
                        .into_iter()
                        .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }

        {
            let (cs, fs) = split_changegroup(
                iter(
                    vec![
                        Part::SectionEnd(Section::Changeset),
                        Part::SectionEnd(Section::Manifest),
                    ]
                    .into_iter()
                    .map(Ok),
                )
                .boxed(),
            );

            assert_equal(cs.try_collect::<Vec<_>>().await.unwrap(), vec![]);
            assert!(fs.try_collect::<Vec<_>>().await.is_err());
        }
    }
}


/*
 * 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::fmt::Debug;
use std::mem;

use anyhow::Error;
use anyhow::Result;
use bytes::Bytes;
use futures::compat::Compat;
use futures::stream::BoxStream;
use futures::TryStreamExt;
use futures_old::Poll;
use futures_old::Stream as OldStream;
use mercurial_bundles::wirepack::converter::WirePackConverter;
use mercurial_bundles::wirepack::converter::WirePackPartProcessor;
use mercurial_bundles::wirepack::DataEntry;
use mercurial_bundles::wirepack::HistoryEntry;
use mercurial_bundles::wirepack::Part;
use mercurial_revlog::manifest::ManifestContent;
use mercurial_types::delta;
use mercurial_types::HgNodeHash;
use mercurial_types::HgNodeKey;
use mercurial_types::RepoPath;
use mercurial_types::NULL_HASH;
use thiserror::Error;

#[derive(Debug, Error)]
pub enum ErrorKind {
    #[error("Malformed treemanifest part: {0}")]
    MalformedTreemanifestPart(String),
}

/// Parser for wirepack tree part. It returns a stream of TreemanifestEntry, that can be used by
/// Mononoke's Commit Api.
///
/// It assumes a few things:
/// 1) all data is sent as a delta from the null revision (i.e. data is basically non-deltaed).
/// 2) there are exactly one history entry and exactly one data entry for each tree.
pub struct TreemanifestBundle2Parser {
    stream: WirePackConverter<Compat<BoxStream<'static, Result<Part>>>, TreemanifestPartProcessor>,
}

impl TreemanifestBundle2Parser {
    pub fn new(part_stream: BoxStream<'static, Result<Part>>) -> Self {
        Self {
            stream: WirePackConverter::new(part_stream.compat(), TreemanifestPartProcessor::new()),
        }
    }
}

impl OldStream for TreemanifestBundle2Parser {
    type Item = TreemanifestEntry;
    type Error = Error;

    fn poll(&mut self) -> Poll<Option<Self::Item>, Error> {
        self.stream.poll()
    }
}

#[derive(Debug, Eq, PartialEq)]
pub struct TreemanifestEntry {
    pub node_key: HgNodeKey,
    pub data: Bytes,
    pub p1: Option<HgNodeHash>,
    pub p2: Option<HgNodeHash>,
    pub manifest_content: ManifestContent,
}

impl TreemanifestEntry {
    fn new(node_key: HgNodeKey, data: Bytes, p1: HgNodeHash, p2: HgNodeHash) -> Result<Self> {
        let manifest_content = ManifestContent::parse(data.as_ref())?;

        Ok(Self {
            node_key,
            data,
            p1: p1.into_option(),
            p2: p2.into_option(),
            manifest_content,
        })
    }
}

struct TreemanifestPartProcessor {
    node: Option<HgNodeHash>,
    p1: Option<HgNodeHash>,
    p2: Option<HgNodeHash>,
    path: Option<RepoPath>,
}

impl TreemanifestPartProcessor {
    fn new() -> Self {
        Self {
            node: None,
            p1: None,
            p2: None,
            path: None,
        }
    }
}

impl WirePackPartProcessor for TreemanifestPartProcessor {
    type Data = TreemanifestEntry;

    fn history_meta(&mut self, path: &RepoPath, entry_count: u32) -> Result<Option<Self::Data>> {
        replace_or_fail_if_exists(&mut self.path, path.clone())?;
        if entry_count != 1 {
            let msg = format!("expected exactly one history entry, got: {}", entry_count);
            return Err(ErrorKind::MalformedTreemanifestPart(msg).into());
        }
        Ok(None)
    }

    fn history(&mut self, entry: &HistoryEntry) -> Result<Option<Self::Data>> {
        replace_or_fail_if_exists(&mut self.node, entry.node.clone())?;
        replace_or_fail_if_exists(&mut self.p1, entry.p1.clone())?;
        replace_or_fail_if_exists(&mut self.p2, entry.p2.clone())?;
        Ok(None)
    }

    fn data_meta(&mut self, path: &RepoPath, entry_count: u32) -> Result<Option<Self::Data>> {
        if Some(path) != self.path.as_ref() {
            let msg = format!("unexpected path: {:?} != {:?}", path, self.path);
            Err(ErrorKind::MalformedTreemanifestPart(msg).into())
        } else if entry_count != 1 {
            let msg = format!("expected exactly one data entry, got: {}", entry_count);
            Err(ErrorKind::MalformedTreemanifestPart(msg).into())
        } else {
            Ok(None)
        }
    }

    fn data(&mut self, data_entry: &DataEntry) -> Result<Option<Self::Data>> {
        if data_entry.delta_base != NULL_HASH {
            let msg = format!("unexpected delta base: {:?}", data_entry.delta_base);
            return Err(ErrorKind::MalformedTreemanifestPart(msg).into());
        }

        let node_key = HgNodeKey {
            path: unwrap_field(&mut self.path, "path")?,
            hash: unwrap_field(&mut self.node, "node")?,
        };
        let bytes = Bytes::from(delta::apply("".as_bytes(), &data_entry.delta)?);
        let p1 = unwrap_field(&mut self.p1, "p1")?;
        let p2 = unwrap_field(&mut self.p2, "p2")?;

        Ok(Some(TreemanifestEntry::new(node_key, bytes, p1, p2)?))
    }

    fn end(&mut self) -> Result<Option<Self::Data>> {
        Ok(None)
    }
}

fn replace_or_fail_if_exists<T: Debug>(existing: &mut Option<T>, new_value: T) -> Result<()> {
    let existing = mem::replace(existing, Some(new_value));
    if existing.is_some() {
        let msg = format!("{:?} was already set", existing);
        Err(ErrorKind::MalformedTreemanifestPart(msg).into())
    } else {
        Ok(())
    }
}

fn unwrap_field<T: Clone>(field: &mut Option<T>, field_name: &str) -> Result<T> {
    field.take().ok_or_else(|| {
        let msg = format!("{} is not set", field_name);
        ErrorKind::MalformedTreemanifestPart(msg).into()
    })
}

#[cfg(test)]
mod test {
    use futures::stream;
    use futures::stream::StreamExt;
    use futures_old::Future;
    use maplit::btreemap;
    use mercurial_revlog::manifest::Details;
    use mercurial_types::manifest::Type;
    use mercurial_types::FileType;
    use mercurial_types::MPath;
    use mercurial_types_mocks::nodehash::*;

    use super::*;

    #[test]
    fn test_simple() {
        let parts = vec![
            get_history_meta(),
            get_history_entry(),
            get_data_meta(),
            get_data_entry(),
            get_history_meta(),
            get_history_entry(),
            get_data_meta(),
            get_data_entry(),
            Part::End,
        ];

        let part_stream = stream::iter(parts.into_iter().map(Ok)).boxed();
        let stream = TreemanifestBundle2Parser::new(part_stream);
        assert_eq!(
            stream.collect().wait().unwrap(),
            vec![get_expected_entry(), get_expected_entry()]
        );
    }

    #[test]
    fn test_broken() {
        let parts = vec![get_history_meta(), get_history_entry(), Part::End];
        assert_fails(parts);
        let parts = vec![
            get_history_meta(),
            get_history_entry(),
            get_data_meta(),
            Part::End,
        ];
        assert_fails(parts);
        let parts = vec![
            get_history_meta(),
            get_history_entry(),
            get_data_entry(),
            get_data_meta(),
            Part::End,
        ];
        assert_fails(parts);

        let parts = vec![
            get_history_meta(),
            get_history_entry(),
            Part::DataMeta {
                path: RepoPath::dir("dir").unwrap(),
                entry_count: 1,
            },
            get_data_entry(),
            Part::End,
        ];
        assert_fails(parts);
    }

    fn get_history_meta() -> Part {
        Part::HistoryMeta {
            path: RepoPath::root(),
            entry_count: 1,
        }
    }

    fn get_history_entry() -> Part {
        let node = ONES_HASH;
        let p1 = TWOS_HASH;
        let p2 = THREES_HASH;
        let linknode = FOURS_HASH;

        Part::History(HistoryEntry {
            node,
            p1,
            p2,
            linknode,
            copy_from: None,
        })
    }

    fn get_data_meta() -> Part {
        Part::DataMeta {
            path: RepoPath::root(),
            entry_count: 1,
        }
    }

    fn get_revlog_manifest_content() -> ManifestContent {
        ManifestContent {
            files: btreemap! {
                MPath::new("test_dir/test_file").unwrap() =>
                Details::new(
                    ONES_HASH,
                    Type::File(FileType::Regular),
                ),
                MPath::new("test_dir2/test_manifest").unwrap() =>
                Details::new(
                    TWOS_HASH,
                    Type::Tree,
                ),
            },
        }
    }

    fn get_data_entry() -> Part {
        let node = ONES_HASH;

        let data = {
            let mut data = Vec::new();
            get_revlog_manifest_content().generate(&mut data).unwrap();
            data
        };

        Part::Data(DataEntry {
            node,
            delta_base: NULL_HASH,
            delta: delta::Delta::new_fulltext(data),
            metadata: None,
        })
    }

    fn assert_fails(parts: Vec<Part>) {
        let part_stream = stream::iter(parts.into_iter().map(Ok)).boxed();
        let stream = TreemanifestBundle2Parser::new(part_stream);
        assert!(stream.collect().wait().is_err());
    }

    fn get_expected_entry() -> TreemanifestEntry {
        let node_key = HgNodeKey {
            path: RepoPath::root(),
            hash: ONES_HASH,
        };
        let p1 = TWOS_HASH;
        let p2 = THREES_HASH;

        let data = {
            let mut data = Vec::new();
            get_revlog_manifest_content().generate(&mut data).unwrap();
            data
        };

        let entry = TreemanifestEntry::new(node_key, Bytes::from(data), p1, p2).unwrap();

        assert_eq!(
            entry.manifest_content,
            get_revlog_manifest_content(),
            "Sanity check for manifest content failed"
        );

        entry
    }
}