/*
 * 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::cmp;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::hash::Hash;
use std::marker::PhantomData;
use std::ops::Add;
use std::sync::atomic::AtomicU32;
use std::sync::atomic::AtomicUsize;
use std::sync::atomic::Ordering;

use ahash::RandomState;
use anyhow::bail;
use anyhow::Error;
use array_init::array_init;
use async_trait::async_trait;
use bonsai_hg_mapping::BonsaiHgMapping;
use bonsai_hg_mapping::BonsaiHgMappingEntry;
use bulkops::Direction;
use context::CoreContext;
use dashmap::mapref::one::Ref;
use dashmap::DashMap;
use futures::future::TryFutureExt;
use itertools::Itertools;
use mercurial_types::HgChangesetId;
use mercurial_types::HgFileNodeId;
use mercurial_types::HgManifestId;
use mononoke_types::BasenameSuffixSkeletonManifestId;
use mononoke_types::ChangesetId;
use mononoke_types::ContentId;
use mononoke_types::DeletedManifestV2Id;
use mononoke_types::FastlogBatchId;
use mononoke_types::FileUnodeId;
use mononoke_types::FsnodeId;
use mononoke_types::ManifestUnodeId;
use mononoke_types::SkeletonManifestId;
use phases::Phase;
use phases::Phases;
use slog::info;
use slog::Logger;
use strum::EnumCount;
use strum::EnumIter;
use strum::EnumString;
use strum::EnumVariantNames;

use crate::detail::graph::EdgeType;
use crate::detail::graph::Node;
use crate::detail::graph::NodeData;
use crate::detail::graph::NodeType;
use crate::detail::graph::UnodeFlags;
use crate::detail::graph::WrappedPath;
use crate::detail::graph::WrappedPathHash;
use crate::detail::log;
use crate::detail::progress::sort_by_string;
use crate::detail::walk::expand_checked_nodes;
use crate::detail::walk::EmptyRoute;
use crate::detail::walk::OutgoingEdge;
use crate::detail::walk::TailingWalkVisitor;
use crate::detail::walk::VisitOne;
use crate::detail::walk::WalkVisitor;

#[derive(Clone, Copy, Default, Debug, PartialEq)]
pub struct StepStats {
    pub error_count: usize,
    pub missing_count: usize,
    pub hash_validation_failure_count: usize,
    pub num_expanded_new: usize,
    pub visited_of_type: usize,
}

impl Add<StepStats> for StepStats {
    type Output = Self;
    fn add(self, other: Self) -> Self {
        Self {
            error_count: self.error_count + other.error_count,
            missing_count: self.missing_count + other.missing_count,
            hash_validation_failure_count: self.hash_validation_failure_count
                + other.hash_validation_failure_count,
            num_expanded_new: self.num_expanded_new + other.num_expanded_new,
            visited_of_type: cmp::max(self.visited_of_type, other.visited_of_type),
        }
    }
}

// So we could change the type later without too much code churn
type InternId = u32;

// Common trait for the interned ids
trait Interned: Clone + Copy + Eq + Hash + PartialEq {
    fn new(id: InternId) -> Self;
}

#[derive(Eq, Hash, PartialEq)]
struct InternedId<K> {
    id: InternId,
    _phantom: PhantomData<K>,
}

// Can't auto-derive as dont want to make K Copy
impl<K> Copy for InternedId<K> {}
impl<K> Clone for InternedId<K> {
    fn clone(&self) -> Self {
        Self {
            id: self.id,
            _phantom: PhantomData,
        }
    }
}

impl<K> Interned for InternedId<K>
where
    K: Eq + Hash + PartialEq,
{
    fn new(id: InternId) -> Self {
        Self {
            id,
            _phantom: PhantomData,
        }
    }
}

struct InternMap<K, V> {
    interned: DashMap<K, V, RandomState>,
    next_id: AtomicU32,
}

impl<K, V> InternMap<K, V>
where
    K: Eq + Hash + Copy + fmt::Debug,
    V: Interned,
{
    fn with_hasher(fac: RandomState) -> Self {
        Self {
            interned: DashMap::with_hasher(fac),
            next_id: AtomicU32::new(0),
        }
    }

    // Intern the key if not already present, returns interned value.
    //
    // From `DashMap::entry()` documentation:
    // **Locking behaviour:** May deadlock if called when holding any sort of reference into the map.
    fn interned(&self, k: &K) -> V {
        // First try the read lock for the interned id, making sure we give up the read guard
        let id: Option<V> = self.interned.get(k).map(|id| *id);
        // Read guard released, escalated to write lock if necessary
        id.unwrap_or_else(|| {
            *self
                .interned
                .entry(*k)
                .or_insert_with(|| {
                    let id = self.next_id.fetch_add(1, Ordering::Release);
                    if id == InternId::MAX {
                        panic!("Intern counter wrapped around for {:?}", k);
                    }
                    V::new(id)
                })
                .value()
        })
    }

    // Get a immutable reference to an entry in the map
    //
    // From `DashMap::get()` documentation:
    // **Locking behaviour:** May deadlock if called when holding a mutable reference into the map.
    fn get(&self, k: &K) -> Option<Ref<K, V, RandomState>> {
        self.interned.get(k)
    }

    fn clear(&self) {
        self.next_id.store(1, Ordering::SeqCst);
        self.interned.clear()
    }
}

type ValueMap<K, V> = DashMap<K, V, RandomState>;

type StateMap<K> = DashMap<K, (), RandomState>;
#[derive(Clone, Copy, PartialEq, Eq, Hash)]
struct UnodeInterned<T> {
    id: InternedId<T>,
    flags: UnodeFlags,
}

pub struct WalkState {
    // Params
    include_node_types: HashSet<NodeType>,
    include_edge_types: HashSet<EdgeType>,
    always_emit_edge_types: HashSet<EdgeType>,
    enable_derive: bool,
    chunk_direction: Option<Direction>,
    // Interning
    bcs_ids: InternMap<ChangesetId, InternedId<ChangesetId>>,
    hg_cs_ids: InternMap<HgChangesetId, InternedId<HgChangesetId>>,
    hg_filenode_ids: InternMap<HgFileNodeId, InternedId<HgFileNodeId>>,
    path_hashes: InternMap<WrappedPathHash, InternedId<WrappedPathHash>>,
    hg_manifest_ids: InternMap<HgManifestId, InternedId<HgManifestId>>,
    unode_file_ids: InternMap<FileUnodeId, InternedId<FileUnodeId>>,
    unode_manifest_ids: InternMap<ManifestUnodeId, InternedId<ManifestUnodeId>>,
    // State
    chunk_bcs: StateMap<InternedId<ChangesetId>>,
    deferred_bcs: ValueMap<InternedId<ChangesetId>, HashSet<OutgoingEdge>>,
    bcs_to_hg: ValueMap<InternedId<ChangesetId>, HgChangesetId>,
    hg_to_bcs: ValueMap<InternedId<HgChangesetId>, ChangesetId>,
    visited_bcs: StateMap<InternedId<ChangesetId>>,
    visited_bcs_mapping: StateMap<InternedId<ChangesetId>>,
    public_not_visited: StateMap<InternedId<ChangesetId>>,
    visited_bcs_phase: StateMap<InternedId<ChangesetId>>,
    visited_file: StateMap<ContentId>,
    visited_hg_cs: StateMap<InternedId<HgChangesetId>>,
    visited_hg_cs_mapping: StateMap<InternedId<HgChangesetId>>,
    visited_hg_cs_via_bonsai: StateMap<InternedId<HgChangesetId>>,
    visited_hg_file_envelope: StateMap<InternedId<HgFileNodeId>>,
    visited_hg_filenode: StateMap<(InternedId<WrappedPathHash>, InternedId<HgFileNodeId>)>,
    visited_hg_manifest_filenode: StateMap<(InternedId<WrappedPathHash>, InternedId<HgFileNodeId>)>,
    visited_hg_manifest: StateMap<(InternedId<WrappedPathHash>, InternedId<HgManifestId>)>,
    // Derived
    visited_blame: StateMap<InternedId<FileUnodeId>>,
    visited_changeset_info: StateMap<InternedId<ChangesetId>>,
    visited_changeset_info_mapping: StateMap<InternedId<ChangesetId>>,
    visited_deleted_manifest_v2: StateMap<DeletedManifestV2Id>,
    visited_deleted_manifest_v2_mapping: StateMap<InternedId<ChangesetId>>,
    visited_fastlog_batch: StateMap<FastlogBatchId>,
    visited_fastlog_dir: StateMap<InternedId<ManifestUnodeId>>,
    visited_fastlog_file: StateMap<InternedId<FileUnodeId>>,
    visited_fsnode: StateMap<FsnodeId>,
    visited_fsnode_mapping: StateMap<InternedId<ChangesetId>>,
    visited_skeleton_manifest: StateMap<SkeletonManifestId>,
    visited_skeleton_manifest_mapping: StateMap<InternedId<ChangesetId>>,
    visited_basename_suffix_skeleton_manifest: StateMap<BasenameSuffixSkeletonManifestId>,
    visited_basename_suffix_skeleton_manifest_mapping: StateMap<InternedId<ChangesetId>>,
    visited_unode_file: StateMap<UnodeInterned<FileUnodeId>>,
    visited_unode_manifest: StateMap<UnodeInterned<ManifestUnodeId>>,
    visited_unode_mapping: StateMap<InternedId<ChangesetId>>,
    // Count
    visit_count: [AtomicUsize; NodeType::COUNT],
}

impl WalkState {
    pub fn new(
        include_node_types: HashSet<NodeType>,
        include_edge_types: HashSet<EdgeType>,
        always_emit_edge_types: HashSet<EdgeType>,
        enable_derive: bool,
        chunk_direction: Option<Direction>,
    ) -> Self {
        let fac = RandomState::default();
        Self {
            // Params
            include_node_types,
            include_edge_types,
            always_emit_edge_types,
            enable_derive,
            chunk_direction,
            // Interning
            bcs_ids: InternMap::with_hasher(fac.clone()),
            hg_cs_ids: InternMap::with_hasher(fac.clone()),
            hg_filenode_ids: InternMap::with_hasher(fac.clone()),
            path_hashes: InternMap::with_hasher(fac.clone()),
            hg_manifest_ids: InternMap::with_hasher(fac.clone()),
            unode_file_ids: InternMap::with_hasher(fac.clone()),
            unode_manifest_ids: InternMap::with_hasher(fac.clone()),
            // State
            chunk_bcs: StateMap::with_hasher(fac.clone()),
            deferred_bcs: ValueMap::with_hasher(fac.clone()),
            bcs_to_hg: ValueMap::with_hasher(fac.clone()),
            hg_to_bcs: ValueMap::with_hasher(fac.clone()),
            visited_bcs: StateMap::with_hasher(fac.clone()),
            visited_bcs_mapping: StateMap::with_hasher(fac.clone()),
            public_not_visited: StateMap::with_hasher(fac.clone()),
            visited_bcs_phase: StateMap::with_hasher(fac.clone()),
            visited_file: StateMap::with_hasher(fac.clone()),
            visited_hg_cs: StateMap::with_hasher(fac.clone()),
            visited_hg_cs_mapping: StateMap::with_hasher(fac.clone()),
            visited_hg_cs_via_bonsai: StateMap::with_hasher(fac.clone()),
            visited_hg_file_envelope: StateMap::with_hasher(fac.clone()),
            visited_hg_filenode: StateMap::with_hasher(fac.clone()),
            visited_hg_manifest_filenode: StateMap::with_hasher(fac.clone()),
            visited_hg_manifest: StateMap::with_hasher(fac.clone()),
            // Derived
            visited_blame: StateMap::with_hasher(fac.clone()),
            visited_changeset_info: StateMap::with_hasher(fac.clone()),
            visited_changeset_info_mapping: StateMap::with_hasher(fac.clone()),
            visited_deleted_manifest_v2: StateMap::with_hasher(fac.clone()),
            visited_deleted_manifest_v2_mapping: StateMap::with_hasher(fac.clone()),
            visited_fastlog_batch: StateMap::with_hasher(fac.clone()),
            visited_fastlog_dir: StateMap::with_hasher(fac.clone()),
            visited_fastlog_file: StateMap::with_hasher(fac.clone()),
            visited_fsnode: StateMap::with_hasher(fac.clone()),
            visited_fsnode_mapping: StateMap::with_hasher(fac.clone()),
            visited_skeleton_manifest: StateMap::with_hasher(fac.clone()),
            visited_skeleton_manifest_mapping: StateMap::with_hasher(fac.clone()),
            visited_basename_suffix_skeleton_manifest: StateMap::with_hasher(fac.clone()),
            visited_basename_suffix_skeleton_manifest_mapping: StateMap::with_hasher(fac.clone()),
            visited_unode_file: StateMap::with_hasher(fac.clone()),
            visited_unode_manifest: StateMap::with_hasher(fac.clone()),
            visited_unode_mapping: StateMap::with_hasher(fac),
            // Count
            visit_count: array_init(|_i| AtomicUsize::new(0)),
        }
    }

    fn record<K>(&self, visited: &StateMap<K>, k: &K) -> bool
    where
        K: Eq + Hash + Clone,
    {
        if visited.contains_key(k) {
            false
        } else {
            visited.insert(k.clone(), ()).is_none()
        }
    }

    fn record_multi<K, V>(&self, multi_map: &ValueMap<K, HashSet<V>>, k: K, v: &V) -> bool
    where
        K: Eq + Hash + Clone,
        V: Eq + Hash + Clone,
    {
        let mut entry = multi_map.entry(k).or_insert_with(HashSet::default);
        let values = entry.value_mut();
        // No insert_with in HashSet, so do it ourselves
        if values.contains(v) {
            false
        } else {
            values.insert(v.clone())
        }
    }

    /// If the state did not have this value present, true is returned.
    fn record_with_path<K>(
        &self,
        visited_with_path: &StateMap<(InternedId<WrappedPathHash>, K)>,
        k: (&WrappedPath, &K),
    ) -> bool
    where
        K: Eq + Hash + Copy,
    {
        let (path, id) = k;
        let path = self.path_hashes.interned(path.get_path_hash());
        let key = (path, *id);
        if visited_with_path.contains_key(&key) {
            false
        } else {
            visited_with_path.insert(key, ()).is_none()
        }
    }

    fn record_resolved_visit(&self, resolved: &OutgoingEdge, node_data: Option<&NodeData>) {
        match (&resolved.target, node_data) {
            // Bonsai
            (Node::PhaseMapping(bcs_id), Some(NodeData::PhaseMapping(Some(Phase::Public)))) => {
                let id = &self.bcs_ids.interned(bcs_id);
                // Only retain visit if already public, otherwise it could mutate between walks.
                self.record(&self.visited_bcs_phase, id);
                // Save some memory, no need to keep an entry in public_not_visited now its in visited_bcs_phase
                self.public_not_visited.remove(id);
            }
            // Hg
            (Node::BonsaiHgMapping(k), Some(_)) => {
                self.record(&self.visited_bcs_mapping, &self.bcs_ids.interned(&k.inner));
            }
            // Derived
            (Node::ChangesetInfoMapping(bcs_id), Some(_)) => {
                self.record(
                    &self.visited_changeset_info_mapping,
                    &self.bcs_ids.interned(bcs_id),
                );
            }
            (Node::DeletedManifestV2Mapping(bcs_id), Some(_)) => {
                self.record(
                    &self.visited_deleted_manifest_v2_mapping,
                    &self.bcs_ids.interned(bcs_id),
                );
            }
            (Node::FsnodeMapping(bcs_id), Some(_)) => {
                self.record(&self.visited_fsnode_mapping, &self.bcs_ids.interned(bcs_id));
            }
            (Node::SkeletonManifestMapping(bcs_id), Some(_)) => {
                self.record(
                    &self.visited_skeleton_manifest_mapping,
                    &self.bcs_ids.interned(bcs_id),
                );
            }
            (Node::BasenameSuffixSkeletonManifestMapping(bcs_id), Some(_)) => {
                self.record(
                    &self.visited_basename_suffix_skeleton_manifest_mapping,
                    &self.bcs_ids.interned(bcs_id),
                );
            }
            (Node::UnodeMapping(bcs_id), Some(_)) => {
                self.record(&self.visited_unode_mapping, &self.bcs_ids.interned(bcs_id));
            }
            _ => {}
        }
    }

    fn retain_edge(&self, outgoing_edge: &OutgoingEdge) -> bool {
        // Retain if a root, or if selected
        outgoing_edge.label.incoming_type().is_none()
            || (self
                .include_node_types
                .contains(&outgoing_edge.target.get_type())
                && self.include_edge_types.contains(&outgoing_edge.label))
    }

    fn get_visit_count(&self, t: &NodeType) -> usize {
        self.visit_count[*t as usize].load(Ordering::Acquire)
    }

    fn chunk_contains(&self, id: InternedId<ChangesetId>) -> bool {
        if self.chunk_bcs.is_empty() {
            true
        } else {
            self.chunk_bcs.contains_key(&id)
        }
    }

    // Clear the the InternedId to X mappings, which also means we must clear all mappings using those mappings.
    // InternedId values will start again from 0,  so we don't want old values to clash.
    //
    // The match clears mappings containing InternedId<Foo> for InternedType::Foo. The mapping was built by
    // looking at the type definition, seeing what other fields reference the InternedId<Foo> provided by the InternMap,
    // and then clearing those too.
    fn clear_interned(&mut self, interned_type: InternedType) {
        match interned_type {
            InternedType::FileUnodeId => {
                self.unode_file_ids.clear();
                self.clear_mapping(NodeType::Blame);
                self.clear_mapping(NodeType::FastlogFile);
                self.clear_mapping(NodeType::UnodeFile);
            }
            InternedType::HgChangesetId => {
                self.hg_cs_ids.clear();
                self.hg_to_bcs.clear();
                self.bcs_to_hg.clear();
                self.clear_mapping(NodeType::HgChangeset);
                self.clear_mapping(NodeType::HgBonsaiMapping);
                self.clear_mapping(NodeType::HgChangesetViaBonsai);
            }
            InternedType::HgFileNodeId => {
                self.hg_filenode_ids.clear();
                self.clear_mapping(NodeType::HgFileEnvelope);
                self.clear_mapping(NodeType::HgFileNode);
            }
            InternedType::HgManifestId => {
                self.hg_manifest_ids.clear();
                self.clear_mapping(NodeType::HgManifest);
            }
            InternedType::ManifestUnodeId => {
                self.unode_manifest_ids.clear();
                self.clear_mapping(NodeType::FastlogDir);
                self.clear_mapping(NodeType::UnodeManifest);
            }
            InternedType::MPathHash => {
                self.path_hashes.clear();
                self.clear_mapping(NodeType::HgFileNode);
                self.clear_mapping(NodeType::HgManifest);
            }
        }
    }

    fn clear_mapping(&mut self, node_type: NodeType) {
        match node_type {
            // Entry points
            NodeType::Root => {}
            NodeType::Bookmark => {}
            NodeType::PublishedBookmarks => {}
            // Bonsai
            NodeType::Changeset => self.visited_bcs.clear(),
            NodeType::BonsaiHgMapping => self.visited_bcs_mapping.clear(),
            NodeType::PhaseMapping => self.visited_bcs_phase.clear(),
            // Hg
            NodeType::HgBonsaiMapping => self.visited_hg_cs_mapping.clear(),
            NodeType::HgChangeset => self.visited_hg_cs.clear(),
            NodeType::HgChangesetViaBonsai => self.visited_hg_cs_via_bonsai.clear(),
            NodeType::HgManifest => self.visited_hg_manifest.clear(),
            NodeType::HgFileNode => self.visited_hg_filenode.clear(),
            NodeType::HgManifestFileNode => self.visited_hg_manifest_filenode.clear(),
            NodeType::HgFileEnvelope => self.visited_hg_file_envelope.clear(),
            // Content
            NodeType::FileContent => self.visited_file.clear(),
            NodeType::FileContentMetadataV2 => {} // reached via expand_checked_nodes
            NodeType::AliasContentMapping => {}   // reached via expand_checked_nodes
            // Derived
            NodeType::Blame => self.visited_blame.clear(),
            NodeType::ChangesetInfo => self.visited_changeset_info.clear(),
            NodeType::ChangesetInfoMapping => self.visited_changeset_info_mapping.clear(),
            NodeType::DeletedManifestV2 => self.visited_deleted_manifest_v2.clear(),
            NodeType::DeletedManifestV2Mapping => self.visited_deleted_manifest_v2_mapping.clear(),
            NodeType::FastlogBatch => self.visited_fastlog_batch.clear(),
            NodeType::FastlogDir => self.visited_fastlog_dir.clear(),
            NodeType::FastlogFile => self.visited_fastlog_file.clear(),
            NodeType::Fsnode => self.visited_fsnode.clear(),
            NodeType::FsnodeMapping => self.visited_fsnode_mapping.clear(),
            NodeType::SkeletonManifest => self.visited_skeleton_manifest.clear(),
            NodeType::SkeletonManifestMapping => self.visited_skeleton_manifest_mapping.clear(),
            NodeType::BasenameSuffixSkeletonManifest => {
                self.visited_basename_suffix_skeleton_manifest.clear()
            }
            NodeType::BasenameSuffixSkeletonManifestMapping => self
                .visited_basename_suffix_skeleton_manifest_mapping
                .clear(),
            NodeType::UnodeFile => self.visited_unode_file.clear(),
            NodeType::UnodeManifest => self.visited_unode_manifest.clear(),
            NodeType::UnodeMapping => self.visited_unode_mapping.clear(),
        }
    }

    fn needs_visit_impl(&self, outgoing: &OutgoingEdge, executing_step: bool) -> bool {
        let target_node: &Node = &outgoing.target;
        let k = target_node.get_type();
        if !executing_step {
            self.visit_count[k as usize].fetch_add(1, Ordering::Release);
        }

        // For types handled by record_resolved_visit logic is same when executing or checking a step
        // For types handled by record() and record_with_path, executing_step returns true.
        match (&target_node, executing_step) {
            // Entry points
            (Node::Root(_), _) => true,
            (Node::Bookmark(_), _) => true,
            (Node::PublishedBookmarks(_), _) => true,
            // Bonsai
            (Node::Changeset(_), true) => true,
            (Node::Changeset(k), false) => {
                let id = self.bcs_ids.interned(&k.inner);
                if self.chunk_contains(id) {
                    self.record(&self.visited_bcs, &id)
                } else {
                    if self.chunk_direction == Some(Direction::NewestFirst)
                        && !self.visited_bcs.contains_key(&id)
                    {
                        self.record_multi(&self.deferred_bcs, id, outgoing);
                    }
                    false
                }
            }
            (Node::BonsaiHgMapping(k), _) => {
                if let Some(id) = self.bcs_ids.get(&k.inner) {
                    // Does not insert, see record_resolved_visit
                    !self.visited_bcs_mapping.contains_key(&id)
                } else {
                    true
                }
            }
            (Node::PhaseMapping(bcs_id), _) => {
                if let Some(id) = self.bcs_ids.get(bcs_id) {
                    // Does not insert, as can only prune visits once data resolved, see record_resolved_visit
                    !self.visited_bcs_phase.contains_key(&id)
                } else {
                    true
                }
            }
            // Hg
            (Node::HgBonsaiMapping(_), true) => true,
            (Node::HgBonsaiMapping(k), false) => self.record(
                &self.visited_hg_cs_mapping,
                &self.hg_cs_ids.interned(&k.inner),
            ),
            (Node::HgChangeset(_), true) => true,
            (Node::HgChangeset(k), false) => {
                self.record(&self.visited_hg_cs, &self.hg_cs_ids.interned(&k.inner))
            }
            (Node::HgChangesetViaBonsai(_), true) => true,
            (Node::HgChangesetViaBonsai(k), false) => self.record(
                &self.visited_hg_cs_via_bonsai,
                &self.hg_cs_ids.interned(&k.inner),
            ),
            (Node::HgManifest(_), true) => true,
            (Node::HgManifest(k), false) => self.record_with_path(
                &self.visited_hg_manifest,
                (&k.path, &self.hg_manifest_ids.interned(&k.id)),
            ),
            (Node::HgFileNode(_), true) => true,
            (Node::HgFileNode(k), false) => self.record_with_path(
                &self.visited_hg_filenode,
                (&k.path, &self.hg_filenode_ids.interned(&k.id)),
            ),
            (Node::HgManifestFileNode(_), true) => true,
            (Node::HgManifestFileNode(k), false) => self.record_with_path(
                &self.visited_hg_manifest_filenode,
                (&k.path, &self.hg_filenode_ids.interned(&k.id)),
            ),
            (Node::HgFileEnvelope(_), true) => true,
            (Node::HgFileEnvelope(id), false) => self.record(
                &self.visited_hg_file_envelope,
                &self.hg_filenode_ids.interned(id),
            ),
            // Content
            (Node::FileContent(_), true) => true,
            (Node::FileContent(content_id), false) => self.record(&self.visited_file, content_id),
            (Node::FileContentMetadataV2(_), _) => true, // reached via expand_checked_nodes
            (Node::AliasContentMapping(_), _) => true,   // reached via expand_checked_nodes
            // Derived
            (Node::Blame(_), true) => true,
            (Node::Blame(k), false) => self.record(
                &self.visited_blame,
                &self.unode_file_ids.interned(k.as_ref()),
            ),
            (Node::ChangesetInfo(_), true) => true,
            (Node::ChangesetInfo(bcs_id), false) => {
                let id = self.bcs_ids.interned(bcs_id);
                if self.chunk_contains(id) {
                    self.record(&self.visited_changeset_info, &id)
                } else {
                    if self.chunk_direction == Some(Direction::NewestFirst)
                        && !self.visited_changeset_info.contains_key(&id)
                    {
                        self.record_multi(&self.deferred_bcs, id, outgoing);
                    }
                    false
                }
            }
            (Node::ChangesetInfoMapping(bcs_id), _) => {
                if let Some(id) = self.bcs_ids.get(bcs_id) {
                    !self.visited_changeset_info_mapping.contains_key(&id) // Does not insert, see record_resolved_visit
                } else {
                    true
                }
            }
            (Node::DeletedManifestV2(_), true) => true,
            (Node::DeletedManifestV2(id), false) => {
                self.record(&self.visited_deleted_manifest_v2, id)
            }
            (Node::DeletedManifestV2Mapping(bcs_id), _) => {
                if let Some(id) = self.bcs_ids.get(bcs_id) {
                    !self.visited_deleted_manifest_v2_mapping.contains_key(&id) // Does not insert, see record_resolved_visit
                } else {
                    true
                }
            }
            (Node::FastlogBatch(_), true) => true,
            (Node::FastlogBatch(k), false) => self.record(&self.visited_fastlog_batch, k),
            (Node::FastlogDir(_), true) => true,
            (Node::FastlogDir(k), false) => self.record(
                &self.visited_fastlog_dir,
                &self.unode_manifest_ids.interned(&k.inner),
            ),
            (Node::FastlogFile(_), true) => true,
            (Node::FastlogFile(k), false) => self.record(
                &self.visited_fastlog_file,
                &self.unode_file_ids.interned(&k.inner),
            ),
            (Node::Fsnode(_), true) => true,
            (Node::Fsnode(id), false) => self.record(&self.visited_fsnode, id),
            (Node::FsnodeMapping(bcs_id), _) => {
                if let Some(id) = self.bcs_ids.get(bcs_id) {
                    !self.visited_fsnode_mapping.contains_key(&id) // Does not insert, see record_resolved_visit
                } else {
                    true
                }
            }
            (Node::SkeletonManifest(_), true) => true,
            (Node::SkeletonManifest(id), false) => self.record(&self.visited_skeleton_manifest, id),
            (Node::SkeletonManifestMapping(bcs_id), _) => {
                if let Some(id) = self.bcs_ids.get(bcs_id) {
                    !self.visited_skeleton_manifest_mapping.contains_key(&id) // Does not insert, see record_resolved_visit
                } else {
                    true
                }
            }
            (Node::BasenameSuffixSkeletonManifest(_), true) => true,
            (Node::BasenameSuffixSkeletonManifest(id), false) => {
                self.record(&self.visited_basename_suffix_skeleton_manifest, id)
            }
            (Node::BasenameSuffixSkeletonManifestMapping(bcs_id), _) => {
                if let Some(id) = self.bcs_ids.get(bcs_id) {
                    !self
                        .visited_basename_suffix_skeleton_manifest_mapping
                        .contains_key(&id) // Does not insert, see record_resolved_visit
                } else {
                    true
                }
            }
            (Node::UnodeFile(_), true) => true,
            (Node::UnodeFile(k), false) => self.record(
                &self.visited_unode_file,
                &UnodeInterned {
                    id: self.unode_file_ids.interned(&k.inner),
                    flags: k.flags,
                },
            ),
            (Node::UnodeManifest(_), true) => true,
            (Node::UnodeManifest(k), false) => self.record(
                &self.visited_unode_manifest,
                &UnodeInterned {
                    id: self.unode_manifest_ids.interned(&k.inner),
                    flags: k.flags,
                },
            ),
            (Node::UnodeMapping(bcs_id), _) => {
                if let Some(id) = self.bcs_ids.get(bcs_id) {
                    !self.visited_unode_mapping.contains_key(&id) // Does not insert, see record_resolved_visit
                } else {
                    true
                }
            }
        }
    }
}

#[derive(
    Clone,
    Copy,
    Debug,
    PartialEq,
    Eq,
    Hash,
    EnumIter,
    EnumString,
    EnumVariantNames
)]
pub enum InternedType {
    // No ChangesetId as that is not flushable as it is used to maintain deferred_bcs
    FileUnodeId,
    HgChangesetId,
    HgFileNodeId,
    HgManifestId,
    ManifestUnodeId,
    MPathHash,
}

#[async_trait]
impl VisitOne for WalkState {
    fn in_chunk(&self, bcs_id: &ChangesetId) -> bool {
        if self.chunk_bcs.is_empty() {
            true
        } else {
            let id = self.bcs_ids.interned(bcs_id);
            self.chunk_bcs.contains_key(&id)
        }
    }

    fn get_hg_from_bonsai(&self, bcs_id: &ChangesetId) -> Option<HgChangesetId> {
        let bcs_int = self.bcs_ids.interned(bcs_id);
        self.bcs_to_hg.get(&bcs_int).map(|v| *v.value())
    }

    fn record_hg_from_bonsai(&self, bcs_id: &ChangesetId, hg_cs_id: HgChangesetId) {
        let bcs_int = self.bcs_ids.interned(bcs_id);
        self.bcs_to_hg.insert(bcs_int, hg_cs_id);
    }

    async fn get_bonsai_from_hg(
        &self,
        ctx: &CoreContext,
        bonsai_hg_mapping: &dyn BonsaiHgMapping,
        hg_cs_id: &HgChangesetId,
    ) -> Result<ChangesetId, Error> {
        let hg_int = self.hg_cs_ids.interned(hg_cs_id);
        let bcs_id = if let Some(bcs_id) = self.hg_to_bcs.get(&hg_int) {
            *bcs_id
        } else {
            let bcs_id = bonsai_hg_mapping
                .get_bonsai_from_hg(ctx, hg_cs_id.clone())
                .await?;
            if let Some(bcs_id) = bcs_id {
                let bcs_int = self.bcs_ids.interned(&bcs_id);
                self.hg_to_bcs.insert(hg_int, bcs_id);
                self.bcs_to_hg.insert(bcs_int, *hg_cs_id);
                bcs_id
            } else {
                bail!("Can't have hg without bonsai for {}", hg_cs_id);
            }
        };
        Ok(bcs_id)
    }

    async fn defer_from_hg(
        &self,
        ctx: &CoreContext,
        bonsai_hg_mapping: &dyn BonsaiHgMapping,
        hg_cs_id: &HgChangesetId,
    ) -> Result<Option<ChangesetId>, Error> {
        if self.chunk_bcs.is_empty() {
            return Ok(None);
        }
        let bcs_id = self
            .get_bonsai_from_hg(ctx, bonsai_hg_mapping, hg_cs_id)
            .await?;
        let id = self.bcs_ids.interned(&bcs_id);
        if self.chunk_bcs.contains_key(&id) {
            Ok(None)
        } else {
            Ok(Some(bcs_id))
        }
    }

    async fn is_public(
        &self,
        ctx: &CoreContext,
        phases_store: &dyn Phases,
        bcs_id: &ChangesetId,
    ) -> Result<bool, Error> {
        // Short circuit if we already know its public
        if let Some(id) = self.bcs_ids.get(bcs_id) {
            if self.visited_bcs_phase.contains_key(&id) || self.public_not_visited.contains_key(&id)
            {
                return Ok(true);
            }
        }

        let public_not_visited = &self.public_not_visited;
        let id = self.bcs_ids.interned(bcs_id);

        let is_public = phases_store
            .get_public(
                ctx,
                vec![*bcs_id],
                !self.enable_derive, /* emphemeral_derive */
            )
            .map_ok(move |public| public.contains(bcs_id))
            .await?;

        // Only record visit in public_not_visited if it is public, as state can't change from that point
        // NB, this puts it in public_not_visited rather than visited_bcs_phase so that we still emit a Phase
        // entry from the stream
        if is_public {
            public_not_visited.insert(id, ());
        }
        Ok(is_public)
    }

    /// If the set did not have this value present, true is returned.
    fn needs_visit(&self, outgoing: &OutgoingEdge) -> bool {
        self.needs_visit_impl(outgoing, false)
    }
}

impl TailingWalkVisitor for WalkState {
    fn start_chunk(
        &mut self,
        new_chunk_bcs: &HashSet<ChangesetId>,
        mapping_prepop: Vec<BonsaiHgMappingEntry>,
    ) -> Result<HashSet<OutgoingEdge>, Error> {
        // Reset self.chunk_bcs
        let mut chunk_interned = HashSet::new();
        for bcs_id in new_chunk_bcs {
            let i = self.bcs_ids.interned(bcs_id);
            chunk_interned.insert(i);
            self.chunk_bcs.insert(i, ());
        }
        self.chunk_bcs.retain(|k, _v| chunk_interned.contains(k));

        // Check for items that were outside the chunk now being inside
        let mut in_new_chunk = HashSet::new();
        for e in self.deferred_bcs.iter() {
            if !chunk_interned.contains(e.key()) {
                continue;
            }
            in_new_chunk.extend(e.value().iter().cloned());
        }
        self.deferred_bcs
            .retain(|k, _v| !chunk_interned.contains(k));

        for i in mapping_prepop {
            let bcs_int = self.bcs_ids.interned(&i.bcs_id);
            let hg_int = self.hg_cs_ids.interned(&i.hg_cs_id);
            self.bcs_to_hg.insert(bcs_int, i.hg_cs_id);
            self.hg_to_bcs.insert(hg_int, i.bcs_id);
        }

        Ok(in_new_chunk)
    }

    fn clear_state(
        &mut self,
        node_types: &HashSet<NodeType>,
        interned_types: &HashSet<InternedType>,
    ) {
        node_types.iter().for_each(|t| self.clear_mapping(*t));
        interned_types.iter().for_each(|t| self.clear_interned(*t));
    }

    fn end_chunks(&mut self, logger: &Logger, contiguous_bounds: bool) -> Result<(), Error> {
        if !self.deferred_bcs.is_empty() {
            let summary: HashMap<EdgeType, usize> = self
                .deferred_bcs
                .iter()
                .flat_map(|e| e.value().clone())
                .group_by(|e| e.label)
                .into_iter()
                .map(|(key, group)| (key, group.count()))
                .collect();

            let summary_msg: String = sort_by_string(summary.keys())
                .iter()
                .map(|k| {
                    let mut s = k.to_string();
                    s.push(':');
                    s.push_str(&summary.get(k).map_or("".to_string(), |v| v.to_string()));
                    s
                })
                .join(" ");

            // Where we load from checkpoints the chunks may not be contiguous,
            // which means that some deferred edges can be covered in the checkpointed
            // section we are not repeating.
            if contiguous_bounds {
                let mut count = 0;
                bail!(
                    "Unexpected remaining edges to walk {}, sample of remaining: {:?}",
                    summary_msg,
                    self.deferred_bcs
                        .iter()
                        .take_while(|_| {
                            count += 1;
                            count < 50
                        })
                        .map(|e| e.value().clone())
                        .collect::<Vec<_>>()
                );
            } else {
                info!(logger, #log::CHUNKING, "Deferred edge counts by type were: {}", summary_msg);
                // Deferrals are only between chunks, clear if all chunks done.
                self.deferred_bcs.clear();
            }
        }
        Ok(())
    }

    fn num_deferred(&self) -> usize {
        self.deferred_bcs.len()
    }
}

impl WalkVisitor<(Node, Option<NodeData>, Option<StepStats>), EmptyRoute> for WalkState {
    fn start_step(
        &self,
        ctx: CoreContext,
        route: Option<&EmptyRoute>,
        step: &OutgoingEdge,
    ) -> Option<CoreContext> {
        if route.is_none() // is it a root
            || step.label.incoming_type().is_none() // is it from a root?
            || self.always_emit_edge_types.contains(&step.label) // always emit?
            || self.needs_visit_impl(step, true)
        {
            Some(ctx)
        } else {
            None
        }
    }

    fn visit(
        &self,
        _ctx: &CoreContext,
        resolved: OutgoingEdge,
        node_data: Option<NodeData>,
        route: Option<EmptyRoute>,
        mut outgoing: Vec<OutgoingEdge>,
    ) -> (
        (Node, Option<NodeData>, Option<StepStats>),
        EmptyRoute,
        Vec<OutgoingEdge>,
    ) {
        let queued_roots = if resolved.label.incoming_type().is_none() {
            1
        } else {
            0
        };
        if route.is_none() || !self.always_emit_edge_types.is_empty() {
            outgoing.retain(|e| {
                if e.label.incoming_type().is_none() {
                    // Make sure stats are updated for root nodes
                    self.needs_visit(e);
                    true
                } else {
                    // Check the always emit edges, outer visitor has now processed them.
                    self.retain_edge(e)
                        && (!self.always_emit_edge_types.contains(&e.label) || self.needs_visit(e))
                }
            });
        }

        let num_outgoing = outgoing.len();
        expand_checked_nodes(&mut outgoing);

        // Make sure we don't expand to types of node and edge not wanted
        if num_outgoing != outgoing.len() {
            outgoing.retain(|e| self.retain_edge(e));
        }

        self.record_resolved_visit(&resolved, node_data.as_ref());

        // Stats
        let num_expanded_new = outgoing.len() + queued_roots;
        let node = resolved.target;

        let mut stats = StepStats {
            error_count: 0,
            missing_count: 0,
            hash_validation_failure_count: 0,
            num_expanded_new,
            visited_of_type: self.get_visit_count(&node.get_type()),
        };
        let node_data = match node_data {
            Some(NodeData::ErrorAsData(_key)) => {
                stats.error_count += 1;
                None
            }
            Some(NodeData::MissingAsData(_key)) => {
                stats.missing_count += 1;
                None
            }
            Some(NodeData::HashValidationFailureAsData(_key)) => {
                stats.hash_validation_failure_count += 1;
                None
            }
            Some(d) => Some(d),
            None => None,
        };

        ((node, node_data, Some(stats)), EmptyRoute {}, outgoing)
    }

    fn defer_visit(
        &self,
        bcs_id: &ChangesetId,
        walk_item: &OutgoingEdge,
        _route: Option<EmptyRoute>,
    ) -> Result<((Node, Option<NodeData>, Option<StepStats>), EmptyRoute), Error> {
        let node_data = match self.chunk_direction {
            Some(Direction::NewestFirst) => {
                let i = self.bcs_ids.interned(bcs_id);
                self.record_multi(&self.deferred_bcs, i, walk_item);
                None
            }
            // We'll never visit backward looking edges when running OldestFirst, so don't record them.
            // returning Some for NodeData tells record_resolved_visit that we don't need to visit this node again if we see it.
            Some(Direction::OldestFirst) => Some(NodeData::OutsideChunk),
            None => bail!(
                "Attempt to defer {:?} step {:?} when not chunking",
                bcs_id,
                walk_item
            ),
        };
        Ok(((walk_item.target.clone(), node_data, None), EmptyRoute {}))
    }
}

#[cfg(test)]
mod tests {
    use std::mem::size_of;

    use super::*;

    #[test]
    fn test_interned_size() {
        // InternedId size is important as we have a lot of them, so test in case it changes
        assert_eq!(4, size_of::<InternedId<ChangesetId>>());
    }
}


/*
 * 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::cmp::max;
use std::cmp::min;
use std::collections::HashSet;
use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Error;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use bonsai_hg_mapping::BonsaiOrHgChangesetIds;
use bulkops::Direction;
use bulkops::PublicChangesetBulkFetch;
use bulkops::MAX_FETCH_STEP;
use changesets::ChangesetsArc;
use cloned::cloned;
use context::CoreContext;
use derived_data_manager::BonsaiDerivable as NewBonsaiDerivable;
use fbinit::FacebookInit;
use filenodes_derivation::FilenodesOnlyPublic;
use futures::future;
use futures::future::Future;
use futures::stream;
use futures::stream::BoxStream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use mercurial_derivation::MappedHgChangesetId;
use mononoke_types::ChangesetId;
use mononoke_types::RepositoryId;
use mononoke_types::Timestamp;
use phases::PhasesArc;
use repo_identity::RepoIdentityRef;
use slog::info;
use slog::Logger;
use strum::IntoEnumIterator;
use tokio::time::Duration;
use tokio::time::Instant;

use crate::commands::JobWalkParams;
use crate::detail::checkpoint::Checkpoint;
use crate::detail::checkpoint::CheckpointsByName;
use crate::detail::graph::ChangesetKey;
use crate::detail::graph::Node;
use crate::detail::graph::NodeType;
use crate::detail::log;
use crate::detail::state::InternedType;
use crate::detail::walk::walk_exact;
use crate::detail::walk::OutgoingEdge;
use crate::detail::walk::RepoWalkParams;
use crate::detail::walk::RepoWalkTypeParams;
use crate::detail::walk::StepRoute;
use crate::detail::walk::TailingWalkVisitor;
use crate::detail::walk::WalkVisitor;

// We can chose to go direct from the ChangesetId to types keyed by it without loading the Changeset
fn roots_for_chunk(
    ids: HashSet<ChangesetId>,
    root_types: &HashSet<NodeType>,
) -> Result<Vec<OutgoingEdge>, Error> {
    let mut edges = vec![];
    for id in ids {
        for r in root_types {
            let n = match r {
                NodeType::BonsaiHgMapping => Node::BonsaiHgMapping(ChangesetKey {
                    inner: id,
                    filenode_known_derived: false,
                }),
                NodeType::PhaseMapping => Node::PhaseMapping(id),
                NodeType::Changeset => Node::Changeset(ChangesetKey {
                    inner: id,
                    filenode_known_derived: false,
                }),
                NodeType::ChangesetInfo => Node::ChangesetInfo(id),
                NodeType::ChangesetInfoMapping => Node::ChangesetInfoMapping(id),
                NodeType::DeletedManifestV2Mapping => Node::DeletedManifestV2Mapping(id),
                NodeType::FsnodeMapping => Node::FsnodeMapping(id),
                NodeType::SkeletonManifestMapping => Node::SkeletonManifestMapping(id),
                NodeType::BasenameSuffixSkeletonManifestMapping => {
                    Node::BasenameSuffixSkeletonManifestMapping(id)
                }
                NodeType::UnodeMapping => Node::UnodeMapping(id),
                _ => bail!("Unsupported root type for chunking {:?}", r),
            };
            if let Some(edge_type) = n.get_type().root_edge_type() {
                let edge = OutgoingEdge::new(edge_type, n);
                edges.push(edge);
            } else {
                bail!("Unsupported node type for root edges {:?}", n.get_type());
            }
        }
    }
    Ok(edges)
}

#[derive(Clone, Debug)]
pub struct ClearStateParams {
    pub sample_rate: u64,
    pub interned_types: HashSet<InternedType>,
    pub node_types: HashSet<NodeType>,
}

#[derive(Clone, Debug)]
pub struct ChunkingParams {
    pub chunk_size: usize,
    pub chunk_by: HashSet<NodeType>,
    pub direction: Direction,
    pub clear_state: Option<ClearStateParams>,
    pub checkpoints: Option<CheckpointsByName>,
    pub allow_remaining_deferred: bool,
    pub repo_lower_bound_override: Option<u64>,
    pub repo_upper_bound_override: Option<u64>,
}

#[derive(Clone, Debug)]
pub struct TailParams {
    pub tail_secs: Option<u64>,
    pub chunking: Option<ChunkingParams>,
    pub state_max_age: Duration,
}

// Represent that only one end of the bound is optional, depending on direction
enum BestBounds {
    NewestFirst(Option<u64>, u64),
    OldestFirst(u64, Option<u64>),
}

impl BestBounds {
    // Checkpoint if necessary given the existing bounds. If there was a change return new bounds and checkpoint.
    async fn checkpoint(
        &self,
        logger: &Logger,
        repo_id: RepositoryId,
        checkpoints: &CheckpointsByName,
        chunk_low: u64,
        chunk_upper: u64,
        checkpoint: Option<Checkpoint>,
        chunk_num: u64,
    ) -> Result<Option<(BestBounds, Checkpoint)>, Error> {
        match self {
            BestBounds::NewestFirst(best_low, repo_high_bound) => {
                let new_best = best_low.map_or_else(
                    || Some(chunk_low),
                    |v| if chunk_low < v { Some(chunk_low) } else { None },
                );
                if let Some(new_best) = new_best {
                    let checkpoint = checkpoints
                        .persist(
                            logger,
                            repo_id,
                            chunk_num,
                            checkpoint,
                            new_best,
                            *repo_high_bound,
                        )
                        .await?;
                    return Ok(Some((
                        BestBounds::NewestFirst(Some(new_best), *repo_high_bound),
                        checkpoint,
                    )));
                }
            }
            BestBounds::OldestFirst(repo_low_bound, best_high) => {
                let new_best = best_high.map_or_else(
                    || Some(chunk_upper),
                    |v| {
                        if chunk_upper > v {
                            Some(chunk_upper)
                        } else {
                            None
                        }
                    },
                );
                if let Some(new_best) = new_best {
                    let checkpoint = checkpoints
                        .persist(
                            logger,
                            repo_id,
                            chunk_num,
                            checkpoint,
                            *repo_low_bound,
                            new_best,
                        )
                        .await?;
                    return Ok(Some((
                        BestBounds::OldestFirst(*repo_low_bound, Some(new_best)),
                        checkpoint,
                    )));
                }
            }
        }
        Ok(None)
    }
}

pub async fn walk_exact_tail<RunFac, SinkFac, SinkOut, V, VOut, Route>(
    fb: FacebookInit,
    job_params: JobWalkParams,
    mut repo_params: RepoWalkParams,
    type_params: RepoWalkTypeParams,
    tail_params: TailParams,
    mut visitor: V,
    make_run: RunFac,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error>
where
    RunFac: 'static + Clone + Send + Sync + FnOnce(&CoreContext, &RepoWalkParams) -> SinkFac,
    SinkFac: 'static
        + FnOnce(BoxStream<'static, Result<VOut, Error>>, Timestamp, u64, Option<String>) -> SinkOut
        + Clone
        + Send,
    SinkOut: Future<Output = Result<(), Error>> + 'static + Send,
    V: 'static + TailingWalkVisitor + WalkVisitor<VOut, Route> + Send + Sync,
    VOut: 'static + Send,
    Route: 'static + Send + Clone + StepRoute,
{
    let repo_id = repo_params.repo.repo_identity().id();

    let mut state_start = Timestamp::now();

    let with_hg = repo_params.include_node_types.iter().any(|n| {
        let n = n.derived_data_name();
        n == Some(MappedHgChangesetId::NAME) || n == Some(FilenodesOnlyPublic::NAME)
    });

    // At every iteration, check if cancellation is requested by the caller.
    while !cancellation_requested.load(Ordering::Relaxed) {
        cloned!(job_params, tail_params, type_params, make_run);
        let tail_secs = tail_params.tail_secs;
        // Each loop get new ctx and thus session id so we can distinguish runs
        let ctx = CoreContext::new_with_logger(fb, repo_params.logger.clone());
        let session_text = ctx.session().metadata().session_id().to_string();
        if !job_params.quiet {
            info!(
                repo_params.logger,
                "Starting walk with session id {}", &session_text
            )
        }
        repo_params.scuba_builder.add("session", session_text);

        let mut checkpoint = if let Some(checkpoints) = tail_params
            .chunking
            .as_ref()
            .and_then(|v| v.checkpoints.as_ref())
        {
            checkpoints.load(repo_id).await?
        } else {
            None
        };

        if let Some(cp) = checkpoint.as_ref() {
            info!(repo_params.logger, #log::CHUNKING, "Found checkpoint with bounds: ({}, {})", cp.lower_bound, cp.upper_bound);
        }

        let fetcher_params = tail_params
            .chunking
            .as_ref()
            .map(|chunking| {
                let heads_fetcher = PublicChangesetBulkFetch::new(
                    repo_params.repo.changesets_arc(),
                    repo_params.repo.phases_arc(),
                )
                .with_read_from_master(false)
                .with_step(MAX_FETCH_STEP);
                heads_fetcher.map(|v| (chunking, v))
            })
            .transpose()?;

        let is_chunking = fetcher_params.is_some();
        let mut run_start = Timestamp::now();
        let mut chunk_smaller_than_fetch = None;

        // Get the chunk stream and whether the bounds it covers are contiguous
        let (contiguous_bounds, mut best_bounds, chunk_stream) = if let Some((
            chunking,
            heads_fetcher,
        )) = &fetcher_params
        {
            if chunking.chunk_size < MAX_FETCH_STEP as usize {
                chunk_smaller_than_fetch = Some(chunking.direction);
            }
            let (mut lower, mut upper) = heads_fetcher.get_repo_bounds(&ctx).await?;
            if let Some(lower_override) = chunking.repo_lower_bound_override {
                lower = lower_override;
            }
            if let Some(upper_override) = chunking.repo_upper_bound_override {
                upper = upper_override;
            }

            info!(repo_params.logger, #log::CHUNKING, "Repo bounds: ({}, {})", lower, upper);

            let (contiguous_bounds, best_bound, catchup_bounds, main_bounds) = if let Some(
                ref mut checkpoint,
            ) = checkpoint
            {
                let age_secs = checkpoint.create_timestamp.since_seconds();
                run_start = checkpoint.create_timestamp;
                if age_secs >= 0 && Duration::from_secs(age_secs as u64) > tail_params.state_max_age
                {
                    info!(repo_params.logger, #log::CHUNKING, "Checkpoint run {} chunk {} is too old at {}s, running from repo bounds",
                        checkpoint.update_run_number, checkpoint.update_chunk_number, age_secs);
                    // Increment checkpoints run, reset chunk and create_timestamp as we're starting again
                    checkpoint.update_run_number += 1;
                    checkpoint.update_chunk_number = 0;
                    checkpoint.create_timestamp = Timestamp::now();
                    run_start = checkpoint.create_timestamp;
                    (true, None, None, Some((lower, upper)))
                } else {
                    let (catchup_bounds, main_bounds) =
                        checkpoint.stream_bounds(lower, upper, chunking.direction)?;

                    let contiguous_bounds = match (chunking.direction, catchup_bounds, main_bounds)
                    {
                        (
                            Direction::NewestFirst,
                            Some((catchup_lower, _)),
                            Some((_, main_upper)),
                        ) => catchup_lower == main_upper,
                        (
                            Direction::OldestFirst,
                            Some((_, catchup_upper)),
                            Some((main_lower, _)),
                        ) => catchup_upper == main_lower,
                        (_, Some(_), None) => false,
                        _ => true,
                    };
                    info!(repo_params.logger, #log::CHUNKING, "Continuing from checkpoint run {} chunk {} with catchup {:?} and main {:?} bounds",
                        checkpoint.update_run_number, checkpoint.update_chunk_number, catchup_bounds, main_bounds);
                    (
                        contiguous_bounds,
                        if chunking.direction == Direction::NewestFirst {
                            Some(checkpoint.lower_bound)
                        } else {
                            Some(checkpoint.upper_bound)
                        },
                        catchup_bounds,
                        main_bounds,
                    )
                }
            } else {
                (true, None, None, Some((lower, upper)))
            };

            let load_ids = |(lower, upper)| {
                heads_fetcher
                    .fetch_ids(&ctx, chunking.direction, Some((lower, upper)))
                    .chunks(chunking.chunk_size)
                    .map(move |v| v.into_iter().collect::<Result<HashSet<_>, Error>>())
            };

            let main_s = if let Some(main_bounds) = main_bounds {
                load_ids(main_bounds).left_stream()
            } else {
                stream::once(future::ok(HashSet::new())).right_stream()
            };

            let s = if let Some(catchup_bounds) = catchup_bounds {
                load_ids(catchup_bounds).chain(main_s).left_stream()
            } else {
                main_s.right_stream()
            };

            let best_bounds = if chunking.direction == Direction::NewestFirst {
                BestBounds::NewestFirst(best_bound, upper)
            } else {
                BestBounds::OldestFirst(lower, best_bound)
            };

            (contiguous_bounds, Some(best_bounds), s.left_stream())
        } else {
            let s = stream::once(future::ok(HashSet::new())).right_stream();
            (true, None, s)
        };

        let mut chunk_num: u64 = 0;
        if let Some(checkpoint) = checkpoint.as_ref() {
            chunk_num = checkpoint.update_chunk_number;
        }

        let mut last_chunk_low = None;
        let mut last_chunk_upper = None;

        futures::pin_mut!(chunk_stream);
        // Before beginning processing, check if the caller expects us to
        // stop. Can directly return from here since no meaningful work has
        // been performed so far.
        if cancellation_requested.load(Ordering::Relaxed) {
            return Ok(());
        }
        while let Some(chunk_members) = chunk_stream.try_next().await? {
            if is_chunking && chunk_members.is_empty() {
                continue;
            }
            chunk_num += 1;

            // convert from stream of (id, bounds) to ids plus overall bounds
            let mut chunk_low: u64 = u64::MAX;
            let mut chunk_upper: u64 = 0;
            let chunk_members: HashSet<ChangesetId> = chunk_members
                .into_iter()
                .map(|((cs_id, id), (fetch_low, fetch_upper))| {
                    if let Some(direction) = chunk_smaller_than_fetch {
                        // Adjust the bounds so it doesn't exceed previous chunk
                        if direction == Direction::NewestFirst {
                            chunk_low = min(chunk_low, id);
                            chunk_upper = max(chunk_upper, fetch_upper);
                            if let Some(last_chunk_low) = last_chunk_low {
                                chunk_upper = min(last_chunk_low, chunk_upper)
                            }
                        } else {
                            chunk_low = min(chunk_low, fetch_low);
                            if let Some(last_chunk_upper) = last_chunk_upper {
                                chunk_low = max(last_chunk_upper, chunk_low)
                            }
                            // Top of range is exclusive, so add one to the found id
                            chunk_upper = max(chunk_upper, id + 1);
                        }
                    } else {
                        // no need to adjust
                        chunk_low = min(chunk_low, fetch_low);
                        chunk_upper = max(chunk_upper, fetch_upper);
                    }
                    cs_id
                })
                .collect();

            cloned!(repo_params.logger);
            if is_chunking {
                match (last_chunk_low, last_chunk_upper) {
                    (Some(last_chunk_low), Some(last_chunk_upper))
                        if last_chunk_low == chunk_low && last_chunk_upper == chunk_upper =>
                    {
                        bail!(
                            "No progress at chunk {} with bounds ({}, {})",
                            chunk_num,
                            chunk_low,
                            chunk_upper
                        )
                    }
                    _ => {
                        info!(logger, #log::CHUNKING, "Starting chunk {} with bounds ({}, {})", chunk_num, chunk_low, chunk_upper)
                    }
                }
            }

            last_chunk_low.replace(chunk_low);
            last_chunk_upper.replace(chunk_upper);

            cloned!(mut repo_params);
            let hg_mapping_prepop = if with_hg && is_chunking {
                // bulk prepopulate the hg/bonsai mappings
                let ids =
                    BonsaiOrHgChangesetIds::Bonsai(chunk_members.clone().into_iter().collect());
                repo_params.repo.bonsai_hg_mapping().get(&ctx, ids).await?
            } else {
                vec![]
            };

            let extra_roots = visitor
                .start_chunk(&chunk_members, hg_mapping_prepop)?
                .into_iter();
            if let Some(chunking) = tail_params.chunking.as_ref() {
                let chunk_roots = roots_for_chunk(chunk_members, &chunking.chunk_by)?;
                repo_params.walk_roots.extend(chunk_roots);
            }
            repo_params.walk_roots.extend(extra_roots);

            cloned!(ctx, job_params, make_run, type_params);
            let make_sink = make_run(&ctx, &repo_params);

            // Walk needs clonable visitor, so wrap in Arc for its duration
            let arc_v = Arc::new(visitor);
            let walk_output =
                walk_exact(ctx, arc_v.clone(), job_params, repo_params, type_params).boxed();
            let cp_name = tail_params
                .chunking
                .as_ref()
                .and_then(|chunking| chunking.checkpoints.as_ref())
                .map(|v| v.name().to_string());
            make_sink(walk_output, run_start, chunk_num, cp_name).await?;
            visitor = Arc::try_unwrap(arc_v).map_err(|_| anyhow!("could not unwrap visitor"))?;

            if let Some(chunking) = tail_params.chunking.as_ref() {
                info!(logger, #log::LOADED, "Deferred: {}", visitor.num_deferred());
                if let Some(clear_state) = chunking.clear_state.as_ref() {
                    if clear_state.sample_rate != 0 && chunk_num % clear_state.sample_rate == 0 {
                        info!(logger, #log::CHUNKING, "Clearing state after chunk {}", chunk_num);
                        visitor.clear_state(&clear_state.node_types, &clear_state.interned_types);
                    }
                }

                // Record checkpoint and update best_bounds
                if let Some(checkpoints) = chunking.checkpoints.as_ref() {
                    if checkpoints.sample_rate != 0 && chunk_num % checkpoints.sample_rate == 0 {
                        let maybe_new = if let Some(best_bounds) = best_bounds.as_ref() {
                            best_bounds
                                .checkpoint(
                                    &logger,
                                    repo_id,
                                    checkpoints,
                                    chunk_low,
                                    chunk_upper,
                                    checkpoint.clone(),
                                    chunk_num,
                                )
                                .await?
                        } else {
                            None
                        };
                        if let Some((new_best, new_cp)) = maybe_new {
                            checkpoint.replace(new_cp);
                            best_bounds.replace(new_best);
                        }
                    }
                }
            }
            // Before processing the next chunk, check if cancellation has been requested.
            // Don't return directly since checkpointing needs to be performed to save the
            // progress made till now.
            // NOTE: This check is not part of the while conditional since using
            // let + condition is currently unstable.
            if cancellation_requested.load(Ordering::Relaxed) {
                break;
            }
        }

        if let Some(chunking) = tail_params.chunking.as_ref() {
            visitor.end_chunks(
                &repo_params.logger,
                    contiguous_bounds
                        && !chunking.allow_remaining_deferred
                        // If lower bound overridden then not contiguous to repo start. Overriding upper bound should not result in deferred.
                        && ((chunking.direction == Direction::NewestFirst && chunking.repo_lower_bound_override.is_none())
                            || (chunking.direction == Direction::OldestFirst && chunking.repo_upper_bound_override.is_none())),
                )?;

            match (chunking.checkpoints.as_ref(), checkpoint.as_ref()) {
                (Some(checkpoints), Some(cp)) => checkpoints.finish(repo_id, cp).await?,
                _ => {}
            }

            info!(
                repo_params.logger, #log::CHUNKING,
                "Completed in {} chunks of size {}", chunk_num, chunking.chunk_size
            );
        };

        match tail_secs {
            Some(interval) => {
                let start = Instant::now();
                let next_iter_deadline = start + Duration::from_secs(interval);
                tokio::time::sleep_until(next_iter_deadline).await;
                let age_secs = state_start.since_seconds();
                if age_secs >= 0 && Duration::from_secs(age_secs as u64) > tail_params.state_max_age
                {
                    // Walk state is too old, clear it.
                    info!(
                        repo_params.logger,
                        "Clearing walk state after {} seconds", age_secs
                    );
                    visitor
                        .clear_state(&NodeType::iter().collect(), &InternedType::iter().collect());
                    state_start = Timestamp::now();
                }
            }
            None => return Ok(()),
        }
    }
    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.
 */

// This module allows the implementation of validating checks over the mononoke graph
// Currently checks are added by
//  1. Add a CheckType variant
//  2. Add CheckType::node_type() and CheckType::enum_type() cases for the new variant
//  3. Add a new validation method
//  4. Add the method to the match/case in ValidatingVisitor::visit()

use std::collections::HashMap;
use std::collections::HashSet;
use std::fmt;
use std::hash::Hash;
use std::sync::atomic::AtomicBool;
use std::sync::Arc;
use std::time::Instant;

use anyhow::Error;
use async_trait::async_trait;
use bonsai_hg_mapping::BonsaiHgMapping;
use bonsai_hg_mapping::BonsaiHgMappingEntry;
use bulkops::Direction;
use cloned::cloned;
use context::CoreContext;
use derive_more::AddAssign;
use fbinit::FacebookInit;
use futures::future::try_join_all;
use futures::stream::TryStreamExt;
use itertools::Itertools;
use maplit::hashset;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;
use mononoke_types::MPath;
use phases::Phase;
use phases::Phases;
use repo_identity::RepoIdentityRef;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::info;
use slog::warn;
use slog::Logger;
use stats::prelude::*;

use crate::commands::JobParams;
use crate::commands::JobWalkParams;
use crate::commands::RepoSubcommandParams;
use crate::detail::graph::EdgeType;
use crate::detail::graph::Node;
use crate::detail::graph::NodeData;
use crate::detail::graph::NodeType;
use crate::detail::graph::WrappedPath;
use crate::detail::log;
use crate::detail::progress::progress_stream;
use crate::detail::progress::report_state;
use crate::detail::progress::sort_by_string;
use crate::detail::progress::ProgressOptions;
use crate::detail::progress::ProgressRecorder;
use crate::detail::progress::ProgressRecorderUnprotected;
use crate::detail::progress::ProgressReporter;
use crate::detail::progress::ProgressReporterUnprotected;
use crate::detail::progress::ProgressStateMutex;
use crate::detail::state::InternedType;
use crate::detail::state::StepStats;
use crate::detail::state::WalkState;
use crate::detail::tail::walk_exact_tail;
use crate::detail::walk::EmptyRoute;
use crate::detail::walk::OutgoingEdge;
use crate::detail::walk::RepoWalkParams;
use crate::detail::walk::RepoWalkTypeParams;
use crate::detail::walk::StepRoute;
use crate::detail::walk::TailingWalkVisitor;
use crate::detail::walk::VisitOne;
use crate::detail::walk::WalkVisitor;

pub const NODES: &str = "nodes";
pub const EDGES: &str = "edges";
pub const PASS: &str = "pass";
pub const FAIL: &str = "fail";
pub const TOTAL: &str = "total";
pub const NODE_KEY: &str = "node_key";
pub const NODE_TYPE: &str = "node_type";
pub const NODE_PATH: &str = "node_path";
pub const EDGE_TYPE: &str = "edge_type";
pub const CHECK_TYPE: &str = "check_type";
pub const CHECK_FAIL: &str = "check_fail";
pub const CHECK_SIZE: &str = "check_size";
pub const WALK_TYPE: &str = "walk_type";
pub const REPO: &str = "repo";
pub const ERROR_MSG: &str = "error_msg";
const SRC_NODE_KEY: &str = "src_node_key";
const SRC_NODE_TYPE: &str = "src_node_type";
const SRC_NODE_PATH: &str = "src_node_path";
const VIA_NODE_KEY: &str = "via_node_key";
const VIA_NODE_TYPE: &str = "via_node_type";
const VIA_NODE_PATH: &str = "via_node_path";

define_stats! {
    prefix = "mononoke.walker.validate";
    // e.g. mononoke.walker.validate.testrepo.hg_link_node_populated.pass
    walker_validate: dynamic_timeseries("{}.{}.{}", (repo: String, check: &'static str, status: &'static str); Rate, Sum),
    last_completed: dynamic_singleton_counter("{}.{}.last_completed.{}", (repo: String, check: &'static str, status: &'static str)),
}

pub const DEFAULT_CHECK_TYPES: &[CheckType] = &[
    CheckType::ChangesetPhaseIsPublic,
    CheckType::HgLinkNodePopulated,
];

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
struct ValidateInfo {
    // Where we stepped from, useful for immediate reproductions with --walk-root
    source_node: Option<Node>,
    // What the check thinks is an interesting node on the route to here (e.g. the affected changeset)
    via_node: Option<Node>,
    // if the walk has extra path tracking enabled, this is the resolved path of the resolved node
    resolved_path: Option<WrappedPath>,
    // if the check wants to report a size
    check_size: Option<u64>,
}

impl ValidateInfo {
    fn new(
        source_node: Option<Node>,
        via_node: Option<Node>,
        resolved_path: Option<WrappedPath>,
        check_size: Option<u64>,
    ) -> Self {
        Self {
            source_node,
            via_node,
            resolved_path,
            check_size,
        }
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
enum CheckStatus {
    Fail(ValidateInfo),
    Pass(Option<ValidateInfo>),
}

define_type_enum! {
enum CheckType {
    ChangesetPhaseIsPublic,
    HgLinkNodePopulated,
    FileContentIsLfs,
}
}

impl CheckType {
    fn stats_key(&self) -> &'static str {
        match self {
            CheckType::ChangesetPhaseIsPublic => "bonsai_phase_is_public",
            CheckType::HgLinkNodePopulated => "hg_link_node_populated",
            CheckType::FileContentIsLfs => "file_content_is_lfs",
        }
    }
    pub fn node_type(&self) -> NodeType {
        match self {
            CheckType::ChangesetPhaseIsPublic => NodeType::PhaseMapping,
            CheckType::HgLinkNodePopulated => NodeType::HgFileNode,
            CheckType::FileContentIsLfs => NodeType::FileContentMetadataV2,
        }
    }
}

impl fmt::Display for CheckType {
    fn fmt(&self, f: &mut fmt::Formatter) -> fmt::Result {
        write!(f, "{:?}", self)
    }
}

#[derive(Debug)]
struct CheckOutput {
    check: CheckType,
    status: CheckStatus,
}

impl CheckOutput {
    fn new(check: CheckType, status: CheckStatus) -> Self {
        Self { check, status }
    }
}

struct ValidatingVisitor {
    repo_stats_key: String,
    inner: WalkState,
    checks_by_node_type: HashMap<NodeType, HashSet<CheckType>>,
    lfs_threshold: Option<u64>,
}

impl ValidatingVisitor {
    pub fn new(
        repo_stats_key: String,
        include_node_types: HashSet<NodeType>,
        include_edge_types: HashSet<EdgeType>,
        include_checks: HashSet<CheckType>,
        always_emit_edge_types: HashSet<EdgeType>,
        enable_derive: bool,
        lfs_threshold: Option<u64>,
        chunk_direction: Option<Direction>,
    ) -> Self {
        Self {
            repo_stats_key,
            inner: WalkState::new(
                include_node_types,
                include_edge_types,
                always_emit_edge_types,
                enable_derive,
                chunk_direction,
            ),
            checks_by_node_type: include_checks
                .into_iter()
                .group_by(|c| c.node_type())
                .into_iter()
                .map(|(key, group)| (key, HashSet::from_iter(group)))
                .collect(),
            lfs_threshold,
        }
    }
}

#[async_trait]
impl VisitOne for ValidatingVisitor {
    fn in_chunk(&self, bcs_id: &ChangesetId) -> bool {
        self.inner.in_chunk(bcs_id)
    }

    fn needs_visit(&self, outgoing: &OutgoingEdge) -> bool {
        self.inner.needs_visit(outgoing)
    }

    async fn is_public(
        &self,
        ctx: &CoreContext,
        phases_store: &dyn Phases,
        bcs_id: &ChangesetId,
    ) -> Result<bool, Error> {
        self.inner.is_public(ctx, phases_store, bcs_id).await
    }

    fn get_hg_from_bonsai(&self, bcs_id: &ChangesetId) -> Option<HgChangesetId> {
        self.inner.get_hg_from_bonsai(bcs_id)
    }

    fn record_hg_from_bonsai(&self, bcs_id: &ChangesetId, hg_cs_id: HgChangesetId) {
        self.inner.record_hg_from_bonsai(bcs_id, hg_cs_id)
    }

    async fn get_bonsai_from_hg(
        &self,
        ctx: &CoreContext,
        bonsai_hg_mapping: &dyn BonsaiHgMapping,
        hg_cs_id: &HgChangesetId,
    ) -> Result<ChangesetId, Error> {
        self.inner
            .get_bonsai_from_hg(ctx, bonsai_hg_mapping, hg_cs_id)
            .await
    }

    async fn defer_from_hg(
        &self,
        ctx: &CoreContext,
        bonsai_hg_mapping: &dyn BonsaiHgMapping,
        hg_cs_id: &HgChangesetId,
    ) -> Result<Option<ChangesetId>, Error> {
        self.inner
            .defer_from_hg(ctx, bonsai_hg_mapping, hg_cs_id)
            .await
    }
}

fn check_bonsai_phase_is_public(
    node: &Node,
    node_data: Option<&NodeData>,
    route: Option<&ValidateRoute>,
) -> CheckStatus {
    match (&node, &node_data) {
        (Node::PhaseMapping(_cs_id), Some(NodeData::PhaseMapping(Some(Phase::Public)))) => {
            CheckStatus::Pass(None)
        }
        (Node::PhaseMapping(non_public_cs_id), Some(NodeData::PhaseMapping(_phase))) => {
            let via = route.and_then(|r| {
                for n in r.via.iter().rev() {
                    match n {
                        Node::HgChangeset(_) => return Some(n.clone()),
                        Node::Changeset(k) => {
                            // Check for most recent non-identical changesethg
                            if &k.inner != non_public_cs_id {
                                return Some(n.clone());
                            }
                        }
                        _ => {}
                    }
                }
                None
            });
            CheckStatus::Fail(ValidateInfo::new(
                route.map(|r| r.src_node.clone()),
                via,
                None,
                None,
            ))
        }
        _ => CheckStatus::Fail(ValidateInfo::new(
            route.map(|r| r.src_node.clone()),
            None,
            None,
            None,
        )),
    }
}

fn check_linknode_populated(
    outgoing: &[OutgoingEdge],
    route: Option<&ValidateRoute>,
) -> CheckStatus {
    if outgoing
        .iter()
        .any(|e| e.label == EdgeType::HgFileNodeToLinkedHgChangeset)
    {
        CheckStatus::Pass(None)
    } else {
        let via = route.and_then(|r| {
            for n in r.via.iter().rev() {
                match n {
                    Node::HgChangeset(_) => return Some(n.clone()),
                    Node::Changeset(_) => return Some(n.clone()),
                    _ => {}
                }
            }
            None
        });
        CheckStatus::Fail(ValidateInfo::new(
            route.map(|r| r.src_node.clone()),
            via,
            None,
            None,
        ))
    }
}

fn check_file_content_is_lfs(
    resolved: &OutgoingEdge,
    node_data: Option<&NodeData>,
    route: Option<&ValidateRoute>,
    lfs_threshold: u64,
) -> CheckStatus {
    match (&resolved.target, &node_data) {
        (
            Node::FileContentMetadataV2(_content_id),
            Some(NodeData::FileContentMetadataV2(Some(content_meta))),
        ) => {
            if content_meta.total_size < lfs_threshold {
                // not interesting to log
                CheckStatus::Pass(None)
            } else {
                let via = route.and_then(|r| {
                    for n in r.via.iter().rev() {
                        match &n {
                            Node::HgChangeset(_) => return Some(n.clone()),
                            Node::HgBonsaiMapping(_) => return Some(n.clone()),
                            _ => {}
                        }
                    }
                    None
                });
                let info = Some(ValidateInfo::new(
                    route.map(|r| r.src_node.clone()),
                    via,
                    resolved.path.clone(),
                    Some(content_meta.total_size),
                ));
                CheckStatus::Pass(info)
            }
        }
        // Unexpected node type
        _ => CheckStatus::Fail(ValidateInfo::new(
            route.map(|r| r.src_node.clone()),
            None,
            None,
            None,
        )),
    }
}

#[derive(AddAssign, Clone, Copy, Default, Debug)]
struct CheckStats {
    pass: u64,
    fail: u64,
    edges: u64,
}

struct CheckData {
    checked: Vec<CheckOutput>,
    stats: CheckStats,
}

#[derive(Clone, Debug)]
struct ValidateRoute {
    src_node: Node,
    via: Vec<Node>,
}

impl ValidateRoute {
    // Keep memory usage bounded
    const MAX_VIA: usize = 2;

    fn next_route(route: Option<Self>, node: Node) -> Self {
        let mut next_via = match route {
            Some(Self {
                src_node: _src_node,
                mut via,
            }) => {
                if via.len() > ValidateRoute::MAX_VIA {
                    via.remove(0);
                }
                via
            }
            None => vec![],
        };

        // Only track changesets for the via information
        match node {
            Node::HgChangeset(_) | Node::HgBonsaiMapping(_) | Node::Changeset(_) => {
                next_via.push(node.clone())
            }
            _ => {}
        };
        Self {
            src_node: node,
            via: next_via,
        }
    }
}

impl StepRoute for ValidateRoute {
    fn source_node(&self) -> Option<&Node> {
        Some(&self.src_node)
    }
    fn via_node(&self) -> Option<&Node> {
        self.via.last()
    }
}

impl TailingWalkVisitor for ValidatingVisitor {
    fn start_chunk(
        &mut self,
        chunk_members: &HashSet<ChangesetId>,
        mapping_prepop: Vec<BonsaiHgMappingEntry>,
    ) -> Result<HashSet<OutgoingEdge>, Error> {
        self.inner.start_chunk(chunk_members, mapping_prepop)
    }

    fn clear_state(
        &mut self,
        node_types: &HashSet<NodeType>,
        interned_types: &HashSet<InternedType>,
    ) {
        self.inner.clear_state(node_types, interned_types)
    }

    fn end_chunks(&mut self, logger: &Logger, contiguous_bounds: bool) -> Result<(), Error> {
        self.inner.end_chunks(logger, contiguous_bounds)
    }

    fn num_deferred(&self) -> usize {
        self.inner.num_deferred()
    }
}

impl WalkVisitor<(Node, Option<CheckData>, Option<StepStats>), ValidateRoute>
    for ValidatingVisitor
{
    fn start_step(
        &self,
        ctx: CoreContext,
        route: Option<&ValidateRoute>,
        step: &OutgoingEdge,
    ) -> Option<CoreContext> {
        self.inner
            .start_step(ctx, route.map(|_| &EmptyRoute {}), step)
    }

    fn visit(
        &self,
        ctx: &CoreContext,
        resolved: OutgoingEdge,
        node_data: Option<NodeData>,
        route: Option<ValidateRoute>,
        outgoing: Vec<OutgoingEdge>,
    ) -> (
        (Node, Option<CheckData>, Option<StepStats>),
        ValidateRoute,
        Vec<OutgoingEdge>,
    ) {
        let checks_to_do: Option<&HashSet<_>> =
            self.checks_by_node_type.get(&resolved.target.get_type());
        // The incoming resolved edge counts as one
        let mut num_edges: u64 = 1;
        let mut pass = 0;
        let mut fail = 0;
        let checked: Vec<_> = checks_to_do
            .map(|set| {
                set.iter().map(|check| {
                    // Lets check!
                    let status = match check {
                        CheckType::ChangesetPhaseIsPublic => check_bonsai_phase_is_public(
                            &resolved.target,
                            node_data.as_ref(),
                            route.as_ref(),
                        ),
                        CheckType::HgLinkNodePopulated => {
                            num_edges += outgoing.len() as u64;
                            check_linknode_populated(&outgoing, route.as_ref())
                        }
                        CheckType::FileContentIsLfs => {
                            num_edges += outgoing.len() as u64;
                            if let Some(lfs_threshold) = self.lfs_threshold {
                                check_file_content_is_lfs(
                                    &resolved,
                                    node_data.as_ref(),
                                    route.as_ref(),
                                    lfs_threshold,
                                )
                            } else {
                                CheckStatus::Pass(None)
                            }
                        }
                    };
                    match &status {
                        CheckStatus::Pass(_) => pass += 1,
                        CheckStatus::Fail(_) => fail += 1,
                    }
                    CheckOutput::new(*check, status)
                })
            })
            .into_iter()
            .flatten()
            .collect();

        STATS::walker_validate.add_value(
            num_edges as i64,
            (self.repo_stats_key.clone(), EDGES, TOTAL),
        );

        // Call inner after checks. otherwise it will prune outgoing edges we wanted to check.
        let ((node, _opt_data, opt_stats), _, outgoing) = self.inner.visit(
            ctx,
            resolved,
            node_data,
            route.as_ref().map(|_| EmptyRoute {}),
            outgoing,
        );

        let vout = (
            node.clone(),
            if checked.is_empty() {
                None
            } else {
                Some(CheckData {
                    checked,
                    stats: CheckStats {
                        pass,
                        fail,
                        edges: num_edges,
                    },
                })
            },
            opt_stats,
        );

        (vout, ValidateRoute::next_route(route, node), outgoing)
    }

    fn defer_visit(
        &self,
        bcs_id: &ChangesetId,
        walk_item: &OutgoingEdge,
        route: Option<ValidateRoute>,
    ) -> Result<((Node, Option<CheckData>, Option<StepStats>), ValidateRoute), Error> {
        let ((node, _node_data, stats), _route) =
            self.inner
                .defer_visit(bcs_id, walk_item, Some(EmptyRoute {}))?;
        Ok((
            (node.clone(), None, stats),
            ValidateRoute::next_route(route, node),
        ))
    }
}

struct ValidateProgressState {
    logger: Logger,
    fb: FacebookInit,
    scuba_builder: MononokeScubaSampleBuilder,
    repo_stats_key: String,
    types_sorted_by_name: Vec<CheckType>,
    stats_by_type: HashMap<CheckType, CheckStats>,
    total_checks: CheckStats,
    checked_nodes: u64,
    passed_nodes: u64,
    failed_nodes: u64,
    throttle_options: ProgressOptions,
    last_update: Instant,
}

impl ValidateProgressState {
    fn new(
        logger: Logger,
        fb: FacebookInit,
        scuba_builder: MononokeScubaSampleBuilder,
        repo_stats_key: String,
        included_types: HashSet<CheckType>,
        throttle_options: ProgressOptions,
    ) -> Self {
        let types_sorted_by_name = sort_by_string(included_types);
        let now = Instant::now();
        Self {
            logger,
            fb,
            scuba_builder,
            repo_stats_key,
            types_sorted_by_name,
            stats_by_type: HashMap::new(),
            total_checks: CheckStats::default(),
            checked_nodes: 0,
            passed_nodes: 0,
            failed_nodes: 0,
            throttle_options,
            last_update: now,
        }
    }

    fn report_progress_log(&self) {
        let detail_by_type = &self
            .types_sorted_by_name
            .iter()
            .map(|t| {
                let d = CheckStats::default();
                let stats = self.stats_by_type.get(t).unwrap_or(&d);
                format!("{}:{},{}", t, stats.pass, stats.fail)
            })
            .collect::<Vec<_>>()
            .join(" ");
        info!(
            self.logger,
            #log::VALIDATE,
            "Nodes,Pass,Fail:{},{},{}; EdgesChecked:{}; CheckType:Pass,Fail Total:{},{} {}",
            self.checked_nodes,
            self.passed_nodes,
            self.failed_nodes,
            self.total_checks.edges,
            self.total_checks.pass,
            self.total_checks.fail,
            detail_by_type,
        );
    }

    fn report_progress_stats(&self) {
        // Per check type
        for (k, v) in self.stats_by_type.iter() {
            for (desc, value) in &[(PASS, v.pass), (FAIL, v.fail), (EDGES, v.edges)] {
                STATS::last_completed.set_value(
                    self.fb,
                    *value as i64,
                    (self.repo_stats_key.clone(), k.stats_key(), desc),
                );
            }
        }
        // Overall by nodes and edges
        for (stat, desc, value) in &[
            (NODES, PASS, self.passed_nodes),
            (NODES, FAIL, self.failed_nodes),
            (NODES, TOTAL, self.checked_nodes),
            (EDGES, TOTAL, self.total_checks.edges),
        ] {
            STATS::last_completed.set_value(
                self.fb,
                *value as i64,
                (self.repo_stats_key.clone(), stat, desc),
            );
        }
    }
}

fn scuba_log_node(
    n: &Node,
    n_path: Option<&WrappedPath>,
    scuba: &mut MononokeScubaSampleBuilder,
    type_key: &'static str,
    key_key: &'static str,
    path_key: &'static str,
) {
    scuba
        .add(type_key, Into::<&'static str>::into(n.get_type()))
        .add(key_key, n.stats_key());
    if let Some(path) = n.stats_path() {
        scuba.add(path_key, MPath::display_opt(path.as_ref()).to_string());
    } else if let Some(path) = n_path {
        scuba.add(path_key, MPath::display_opt(path.as_ref()).to_string());
    }
}

pub fn add_node_to_scuba(
    source_node: Option<&Node>,
    via_node: Option<&Node>,
    resolved_node: &Node,
    resolved_path: Option<&WrappedPath>,
    scuba: &mut MononokeScubaSampleBuilder,
) {
    scuba_log_node(
        resolved_node,
        resolved_path,
        scuba,
        NODE_TYPE,
        NODE_KEY,
        NODE_PATH,
    );
    if let Some(src_node) = source_node {
        scuba_log_node(
            src_node,
            None,
            scuba,
            SRC_NODE_TYPE,
            SRC_NODE_KEY,
            SRC_NODE_PATH,
        );
    }
    if let Some(via_node) = via_node {
        scuba_log_node(
            via_node,
            None,
            scuba,
            VIA_NODE_TYPE,
            VIA_NODE_KEY,
            VIA_NODE_PATH,
        );
    }
}

impl ProgressRecorderUnprotected<CheckData> for ValidateProgressState {
    fn set_sample_builder(&mut self, s: MononokeScubaSampleBuilder) {
        self.scuba_builder = s;
    }

    fn record_step(&mut self, resolved_node: &Node, checkdata: Option<&CheckData>) {
        self.checked_nodes += 1;
        let mut had_pass = false;
        let mut had_fail = false;
        if let Some(checkdata) = checkdata {
            // By node. One fail is enough for a Node to be failed.
            if checkdata.stats.fail > 0 {
                had_fail = true;
            } else if !had_fail && checkdata.stats.pass > 0 {
                had_pass = true;
            }
            // total
            self.total_checks += checkdata.stats;
            // By type
            for c in &checkdata.checked {
                let k = c.check;
                let stats = self
                    .stats_by_type
                    .entry(k)
                    .or_insert_with(CheckStats::default);
                let (validate_info, check_fail) = match &c.status {
                    CheckStatus::Pass(validate_info) => {
                        stats.pass += 1;
                        STATS::walker_validate
                            .add_value(1, (self.repo_stats_key.clone(), k.stats_key(), PASS));
                        (validate_info.into(), 0)
                    }
                    CheckStatus::Fail(validate_info) => {
                        STATS::walker_validate
                            .add_value(1, (self.repo_stats_key.clone(), k.stats_key(), FAIL));
                        stats.fail += 1;
                        (Some(validate_info), 1)
                    }
                };
                if let Some(validate_info) = validate_info {
                    let mut scuba = self.scuba_builder.clone();
                    add_node_to_scuba(
                        validate_info.source_node.as_ref(),
                        validate_info.via_node.as_ref(),
                        resolved_node,
                        validate_info.resolved_path.as_ref(),
                        &mut scuba,
                    );

                    if let Some(check_size) = validate_info.check_size {
                        scuba.add(CHECK_SIZE, check_size);
                    }

                    scuba
                        .add(CHECK_TYPE, k.stats_key())
                        .add(CHECK_FAIL, check_fail)
                        .log();
                    if check_fail > 0 {
                        if let Ok(json) = scuba.get_sample().to_json() {
                            warn!(self.logger, "Validation failed: {}", json)
                        }
                    }
                }
            }
        }

        if had_pass {
            self.passed_nodes += 1;
            STATS::walker_validate.add_value(1, (self.repo_stats_key.clone(), NODES, PASS));
        } else if had_fail {
            self.failed_nodes += 1;
            STATS::walker_validate.add_value(1, (self.repo_stats_key.clone(), NODES, FAIL));
        }
        STATS::walker_validate.add_value(1, (self.repo_stats_key.clone(), NODES, TOTAL));
    }
}

impl ProgressReporterUnprotected for ValidateProgressState {
    fn report_progress(&mut self) {
        self.report_progress_log();
        self.report_progress_stats();
    }

    fn report_throttled(&mut self) {
        if self.checked_nodes % self.throttle_options.sample_rate == 0 {
            let new_update = Instant::now();
            let delta_time = new_update.duration_since(self.last_update);
            if delta_time >= self.throttle_options.interval {
                self.report_progress_log();
                self.last_update = new_update;
            }
        }
    }
}

#[derive(Clone)]
pub struct ValidateCommand {
    pub include_check_types: HashSet<CheckType>,
    pub progress_options: ProgressOptions,
}

impl ValidateCommand {
    fn apply_repo(&mut self, repo_params: &RepoWalkParams) {
        self.include_check_types
            .retain(|t| repo_params.include_node_types.contains(&t.node_type()));
    }
}

// Subcommand entry point for validation of mononoke commit graph and dependent data
pub async fn validate(
    fb: FacebookInit,
    job_params: JobParams,
    command: ValidateCommand,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error> {
    let JobParams {
        walk_params,
        per_repo,
    } = job_params;

    let mut all_walks = Vec::new();
    for (sub_params, repo_params) in per_repo {
        cloned!(mut command, walk_params);

        command.apply_repo(&repo_params);

        let walk = run_one(
            fb,
            walk_params,
            sub_params,
            repo_params,
            command,
            Arc::clone(&cancellation_requested),
        );
        all_walks.push(walk);
    }
    try_join_all(all_walks).await.map(|_| ())
}

async fn run_one(
    fb: FacebookInit,
    job_params: JobWalkParams,
    sub_params: RepoSubcommandParams,
    repo_params: RepoWalkParams,
    command: ValidateCommand,
    cancellation_requested: Arc<AtomicBool>,
) -> Result<(), Error> {
    info!(
        repo_params.logger,
        #log::VALIDATE,
        "Performing check types {:?}",
        sort_by_string(&command.include_check_types)
    );

    let validate_progress_state = ProgressStateMutex::new(ValidateProgressState::new(
        repo_params.logger.clone(),
        fb,
        repo_params.scuba_builder.clone(),
        repo_params.repo.repo_identity().name().to_string(),
        command.include_check_types.clone(),
        command.progress_options,
    ));

    cloned!(job_params.quiet, sub_params.progress_state);
    let make_sink = move |ctx: &CoreContext, repo_params: &RepoWalkParams| {
        cloned!(ctx);
        validate_progress_state.set_sample_builder(repo_params.scuba_builder.clone());
        async move |walk_output, _run_start, _chunk_num, _checkpoint_name| {
            cloned!(ctx, progress_state, validate_progress_state);
            let walk_progress =
                progress_stream(quiet, &progress_state, walk_output).map_ok(|(n, d, s)| {
                    // swap stats and data round
                    (n, s, d)
                });

            let validate_progress = progress_stream(quiet, &validate_progress_state, walk_progress);

            report_state(ctx, validate_progress).await?;
            progress_state.report_progress();
            validate_progress_state.report_progress();
            Ok(())
        }
    };

    let always_emit_edge_types =
        HashSet::from_iter(vec![EdgeType::HgFileNodeToLinkedHgChangeset].into_iter());

    let mut required_node_data_types = hashset![NodeType::PhaseMapping];
    let mut keep_edge_paths = false;
    if command
        .include_check_types
        .contains(&CheckType::FileContentIsLfs)
    {
        required_node_data_types.insert(NodeType::FileContentMetadataV2);
        keep_edge_paths = true;
    }

    let stateful_visitor = ValidatingVisitor::new(
        repo_params.repo.repo_identity().name().to_string(),
        repo_params.include_node_types.clone(),
        repo_params.include_edge_types.clone(),
        command.include_check_types,
        always_emit_edge_types.clone(),
        job_params.enable_derive,
        sub_params.lfs_threshold,
        sub_params
            .tail_params
            .chunking
            .as_ref()
            .map(|v| v.direction),
    );

    let type_params = RepoWalkTypeParams {
        required_node_data_types,
        always_emit_edge_types,
        keep_edge_paths,
    };

    walk_exact_tail(
        fb,
        job_params,
        repo_params,
        type_params,
        sub_params.tail_params,
        stateful_visitor,
        make_sink,
        cancellation_requested,
    )
    .await
}