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

use std::collections::HashMap;
use std::collections::HashSet;
use std::num::NonZeroU64;
use std::sync::Arc;

use anyhow::bail;
use anyhow::format_err;
use anyhow::Context;
use anyhow::Error;
use blobrepo::BlobRepo;
use blobstore::Blobstore;
use blobstore_factory::ScrubHandler;
use cloned::cloned;
use fbinit::FacebookInit;
use metaconfig_types::CommonConfig;
use metaconfig_types::MetadataDatabaseConfig;
use metaconfig_types::Redaction;
use metaconfig_types::RepoConfig;
use metaconfig_types::WalkerJobParams;
use metaconfig_types::WalkerJobType;
use mononoke_app::args::MultiRepoArgs;
use mononoke_app::MononokeApp;
use mononoke_types::repo::RepositoryId;
use newfilenodes::NewFilenodesBuilder;
use repo_factory::RepoFactory;
use samplingblob::ComponentSamplingHandler;
use samplingblob::SamplingBlobstore;
use samplingblob::SamplingHandler;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::info;
use slog::o;
use slog::warn;
use slog::Logger;
use sql_ext::facebook::MysqlOptions;

use crate::args::NodeTypeArg;
use crate::args::TailArgs;
use crate::args::WalkerCommonArgs;
use crate::args::WalkerGraphParams;
use crate::commands::JobParams;
use crate::commands::JobWalkParams;
use crate::commands::RepoSubcommandParams;
use crate::detail::blobstore::replace_blobconfig;
use crate::detail::blobstore::StatsScrubHandler;
use crate::detail::graph::EdgeType;
use crate::detail::graph::NodeType;
use crate::detail::graph::SqlShardInfo;
use crate::detail::log;
use crate::detail::progress::sort_by_string;
use crate::detail::progress::ProgressOptions;
use crate::detail::progress::ProgressStateCountByType;
use crate::detail::progress::ProgressStateMutex;
use crate::detail::tail::TailParams;
use crate::detail::validate::REPO;
use crate::detail::validate::WALK_TYPE;
use crate::detail::walk::OutgoingEdge;
use crate::detail::walk::RepoWalkParams;
use crate::WalkerArgs;
const CHECKPOINT_PREFIX: &str = "mononoke_sharded_walker";

pub async fn setup_common<'a>(
    walk_stats_key: &'static str,
    app: &MononokeApp,
    repo_args: &MultiRepoArgs,
    common_args: &WalkerCommonArgs,
    blobstore_sampler: Option<Arc<dyn SamplingHandler>>,
    blobstore_component_sampler: Option<Arc<dyn ComponentSamplingHandler>>,
) -> Result<JobParams, Error> {
    let logger = app.logger();
    let mut scuba_builder = app.environment().scuba_sample_builder.clone();
    let walker_type = app.args::<WalkerArgs>()?.walker_type;
    scuba_builder.add(WALK_TYPE, walk_stats_key);

    let WalkerGraphParams {
        include_node_types,
        include_edge_types,
        error_as_data_node_types,
        error_as_data_edge_types,
    } = common_args.graph_params.parse_args()?;

    if !error_as_data_node_types.is_empty() || !error_as_data_edge_types.is_empty() {
        if !app.readonly_storage().0 {
            return Err(format_err!(
                "Error as data could mean internal state is invalid, run with --with-readonly-storage=true to ensure no risk of persisting it"
            ));
        }
        warn!(
            logger,
            "Error as data enabled, walk results may not be complete. Errors as data enabled for node types {:?} edge types {:?}",
            sort_by_string(&error_as_data_node_types),
            sort_by_string(&error_as_data_edge_types)
        );
    }

    // There is no need to check if repos is empty: at least one repo arg
    // is required when running the command.
    let repos = app.multi_repo_configs(repo_args.ids_or_names()?)?;
    let repo_count = repos.len();
    if repo_count > 1 {
        info!(
            logger,
            "Walking repos {:?}",
            repos.iter().map(|repo| &repo.0).collect::<Vec<_>>()
        );
    }

    let repos = override_repo_configs(walk_stats_key, app, common_args, repos)?;

    // configure repo factory
    let repo_id_to_name: HashMap<_, _> = repos
        .iter()
        .map(|(name, conf)| (conf.repoid, name.clone()))
        .collect();
    let repo_factory = setup_repo_factory(
        walk_stats_key,
        app,
        repo_id_to_name,
        blobstore_sampler,
        blobstore_component_sampler,
        scuba_builder.clone(),
        logger,
        common_args.quiet,
    );

    let progress_options = common_args.progress.parse_args();
    let hash_validation_node_types = common_args.hash_validation.parse_args();

    let mysql_options = app.mysql_options();

    let walk_roots = common_args.walk_roots.parse_args()?;
    let mut parsed_tail_params = parse_tail_params(
        app.fb,
        &common_args.tailing,
        mysql_options,
        &repos,
        &walk_roots,
    )?;

    let mut per_repo = Vec::new();
    let mut error_as_data_node_types_for_all_repos = error_as_data_node_types;
    for (repo, repo_conf) in repos {
        let metadatadb_config = &repo_conf.storage_config.metadata;
        let tail_params = parsed_tail_params
            .get_mut(metadatadb_config)
            .ok_or_else(|| format_err!("No tail params for {}", repo))?;

        // repo factory reuses sql factory if one was already initiated for the config
        let sql_factory = repo_factory.sql_factory(metadatadb_config).await?;
        let sql_shard_info = SqlShardInfo {
            filenodes: sql_factory.tier_info_shardable::<NewFilenodesBuilder>()?,
            active_keys_per_shard: mysql_options.per_key_limit(),
        };

        let walker_config_params = walker_type
            .as_ref()
            .and_then(|job_type| walker_config_params(&repo_conf, job_type));
        // Concurrency is primarily provided by config and then by
        // CLI in case config value is absent.
        let scheduled_max_concurrency = walker_config_params
            .and_then(|p| p.scheduled_max_concurrency.map(|i| i as usize))
            .unwrap_or(common_args.scheduled_max);
        // Exclude nodes that might be provided as part of walker config.
        let included_nodes = walker_config_params
            .and_then(|p| p.exclude_node_type.as_ref())
            .map(|s| s.parse::<NodeTypeArg>())
            .transpose()?
            .map(|n| HashSet::<NodeType>::from_iter(n.0.iter().cloned()))
            .map_or_else(
                || include_node_types.clone(),
                |excluded_nodes| {
                    include_node_types
                        .difference(&excluded_nodes)
                        .copied()
                        .collect()
                },
            );

        if let Some(ref mut chunking) = tail_params.chunking {
            // Allow Remaining Deferred = True if either the CLI or the walker
            // config say so.
            chunking.allow_remaining_deferred |=
                walker_config_params.map_or(false, |p| p.allow_remaining_deferred);
            // If the type of walker is specified, the checkpoint name should be
            // a combination of checkpoint_prefix + walker_type + repo_name.
            if let Some(walker_type) = walker_type {
                if let Some(ref mut checkpoints) = chunking.checkpoints {
                    checkpoints.checkpoint_name =
                        format!("{}_{}_{}", CHECKPOINT_PREFIX, walker_type, repo);
                }
            }
        }
        // NOTE: error_as_data_node_types is an argument that can be specified for
        // individual repos but the walker just assumes one univeral value for it even
        // when executing for multiple repos. For sharded execution, having per-repo and
        // all-repo value for error_as_data_node_types will behave the same since the entire
        // setup is done once-per-repo in sharded setting. In CLI, this behavior is enforced
        // by requiring all repos executing together to have the same value for this field.
        error_as_data_node_types_for_all_repos = walker_config_params
            .and_then(|p| p.error_as_node_data_type.as_ref())
            .map(|s| s.parse::<NodeTypeArg>())
            .transpose()?
            .map_or(error_as_data_node_types_for_all_repos, |n| {
                HashSet::<NodeType>::from_iter(n.0.iter().cloned())
            });

        let common_config = app.repo_configs().common.clone();
        let one_repo = setup_repo(
            walk_stats_key,
            app.fb,
            logger,
            &repo_factory,
            scuba_builder.clone(),
            sql_shard_info,
            scheduled_max_concurrency,
            repo_count,
            repo.clone(),
            &repo_conf,
            walk_roots.clone(),
            tail_params.clone(),
            include_edge_types.clone(),
            included_nodes,
            hash_validation_node_types.clone(),
            progress_options,
            common_config,
        )
        .await?;
        per_repo.push(one_repo);
    }

    Ok(JobParams {
        walk_params: JobWalkParams {
            enable_derive: common_args.enable_derive,
            quiet: common_args.quiet,
            error_as_data_node_types: error_as_data_node_types_for_all_repos,
            error_as_data_edge_types,
            repo_count,
        },
        per_repo,
    })
}

/// Fetch the configuration parameters specific to a particular variant
/// of the walker job.
pub fn walker_config_params<'a>(
    repo_config: &'a RepoConfig,
    job_type: &'a WalkerJobType,
) -> Option<&'a WalkerJobParams> {
    if let Some(config) = repo_config.walker_config.as_ref() {
        if let Some(params) = config.params.as_ref() {
            return params
                .iter()
                .find_map(|(k, v)| if *k == *job_type { Some(v) } else { None });
        }
    }
    None
}

// Override the blobstore config so we can do things like run on one side of a multiplex
fn override_repo_configs(
    walk_stats_key: &'static str,
    app: &MononokeApp,
    common_args: &WalkerCommonArgs,
    mut repos: Vec<(String, RepoConfig)>,
) -> Result<Vec<(String, RepoConfig)>, Error> {
    let storage_override = if let Some(storage_id) = &common_args.storage_id {
        let storage_configs = &app.storage_configs().storage;
        let config_opt = storage_configs.get(storage_id).cloned();
        let storage_config = config_opt.ok_or_else(|| {
            format_err!(
                "Storage id `{}` not found in {:?}",
                storage_id,
                storage_configs.keys()
            )
        })?;
        Some(storage_config)
    } else {
        None
    };

    for (name, config) in &mut repos {
        if let Some(storage_config) = storage_override.clone() {
            config.storage_config = storage_config;
        }
        replace_blobconfig(
            &mut config.storage_config.blobstore,
            common_args.inner_blobstore_id,
            name,
            walk_stats_key,
            app.blobstore_options().scrub_options.is_some(),
        )?;

        let sampling_multiplier =
            NonZeroU64::new(common_args.blobstore_sampling_multiplier).context("Cannot be zero")?;
        config
            .storage_config
            .blobstore
            .apply_sampling_multiplier(sampling_multiplier);
    }

    // Disable redaction unless we are running with it enabled.
    if !common_args.enable_redaction {
        for (_repo, config) in &mut repos {
            config.redaction = Redaction::Disabled;
        }
    };

    Ok(repos)
}

fn setup_repo_factory<'a>(
    walk_stats_key: &'static str,
    app: &MononokeApp,
    repo_id_to_name: HashMap<RepositoryId, String>,
    blobstore_sampler: Option<Arc<dyn SamplingHandler>>,
    blobstore_component_sampler: Option<Arc<dyn ComponentSamplingHandler>>,
    scuba_builder: MononokeScubaSampleBuilder,
    logger: &'a Logger,
    quiet: bool,
) -> RepoFactory {
    // We want to customize the repo factory, so take a deep clone
    // of the factory on the App.
    let mut repo_factory = RepoFactory::clone(app.repo_factory());
    if let Some(blobstore_sampler) = blobstore_sampler.clone() {
        repo_factory.with_blobstore_override({
            cloned!(logger);
            move |blobstore| -> Arc<dyn Blobstore> {
                if !quiet {
                    info!(logger, "Sampling from blobstore: {}", blobstore);
                }
                Arc::new(SamplingBlobstore::new(blobstore, blobstore_sampler.clone()))
            }
        });
    }

    if let Some(sampler) = blobstore_component_sampler {
        repo_factory.with_blobstore_component_sampler(sampler);
    }

    repo_factory.with_scrub_handler(Arc::new(StatsScrubHandler::new(
        false,
        scuba_builder,
        walk_stats_key,
        repo_id_to_name,
    )) as Arc<dyn ScrubHandler>);

    repo_factory
}

fn parse_tail_params(
    fb: FacebookInit,
    tail_args: &TailArgs,
    mysql_options: &MysqlOptions,
    repos: &[(String, RepoConfig)],
    walk_roots: &[OutgoingEdge],
) -> Result<HashMap<MetadataDatabaseConfig, TailParams>, Error> {
    let mut parsed_tail_params: HashMap<MetadataDatabaseConfig, TailParams> = HashMap::new();
    for (_repo, repo_conf) in repos {
        let metadatadb_config = &repo_conf.storage_config.metadata;
        let tail_params = match parsed_tail_params.get(metadatadb_config) {
            Some(tail_params) => tail_params.clone(),
            None => {
                let tail_params = tail_args.parse_args(fb, metadatadb_config, mysql_options)?;
                parsed_tail_params.insert(metadatadb_config.clone(), tail_params.clone());
                tail_params
            }
        };

        if tail_params.chunking.is_none() && walk_roots.is_empty() {
            bail!(
                "No walk roots provided, pass with  --bookmark, --walk-root or --chunk-by-public",
            );
        }
    }

    Ok(parsed_tail_params)
}

// Setup for just one repo. Try and keep clap parsing out of here, should be done beforehand
async fn setup_repo<'a>(
    walk_stats_key: &'static str,
    fb: FacebookInit,
    logger: &'a Logger,
    repo_factory: &'a RepoFactory,
    mut scuba_builder: MononokeScubaSampleBuilder,
    sql_shard_info: SqlShardInfo,
    scheduled_max: usize,
    repo_count: usize,
    repo_name: String,
    repo_config: &'a RepoConfig,
    walk_roots: Vec<OutgoingEdge>,
    mut tail_params: TailParams,
    include_edge_types: HashSet<EdgeType>,
    mut include_node_types: HashSet<NodeType>,
    hash_validation_node_types: HashSet<NodeType>,
    progress_options: ProgressOptions,
    common_config: CommonConfig,
) -> Result<(RepoSubcommandParams, RepoWalkParams), Error> {
    let logger = logger.new(o!("repo" => repo_name.clone()));

    let scheduled_max = scheduled_max / repo_count;
    scuba_builder.add(REPO, repo_name.clone());

    // Only walk derived node types that the repo is configured to contain
    include_node_types.retain(|t| {
        if let Some(t) = t.derived_data_name() {
            repo_config.derived_data_config.is_enabled(t)
        } else {
            true
        }
    });

    let mut root_node_types: HashSet<_> =
        walk_roots.iter().map(|e| e.label.outgoing_type()).collect();

    if let Some(ref mut chunking) = tail_params.chunking {
        chunking.chunk_by.retain(|t| {
            if let Some(t) = t.derived_data_name() {
                repo_config.derived_data_config.is_enabled(t)
            } else {
                true
            }
        });

        root_node_types.extend(chunking.chunk_by.iter().cloned());
    }

    let (include_edge_types, include_node_types) =
        reachable_graph_elements(include_edge_types, include_node_types, &root_node_types);
    info!(
        logger,
        #log::GRAPH,
        "Walking edge types {:?}",
        sort_by_string(&include_edge_types)
    );
    info!(
        logger,
        #log::GRAPH,
        "Walking node types {:?}",
        sort_by_string(&include_node_types)
    );

    scuba_builder.add(REPO, repo_name.clone());

    let mut progress_node_types = include_node_types.clone();
    for e in &walk_roots {
        progress_node_types.insert(e.target.get_type());
    }

    let progress_state = ProgressStateMutex::new(ProgressStateCountByType::new(
        fb,
        logger.clone(),
        walk_stats_key,
        repo_name.clone(),
        progress_node_types,
        progress_options,
    ));

    let repo: BlobRepo = repo_factory
        .build(repo_name.clone(), repo_config.clone(), common_config)
        .await?;

    Ok((
        RepoSubcommandParams {
            progress_state,
            tail_params,
            lfs_threshold: repo_config.lfs.threshold,
        },
        RepoWalkParams {
            repo,
            logger: logger.clone(),
            scheduled_max,
            sql_shard_info,
            walk_roots,
            include_node_types,
            include_edge_types,
            hash_validation_node_types,
            scuba_builder,
        },
    ))
}

fn reachable_graph_elements(
    mut include_edge_types: HashSet<EdgeType>,
    mut include_node_types: HashSet<NodeType>,
    root_node_types: &HashSet<NodeType>,
) -> (HashSet<EdgeType>, HashSet<NodeType>) {
    // This stops us logging that we're walking unreachable edge/node types
    let mut param_count = include_edge_types.len() + include_node_types.len();
    let mut last_param_count = 0;
    while param_count != last_param_count {
        let include_edge_types_stable = include_edge_types.clone();
        // Only retain edge types that are traversable
        include_edge_types.retain(|e| {
            e.incoming_type()
                .map_or(true, |t|
                    // its an incoming_type we want
                    (include_node_types.contains(&t) || root_node_types.contains(&t)) &&
                    // Another existing edge can get us to this node type
                    (root_node_types.contains(&t) || include_edge_types_stable.iter().any(|o| o.outgoing_type() == t)))
                // its an outgoing_type we want
                && include_node_types.contains(&e.outgoing_type())
        });
        // Only retain node types we expect to step to after graph entry
        include_node_types.retain(|t| {
            include_edge_types
                .iter()
                .any(|e| &e.outgoing_type() == t || e.incoming_type().map_or(false, |ot| &ot == t))
        });
        last_param_count = param_count;
        param_count = include_edge_types.len() + include_node_types.len();
    }
    (include_edge_types, include_node_types)
}


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

use std::collections::HashSet;

use clap::ValueEnum;
use once_cell::sync::Lazy;
use strum::AsRefStr;
use strum::EnumString;
use strum::EnumVariantNames;
use strum::IntoEnumIterator;

use crate::detail::graph::NodeType;
use crate::detail::state::InternedType;

#[derive(Debug, Clone, Copy, ValueEnum, AsRefStr, EnumString, EnumVariantNames)]
// Forcing backward compatibility with clap-3 for user facing CLI arguments
#[clap(rename_all = "PascalCase")]
pub enum InternedTypeArg {
    All,
    FileUnodeId,
    HgChangesetId,
    HgFileNodeId,
    HgManifestId,
    ManifestUnodeId,
    MPathHash,
}

impl InternedTypeArg {
    pub fn parse_args(args: &[Self]) -> HashSet<InternedType> {
        let mut int_types = HashSet::new();
        for arg in args {
            match *arg {
                InternedTypeArg::FileUnodeId => {
                    int_types.insert(InternedType::FileUnodeId);
                }
                InternedTypeArg::HgChangesetId => {
                    int_types.insert(InternedType::HgChangesetId);
                }
                InternedTypeArg::HgFileNodeId => {
                    int_types.insert(InternedType::HgFileNodeId);
                }
                InternedTypeArg::HgManifestId => {
                    int_types.insert(InternedType::HgManifestId);
                }
                InternedTypeArg::ManifestUnodeId => {
                    int_types.insert(InternedType::ManifestUnodeId);
                }
                InternedTypeArg::MPathHash => {
                    int_types.insert(InternedType::MPathHash);
                }
                InternedTypeArg::All => {
                    int_types.extend(InternedType::iter());
                }
            }
        }
        int_types
    }
}

/// Default to clearing out all except HgChangesets
pub const DEFAULT_INTERNED_TYPES: &[InternedTypeArg] = &[
    InternedTypeArg::FileUnodeId,
    InternedTypeArg::HgFileNodeId,
    InternedTypeArg::HgManifestId,
    InternedTypeArg::ManifestUnodeId,
    InternedTypeArg::MPathHash,
];

// clap doesn't allow to pass typed default values for some reason, let's convert them
pub static DEFAULT_INTERNED_TYPES_STR: Lazy<Vec<&'static str>> = Lazy::new(|| {
    DEFAULT_INTERNED_TYPES
        .iter()
        .map(|int_type| int_type.as_ref())
        .collect()
});

/// We can jump for ChangesetId to all of these
#[derive(Debug, Clone, Copy, ValueEnum, AsRefStr, EnumString, EnumVariantNames)]
// Forcing backward compatibility with clap-3 for user facing CLI arguments
#[clap(rename_all = "PascalCase")]
pub enum ChunkByPublicArg {
    BonsaiHgMapping,
    PhaseMapping,
    Changeset,
    ChangesetInfo,
    ChangesetInfoMapping,
    DeletedManifestV2Mapping,
    FsnodeMapping,
    SkeletonManifestMapping,
    BasenameSuffixSkeletonManifestMapping,
    UnodeMapping,
}

impl From<ChunkByPublicArg> for NodeType {
    fn from(arg: ChunkByPublicArg) -> NodeType {
        match arg {
            ChunkByPublicArg::BonsaiHgMapping => NodeType::BonsaiHgMapping,
            ChunkByPublicArg::PhaseMapping => NodeType::PhaseMapping,
            ChunkByPublicArg::Changeset => NodeType::Changeset,
            ChunkByPublicArg::ChangesetInfo => NodeType::ChangesetInfo,
            ChunkByPublicArg::ChangesetInfoMapping => NodeType::ChangesetInfoMapping,
            ChunkByPublicArg::DeletedManifestV2Mapping => NodeType::DeletedManifestV2Mapping,
            ChunkByPublicArg::FsnodeMapping => NodeType::FsnodeMapping,
            ChunkByPublicArg::SkeletonManifestMapping => NodeType::SkeletonManifestMapping,
            ChunkByPublicArg::BasenameSuffixSkeletonManifestMapping => {
                NodeType::BasenameSuffixSkeletonManifestMapping
            }
            ChunkByPublicArg::UnodeMapping => NodeType::UnodeMapping,
        }
    }
}

impl ChunkByPublicArg {
    pub fn parse_args(args: &[Self]) -> HashSet<NodeType> {
        args.iter().cloned().map(NodeType::from).collect()
    }
}


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

use std::collections::HashMap;
use std::collections::HashSet;
use std::hash::Hash;
use std::str::FromStr;

use anyhow::format_err;
use anyhow::Context as _;
use anyhow::Error;
use derived_data_manager::derivable::BonsaiDerivable;
use filenodes_derivation::FilenodesOnlyPublic;
use mercurial_derivation::MappedHgChangesetId;
use once_cell::sync::Lazy;
use strum::IntoEnumIterator;

use crate::detail::graph::EdgeType;
use crate::detail::graph::NodeType;

const ALL: &str = "all";
const BONSAI: &str = "bonsai";
pub const DEFAULT: &str = "default";
const DERIVED: &str = "derived";
const HG: &str = "hg";
pub const DEEP: &str = "deep";
const SHALLOW: &str = "shallow";
const CONTENTMETA: &str = "contentmeta";
const MARKER: &str = "marker";

/* NodeType */

const DEFAULT_INCLUDE_NODE_TYPES: &[NodeType] = &[
    NodeType::Bookmark,
    NodeType::Changeset,
    NodeType::BonsaiHgMapping,
    NodeType::PhaseMapping,
    NodeType::PublishedBookmarks,
    NodeType::HgBonsaiMapping,
    NodeType::HgChangeset,
    NodeType::HgChangesetViaBonsai,
    NodeType::HgManifest,
    NodeType::HgFileEnvelope,
    NodeType::HgFileNode,
    NodeType::FileContent,
    NodeType::FileContentMetadataV2,
    NodeType::AliasContentMapping,
];

const BONSAI_NODE_TYPES: &[NodeType] = &[NodeType::Bookmark, NodeType::Changeset];
const HG_DERIVED_TYPES: &[&str] = &[MappedHgChangesetId::NAME, FilenodesOnlyPublic::NAME];

const DERIVED_PREFIX: &str = "derived_";

static DERIVED_DATA_NODE_TYPES: Lazy<HashMap<String, Vec<NodeType>>> = Lazy::new(|| {
    let mut m: HashMap<String, Vec<NodeType>> = HashMap::new();
    for t in NodeType::iter() {
        if let Some(name) = t.derived_data_name() {
            m.entry(format!("{}{}", DERIVED_PREFIX, name))
                .or_default()
                .push(t);
        }
    }
    m
});

pub type NodeTypeArg = GraphTypeArg<NodeType>;

impl FromStr for NodeTypeArg {
    type Err = Error;

    fn from_str(arg: &str) -> Result<NodeTypeArg, Error> {
        Ok(match arg {
            ALL => GraphTypeArg(NodeType::iter().collect()),
            BONSAI => NodeTypeArg::new(BONSAI_NODE_TYPES.iter()),
            DEFAULT => NodeTypeArg::new(DEFAULT_INCLUDE_NODE_TYPES.iter()),
            DERIVED => NodeTypeArg::new(DERIVED_DATA_NODE_TYPES.values().flatten()),
            HG => {
                let mut node_types = vec![];
                for hg_derived in HG_DERIVED_TYPES {
                    let hg_derived = format!("{}{}", DERIVED_PREFIX, hg_derived);
                    let nodes_derived = DERIVED_DATA_NODE_TYPES.get(&hg_derived);
                    if let Some(nd) = nodes_derived {
                        nd.iter().for_each(|node| node_types.push(node.clone()));
                    }
                }
                GraphTypeArg(node_types)
            }
            _ => {
                if let Some(node_types) = DERIVED_DATA_NODE_TYPES.get(arg) {
                    GraphTypeArg(node_types.clone())
                } else {
                    NodeType::from_str(arg)
                        .map(|e| GraphTypeArg(vec![e]))
                        .with_context(|| format_err!("Unknown NodeType {}", arg))?
                }
            }
        })
    }
}

/* EdgeType */

// Goes as far into history as it can
const DEEP_INCLUDE_EDGE_TYPES: &[EdgeType] = &[
    // Bonsai
    EdgeType::BookmarkToChangeset,
    EdgeType::ChangesetToFileContent,
    EdgeType::ChangesetToBonsaiParent,
    EdgeType::ChangesetToBonsaiHgMapping,
    EdgeType::BonsaiHgMappingToHgChangesetViaBonsai,
    EdgeType::PublishedBookmarksToChangeset,
    EdgeType::PublishedBookmarksToBonsaiHgMapping,
    EdgeType::ChangesetToChangesetInfoMapping,
    EdgeType::ChangesetToDeletedManifestV2Mapping,
    EdgeType::ChangesetToFsnodeMapping,
    EdgeType::ChangesetToSkeletonManifestMapping,
    EdgeType::ChangesetToBasenameSuffixSkeletonManifestMapping,
    EdgeType::ChangesetToUnodeMapping,
    // Hg
    EdgeType::HgBonsaiMappingToChangeset,
    EdgeType::HgChangesetToHgParent,
    EdgeType::HgChangesetToHgManifest,
    EdgeType::HgChangesetToHgManifestFileNode,
    EdgeType::HgChangesetViaBonsaiToHgChangeset,
    EdgeType::HgManifestToHgFileEnvelope,
    EdgeType::HgManifestToHgFileNode,
    EdgeType::HgManifestToChildHgManifest,
    EdgeType::HgFileEnvelopeToFileContent,
    EdgeType::HgFileNodeToLinkedHgBonsaiMapping,
    EdgeType::HgFileNodeToLinkedHgChangeset,
    EdgeType::HgFileNodeToHgParentFileNode,
    EdgeType::HgFileNodeToHgCopyfromFileNode,
    EdgeType::HgManifestFileNodeToLinkedHgBonsaiMapping,
    EdgeType::HgManifestFileNodeToLinkedHgChangeset,
    EdgeType::HgManifestFileNodeToHgParentFileNode,
    EdgeType::HgManifestFileNodeToHgCopyfromFileNode,
    // Content
    EdgeType::FileContentToFileContentMetadataV2,
    EdgeType::FileContentMetadataV2ToSha1Alias,
    EdgeType::FileContentMetadataV2ToSha256Alias,
    EdgeType::FileContentMetadataV2ToGitSha1Alias,
    EdgeType::FileContentMetadataV2ToSeededBlake3Alias,
    EdgeType::AliasContentMappingToFileContent,
    // Derived data
    EdgeType::BlameToChangeset,
    EdgeType::ChangesetInfoMappingToChangesetInfo,
    EdgeType::ChangesetInfoToChangesetInfoParent,
    EdgeType::DeletedManifestV2MappingToRootDeletedManifestV2,
    EdgeType::DeletedManifestV2ToDeletedManifestV2Child,
    EdgeType::DeletedManifestV2ToLinkedChangeset,
    EdgeType::FastlogBatchToChangeset,
    EdgeType::FastlogBatchToPreviousBatch,
    EdgeType::FastlogDirToChangeset,
    EdgeType::FastlogDirToPreviousBatch,
    EdgeType::FastlogFileToChangeset,
    EdgeType::FastlogFileToPreviousBatch,
    EdgeType::FsnodeMappingToRootFsnode,
    EdgeType::FsnodeToChildFsnode,
    EdgeType::FsnodeToFileContent,
    EdgeType::SkeletonManifestMappingToRootSkeletonManifest,
    EdgeType::SkeletonManifestToSkeletonManifestChild,
    EdgeType::BasenameSuffixSkeletonManifestMappingToRootBasenameSuffixSkeletonManifest,
    EdgeType::BasenameSuffixSkeletonManifestToBasenameSuffixSkeletonManifestChild,
    EdgeType::UnodeFileToBlame,
    EdgeType::UnodeFileToFastlogFile,
    EdgeType::UnodeFileToFileContent,
    EdgeType::UnodeFileToLinkedChangeset,
    EdgeType::UnodeFileToUnodeFileParent,
    EdgeType::UnodeManifestToFastlogDir,
    EdgeType::UnodeManifestToLinkedChangeset,
    EdgeType::UnodeManifestToUnodeManifestParent,
    EdgeType::UnodeManifestToUnodeFileChild,
    EdgeType::UnodeManifestToUnodeManifestChild,
    EdgeType::UnodeMappingToRootUnodeManifest,
];

// Does not recurse into history, edges to parents excluded
const SHALLOW_INCLUDE_EDGE_TYPES: &[EdgeType] = &[
    // Bonsai
    EdgeType::BookmarkToChangeset,
    EdgeType::ChangesetToFileContent,
    EdgeType::ChangesetToBonsaiHgMapping,
    EdgeType::BonsaiHgMappingToHgChangesetViaBonsai,
    EdgeType::PublishedBookmarksToChangeset,
    EdgeType::PublishedBookmarksToBonsaiHgMapping,
    EdgeType::ChangesetToChangesetInfoMapping,
    EdgeType::ChangesetToDeletedManifestV2Mapping,
    EdgeType::ChangesetToFsnodeMapping,
    EdgeType::ChangesetToSkeletonManifestMapping,
    EdgeType::ChangesetToBasenameSuffixSkeletonManifestMapping,
    EdgeType::ChangesetToUnodeMapping,
    // Hg
    EdgeType::HgBonsaiMappingToChangeset,
    EdgeType::HgChangesetToHgManifest,
    EdgeType::HgChangesetToHgManifestFileNode,
    EdgeType::HgChangesetViaBonsaiToHgChangeset,
    EdgeType::HgManifestToHgFileEnvelope,
    EdgeType::HgManifestToHgFileNode,
    EdgeType::HgManifestToHgManifestFileNode,
    EdgeType::HgManifestToChildHgManifest,
    EdgeType::HgFileEnvelopeToFileContent,
    // Content
    EdgeType::FileContentToFileContentMetadataV2,
    EdgeType::FileContentMetadataV2ToSha1Alias,
    EdgeType::FileContentMetadataV2ToSha256Alias,
    EdgeType::FileContentMetadataV2ToGitSha1Alias,
    EdgeType::FileContentMetadataV2ToSeededBlake3Alias,
    EdgeType::AliasContentMappingToFileContent,
    // Derived data
    EdgeType::ChangesetInfoMappingToChangesetInfo,
    EdgeType::DeletedManifestV2MappingToRootDeletedManifestV2,
    EdgeType::DeletedManifestV2ToDeletedManifestV2Child,
    EdgeType::FastlogBatchToPreviousBatch,
    EdgeType::FastlogDirToPreviousBatch,
    EdgeType::FastlogFileToPreviousBatch,
    EdgeType::FsnodeToChildFsnode,
    EdgeType::FsnodeToFileContent,
    EdgeType::FsnodeMappingToRootFsnode,
    EdgeType::SkeletonManifestMappingToRootSkeletonManifest,
    EdgeType::SkeletonManifestToSkeletonManifestChild,
    EdgeType::BasenameSuffixSkeletonManifestMappingToRootBasenameSuffixSkeletonManifest,
    EdgeType::BasenameSuffixSkeletonManifestToBasenameSuffixSkeletonManifestChild,
    EdgeType::UnodeFileToBlame,
    EdgeType::UnodeFileToFastlogFile,
    EdgeType::UnodeFileToFileContent,
    EdgeType::UnodeManifestToFastlogDir,
    EdgeType::UnodeManifestToUnodeFileChild,
    EdgeType::UnodeManifestToUnodeManifestChild,
    EdgeType::UnodeMappingToRootUnodeManifest,
];

// Types that can result in loading hg data.  Useful for excludes.
const HG_EDGE_TYPES: &[EdgeType] = &[
    // Bonsai to Hg
    EdgeType::BookmarkToBonsaiHgMapping,
    EdgeType::BonsaiHgMappingToHgChangesetViaBonsai,
    EdgeType::PublishedBookmarksToBonsaiHgMapping,
    // Hg
    EdgeType::HgChangesetToHgManifest,
    EdgeType::HgChangesetToHgManifestFileNode,
    EdgeType::HgChangesetToHgParent,
    EdgeType::HgChangesetViaBonsaiToHgChangeset,
    EdgeType::HgManifestToHgFileEnvelope,
    EdgeType::HgManifestToHgFileNode,
    EdgeType::HgManifestToChildHgManifest,
    EdgeType::HgFileEnvelopeToFileContent,
    EdgeType::HgFileNodeToLinkedHgChangeset,
    EdgeType::HgFileNodeToHgParentFileNode,
    EdgeType::HgFileNodeToHgCopyfromFileNode,
    EdgeType::HgManifestFileNodeToLinkedHgChangeset,
    EdgeType::HgManifestFileNodeToHgParentFileNode,
    EdgeType::HgManifestFileNodeToHgCopyfromFileNode,
];

// Types that can result in loading bonsai data
const BONSAI_EDGE_TYPES: &[EdgeType] = &[
    // Bonsai
    EdgeType::BookmarkToChangeset,
    EdgeType::ChangesetToFileContent,
    EdgeType::ChangesetToBonsaiParent,
    EdgeType::PublishedBookmarksToChangeset,
];

const CONTENT_META_EDGE_TYPES: &[EdgeType] = &[
    // Content
    EdgeType::FileContentToFileContentMetadataV2,
    EdgeType::FileContentMetadataV2ToSha1Alias,
    EdgeType::FileContentMetadataV2ToSha256Alias,
    EdgeType::FileContentMetadataV2ToGitSha1Alias,
    EdgeType::FileContentMetadataV2ToSeededBlake3Alias,
    EdgeType::AliasContentMappingToFileContent,
];

// Things like phases and obs markers will go here
const MARKER_EDGE_TYPES: &[EdgeType] = &[EdgeType::ChangesetToPhaseMapping];

pub type EdgeTypeArg = GraphTypeArg<EdgeType>;

impl FromStr for EdgeTypeArg {
    type Err = Error;

    fn from_str(arg: &str) -> Result<EdgeTypeArg, Error> {
        Ok(match arg {
            ALL => GraphTypeArg(EdgeType::iter().collect()),
            DEEP => EdgeTypeArg::new(DEEP_INCLUDE_EDGE_TYPES.iter()),
            SHALLOW => EdgeTypeArg::new(SHALLOW_INCLUDE_EDGE_TYPES.iter()),
            CONTENTMETA => EdgeTypeArg::new(CONTENT_META_EDGE_TYPES.iter()),
            MARKER => EdgeTypeArg::new(MARKER_EDGE_TYPES.iter()),
            BONSAI => EdgeTypeArg::new(BONSAI_EDGE_TYPES.iter()),
            HG => EdgeTypeArg::new(HG_EDGE_TYPES.iter()),
            _ => EdgeType::from_str(arg)
                .map(|e| GraphTypeArg(vec![e]))
                .with_context(|| format_err!("Unknown EdgeType {}", arg))?,
        })
    }
}

#[derive(Debug, Clone)]
pub struct GraphTypeArg<T>(pub Vec<T>);

impl<'a, T: 'a + Clone + Eq + Hash> GraphTypeArg<T> {
    pub fn new(it: impl Iterator<Item = &'a T>) -> Self {
        GraphTypeArg(it.cloned().collect())
    }

    pub fn parse_args(args: &[Self]) -> HashSet<T> {
        HashSet::from_iter(args.iter().flat_map(|arg| arg.0.clone()))
    }

    pub fn filter(include: &[Self], exclude: &[Self]) -> HashSet<T> {
        let mut include = Self::parse_args(include);
        let exclude = Self::parse_args(exclude);
        include.retain(|x| !exclude.contains(x));
        include
    }
}


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

use std::collections::HashSet;

use clap::Args;
use clap::ValueEnum;
use strum::AsRefStr;
use strum::EnumString;
use strum::EnumVariantNames;

use crate::detail::graph::NodeType;

#[derive(Args, Debug)]
pub struct HashValidationArgs {
    /// Node types for which we don't want to do hash validation.
    #[clap(long)]
    pub exclude_hash_validation_node_type: Vec<HashValidationArg>,
    /// Node types for which we want to do hash validation.
    #[clap(long)]
    pub include_hash_validation_node_type: Vec<HashValidationArg>,
}

impl HashValidationArgs {
    pub fn parse_args(&self) -> HashSet<NodeType> {
        let mut include_types =
            HashValidationArg::parse_args(&self.include_hash_validation_node_type);
        let exclude_types = HashValidationArg::parse_args(&self.exclude_hash_validation_node_type);
        include_types.retain(|x| !exclude_types.contains(x));
        include_types
    }
}

#[derive(Debug, Clone, Copy, ValueEnum, AsRefStr, EnumString, EnumVariantNames)]
// Forcing backward compatibility with clap-3 for user facing CLI arguments
#[clap(rename_all = "PascalCase")]
pub enum HashValidationArg {
    HgFileEnvelope,
}

impl HashValidationArg {
    pub fn parse_args(args: &[Self]) -> HashSet<NodeType> {
        args.iter().cloned().map(NodeType::from).collect()
    }
}

impl From<HashValidationArg> for NodeType {
    fn from(value: HashValidationArg) -> NodeType {
        match value {
            HashValidationArg::HgFileEnvelope => NodeType::HgFileEnvelope,
        }
    }
}


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

pub mod arg_types;
mod graph_arg_types;
pub mod hash_validation;
pub mod progress;
pub mod sampling;
pub mod scrub;
pub mod tail_args;
pub mod validate;
pub mod walk_params;
pub mod walk_root;

use clap::Args;
pub use graph_arg_types::NodeTypeArg;
pub use hash_validation::HashValidationArgs;
pub use progress::ProgressArgs;
pub use sampling::SamplingArgs;
pub use scrub::ScrubOutputNodeArgs;
pub use scrub::ScrubPackLogArgs;
use strum::AsRefStr;
use strum::EnumString;
use strum::EnumVariantNames;
pub use tail_args::CheckpointArgs;
pub use tail_args::ChunkingArgs;
pub use tail_args::TailArgs;
pub use validate::ValidateCheckTypeArgs;
pub use walk_params::WalkerGraphArgs;
pub use walk_params::WalkerGraphParams;
pub use walk_root::WalkRootArgs;

#[derive(Args, Debug)]
pub struct WalkerCommonArgs {
    /// Log a lot less
    #[clap(long, short = 'q')]
    pub quiet: bool,
    /// Use redaction from config. Default is redaction off.
    #[clap(long)]
    pub enable_redaction: bool,
    /// Maximum number of walk step tasks to attempt to execute at once.
    #[clap(long, default_value = "4096")]
    pub scheduled_max: usize,
    /// Enable derivation of data (e.g. hg, file metadata).
    #[clap(long)]
    pub enable_derive: bool,
    /// Limit the amount of data fetched from stores, by not streaming
    /// large files to the end. Only used by `scrub` subcommand.
    #[clap(long)]
    pub limit_data_fetch: bool,

    /// Id of a storage group to operate over, e.g. manifold_xdb_multiplex
    #[clap(long)]
    pub storage_id: Option<String>,
    /// If main blobstore in the storage config is a multiplexed one,
    /// use inner blobstore with this id.
    #[clap(long)]
    pub inner_blobstore_id: Option<u64>,
    /// Add a multiplier on sampling requests
    #[clap(long, default_value = "100")]
    pub blobstore_sampling_multiplier: u64,

    #[clap(flatten, next_help_heading = "WALKING ROOTS")]
    pub walk_roots: WalkRootArgs,
    #[clap(flatten, next_help_heading = "GRAPH OPTIONS")]
    pub graph_params: WalkerGraphArgs,
    #[clap(flatten, next_help_heading = "HASH VALIDATION OPTIONS")]
    pub hash_validation: HashValidationArgs,
    #[clap(flatten, next_help_heading = "PROGRESS OPTIONS")]
    pub progress: ProgressArgs,
    #[clap(flatten, next_help_heading = "TAILING OPTIONS")]
    pub tailing: TailArgs,
}

#[derive(Clone, Debug, PartialEq, Eq, AsRefStr, EnumVariantNames, EnumString)]
pub enum OutputFormat {
    Debug,
    PrettyDebug,
}


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

use clap::Args;

use crate::detail::progress::ProgressOptions;

#[derive(Args, Debug)]
pub struct ProgressArgs {
    /// Minimum interval between progress reports in seconds.
    #[clap(long, default_value_t = 5)]
    pub progress_interval: u64,
    /// Sample the walk output stream for progress roughly 1 in N steps.
    /// Only log if progress-interval has passed.
    #[clap(long, default_value_t = 100)]
    pub progress_sample_rate: u64,
}

impl ProgressArgs {
    pub fn parse_args(&self) -> ProgressOptions {
        ProgressOptions {
            sample_rate: self.progress_sample_rate,
            interval: Duration::from_secs(self.progress_interval),
        }
    }
}


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

use anyhow::Error;
use clap::Args;
use regex::Regex;

use crate::args::graph_arg_types::NodeTypeArg;
use crate::detail::sampling::SamplingOptions;

#[derive(Args, Debug)]
pub struct SamplingArgs {
    /// Pass 1 to try all nodes, 120 to do 1 in 120, etc.
    #[clap(long)]
    pub sample_rate: Option<u64>,
    /// Offset to apply to the sampling fingerprint for each node, can be used
    /// to cycle through an entire repo in N pieces.
    #[clap(long, default_value = "0")]
    pub sample_offset: u64,
    /// Node types to exclude from the sample.
    #[clap(long, short = 'S')]
    pub exclude_sample_node_type: Vec<NodeTypeArg>,
    /// Node types to include in the sample, defaults to same as the walk.
    #[clap(long, short = 's')]
    pub include_sample_node_type: Vec<NodeTypeArg>,
    /// If provided, only sample paths that match.
    #[clap(long)]
    pub sample_path_regex: Option<Regex>,
}

impl SamplingArgs {
    pub fn parse_args(&self, default_sample_rate: u64) -> Result<SamplingOptions, Error> {
        let sample_rate = self.sample_rate.clone().unwrap_or(default_sample_rate);
        let node_types = NodeTypeArg::filter(
            &self.include_sample_node_type,
            &self.exclude_sample_node_type,
        );
        let exclude_types = NodeTypeArg::parse_args(&self.exclude_sample_node_type);
        Ok(SamplingOptions {
            sample_rate,
            sample_offset: self.sample_offset,
            node_types,
            exclude_types,
        })
    }
}


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

use std::collections::HashSet;

use anyhow::Error;
use clap::Args;
use fbinit::FacebookInit;
use scuba_ext::MononokeScubaSampleBuilder;

use crate::args::graph_arg_types::NodeTypeArg;
use crate::detail::graph::NodeType;
use crate::detail::pack::PackInfoLogOptions;

#[derive(Args, Debug)]
pub struct ScrubOutputNodeArgs {
    /// Node types not to output in debug stdout
    #[clap(long, short = 'O')]
    pub exclude_output_node_type: Vec<NodeTypeArg>,
    /// Node types to output in debug stdout
    #[clap(long, short = 'o')]
    pub include_output_node_type: Vec<NodeTypeArg>,
}

impl ScrubOutputNodeArgs {
    pub fn parse_args(&self) -> HashSet<NodeType> {
        NodeTypeArg::filter(
            &self.include_output_node_type,
            &self.exclude_output_node_type,
        )
    }
}

#[derive(Args, Debug)]
pub struct ScrubPackLogArgs {
    /// Node types not to log pack info for
    #[clap(long, short = 'A')]
    pub exclude_pack_log_node_type: Vec<NodeTypeArg>,
    /// Node types to log pack info for
    #[clap(long, short = 'a')]
    pub include_pack_log_node_type: Vec<NodeTypeArg>,
    /// Scuba table for logging pack info data to. e.g. mononoke_packinfo
    #[clap(long, requires = "include_pack_log_node_type")]
    pub pack_log_scuba_table: Option<String>,
    /// A log file to write Scuba pack info logs to (primarily useful in testing)
    #[clap(long, requires = "include_pack_log_node_type")]
    pub pack_log_scuba_file: Option<String>,
}

impl ScrubPackLogArgs {
    pub fn parse_args(&self, fb: FacebookInit) -> Result<Option<PackInfoLogOptions>, Error> {
        let log_node_types = NodeTypeArg::filter(
            &self.include_pack_log_node_type,
            &self.exclude_pack_log_node_type,
        );

        if !log_node_types.is_empty() {
            let mut scuba_builder =
                MononokeScubaSampleBuilder::with_opt_table(fb, self.pack_log_scuba_table.clone())?;
            if let Some(scuba_log_file) = &self.pack_log_scuba_file {
                scuba_builder = scuba_builder.with_log_file(scuba_log_file)?;
            }

            return Ok(Some(PackInfoLogOptions {
                log_node_types,
                log_dest: scuba_builder,
            }));
        }

        Ok(None)
    }
}


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

use anyhow::Error;
use bulkops::Direction;
use clap::Args;
use fbinit::FacebookInit;
use metaconfig_types::MetadataDatabaseConfig;
use sql_construct::SqlConstruct;
use sql_construct::SqlConstructFromMetadataDatabaseConfig;
use sql_ext::facebook::MysqlOptions;

use crate::args::arg_types::ChunkByPublicArg;
use crate::args::arg_types::InternedTypeArg;
use crate::args::arg_types::DEFAULT_INTERNED_TYPES_STR;
use crate::args::graph_arg_types::NodeTypeArg;
use crate::detail::checkpoint::CheckpointsByName;
use crate::detail::checkpoint::SqlCheckpoints;
use crate::detail::tail::ChunkingParams;
use crate::detail::tail::ClearStateParams;
use crate::detail::tail::TailParams;

#[derive(Args, Debug)]
pub struct TailArgs {
    /// Max age of walk state held internally ot loaded from checkpoint that
    /// we will attempt to continue from, in seconds. Default is set to 5 days.
    // 5 days = 5 * 24 * 3600 seconds = 432000
    #[clap(long, default_value = "432000")]
    pub state_max_age: u64,
    /// Tail by polling the entry points at interval of TAIL seconds.
    #[clap(long)]
    pub tail_interval: Option<u64>,

    #[clap(flatten)]
    pub chunking: ChunkingArgs,
}

impl TailArgs {
    pub fn parse_args(
        &self,
        fb: FacebookInit,
        dbconfig: &MetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
    ) -> Result<TailParams, Error> {
        Ok(TailParams {
            tail_secs: self.tail_interval.clone(),
            chunking: self.chunking.parse_args(fb, dbconfig, mysql_options)?,
            state_max_age: Duration::from_secs(self.state_max_age),
        })
    }
}

#[derive(Args, Debug)]
pub struct ChunkingArgs {
    /// Traverse using chunks of public changesets as roots to the specified node type
    #[clap(long, short = 'p')]
    pub chunk_by_public: Vec<ChunkByPublicArg>,
    /// Set the direction to proceed through changesets
    #[clap(long, short = 'd', requires = "chunk_by_public")]
    pub chunk_direction: Option<Direction>,
    /// How many changesets to include in a chunk.
    #[clap(long, short = 'k', default_value = "100000")]
    pub chunk_size: usize,
    /// Clear the saved walk state 1 in N steps.
    #[clap(long, short = 'K')]
    pub chunk_clear_sample_rate: Option<u64>,
    /// Whether to allow remaining deferred edges after chunks complete.
    /// Well structured repos should have none.
    #[clap(long)]
    pub allow_remaining_deferred: bool,

    /// Include in InternedTypes to flush between chunks
    #[clap(long, short = 't', default_values = &*DEFAULT_INTERNED_TYPES_STR)]
    pub include_chunk_clear_interned_type: Vec<InternedTypeArg>,
    /// Exclude from InternedTypes to flush between chunks
    #[clap(long, short = 'T')]
    pub exclude_chunk_clear_interned_type: Vec<InternedTypeArg>,

    /// Include in NodeTypes to flush between chunks
    #[clap(long, short = 'n')]
    pub include_chunk_clear_node_type: Vec<NodeTypeArg>,
    /// Exclude from NodeTypes to flush between chunks
    #[clap(long, short = 'N')]
    pub exclude_chunk_clear_node_type: Vec<NodeTypeArg>,

    /// Set the repo upper bound used by chunking instead of loading it.
    /// Inclusive. Useful for reproducing issues from a particular chunk.
    #[clap(long, requires = "chunk_by_public")]
    pub repo_lower_bound: Option<u64>,
    /// Set the repo lower bound used by chunking instead of loading it.
    /// Exclusive (used in rust ranges). Useful for reproducing issues
    /// from a particular chunk.
    #[clap(long, requires = "chunk_by_public")]
    pub repo_upper_bound: Option<u64>,

    #[clap(flatten)]
    pub checkpoint: CheckpointArgs,
}

impl ChunkingArgs {
    pub fn parse_args(
        &self,
        fb: FacebookInit,
        dbconfig: &MetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
    ) -> Result<Option<ChunkingParams>, Error> {
        if self.chunk_by_public.is_empty() {
            return Ok(None);
        }

        let clear_state = if let Some(sample_rate) = &self.chunk_clear_sample_rate {
            let mut include_int_types =
                InternedTypeArg::parse_args(&self.include_chunk_clear_interned_type);
            let exclude_int_types =
                InternedTypeArg::parse_args(&self.exclude_chunk_clear_interned_type);
            include_int_types.retain(|x| !exclude_int_types.contains(x));

            let include_nodes = NodeTypeArg::filter(
                &self.include_chunk_clear_node_type,
                &self.exclude_chunk_clear_node_type,
            );

            Some(ClearStateParams {
                sample_rate: *sample_rate,
                interned_types: include_int_types,
                node_types: include_nodes,
            })
        } else {
            None
        };

        let direction = self
            .chunk_direction
            .clone()
            .unwrap_or(Direction::NewestFirst);

        Ok(Some(ChunkingParams {
            chunk_by: ChunkByPublicArg::parse_args(&self.chunk_by_public),
            chunk_size: self.chunk_size,
            direction,
            clear_state,
            checkpoints: self.checkpoint.parse_args(fb, dbconfig, mysql_options)?,
            allow_remaining_deferred: self.allow_remaining_deferred,
            repo_lower_bound_override: self.repo_lower_bound,
            repo_upper_bound_override: self.repo_upper_bound,
        }))
    }
}

#[derive(Args, Debug)]
pub struct CheckpointArgs {
    /// Name of checkpoint.
    #[clap(long)]
    pub checkpoint_name: Option<String>,
    /// Path for sqlite checkpoint db if using sqlite
    #[clap(long, requires = "checkpoint_name")]
    pub checkpoint_path: Option<String>,
    /// Checkpoint the walk covered bounds 1 in N steps.
    #[clap(long, default_value = "1")]
    pub checkpoint_sample_rate: u64,
}

impl CheckpointArgs {
    pub fn parse_args(
        &self,
        fb: FacebookInit,
        dbconfig: &MetadataDatabaseConfig,
        mysql_options: &MysqlOptions,
    ) -> Result<Option<CheckpointsByName>, Error> {
        if let Some(checkpoint_name) = &self.checkpoint_name {
            let sql_checkpoints = if let Some(checkpoint_path) = &self.checkpoint_path {
                SqlCheckpoints::with_sqlite_path(checkpoint_path, false)?
            } else {
                SqlCheckpoints::with_metadata_database_config(fb, dbconfig, mysql_options, false)?
            };

            Ok(Some(CheckpointsByName::new(
                checkpoint_name.clone(),
                sql_checkpoints,
                self.checkpoint_sample_rate,
            )))
        } else {
            Ok(None)
        }
    }
}


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

use std::collections::HashSet;

use clap::Args;
use clap::ValueEnum;
use strum::AsRefStr;
use strum::EnumString;
use strum::EnumVariantNames;

use crate::detail::validate::CheckType;
use crate::detail::validate::DEFAULT_CHECK_TYPES;

#[derive(Args, Debug)]
pub struct ValidateCheckTypeArgs {
    /// Check types to exclude
    #[clap(long, short = 'C')]
    pub exclude_check_type: Vec<CheckTypeArg>,
    /// Check types to include
    // TODO(aida): After the D26853691 it doesn't actually default to all
    // possible values, `CheckType::FileContentIsLfs` is missing.
    // We need to fix this (then we can also get rid of DEFAULT_CHECK_TYPES).
    #[clap(long, short = 'c', default_values = &[CheckTypeArg::All.as_ref()])]
    pub include_check_type: Vec<CheckTypeArg>,
}

impl ValidateCheckTypeArgs {
    pub fn parse_args(&self) -> HashSet<CheckType> {
        let mut include_types = parse_check_type_args(&self.include_check_type);
        let exclude_types = parse_check_type_args(&self.exclude_check_type);
        include_types.retain(|x| !exclude_types.contains(x));
        include_types
    }
}

#[derive(Debug, Clone, Copy, ValueEnum, AsRefStr, EnumString, EnumVariantNames)]
// Forcing backward compatibility with clap-3 for user facing CLI arguments
#[clap(rename_all = "PascalCase")]
pub enum CheckTypeArg {
    All,
    ChangesetPhaseIsPublic,
    HgLinkNodePopulated,
    FileContentIsLfs,
}

fn parse_check_type_args(check_type_args: &[CheckTypeArg]) -> HashSet<CheckType> {
    let mut check_types = HashSet::new();
    for arg in check_type_args {
        match arg {
            CheckTypeArg::All => {
                for default_type in DEFAULT_CHECK_TYPES {
                    check_types.insert(default_type.clone());
                }
            }
            CheckTypeArg::ChangesetPhaseIsPublic => {
                check_types.insert(CheckType::ChangesetPhaseIsPublic);
            }
            CheckTypeArg::HgLinkNodePopulated => {
                check_types.insert(CheckType::HgLinkNodePopulated);
            }
            CheckTypeArg::FileContentIsLfs => {
                check_types.insert(CheckType::FileContentIsLfs);
            }
        }
    }
    check_types
}


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

use std::collections::HashSet;

use anyhow::Error;
use clap::Args;

use crate::args::graph_arg_types::EdgeTypeArg;
use crate::args::graph_arg_types::NodeTypeArg;
use crate::args::graph_arg_types::DEEP;
use crate::args::graph_arg_types::DEFAULT;
use crate::detail::graph::EdgeType;
use crate::detail::graph::NodeType;

#[derive(Args, Debug)]
pub struct WalkerGraphArgs {
    /// Graph node types to exclude from walk. They are removed from
    /// the include node types.
    #[clap(long, short = 'x')]
    pub exclude_node_type: Vec<NodeTypeArg>,
    /// Graph node types we want to step to in the walk.
    #[clap(long, short = 'i', default_values = &[DEFAULT])]
    pub include_node_type: Vec<NodeTypeArg>,

    /// Graph edge types to exclude from walk. Can pass pre-configured sets
    /// via deep, shallow, hg, bonsai, etc as well as individual types.
    #[clap(long, short = 'X')]
    pub exclude_edge_type: Vec<EdgeTypeArg>,
    /// Graph edge types to include in the walk. Defaults to deep traversal.
    #[clap(long, short = 'I', default_values = &[DEEP])]
    pub include_edge_type: Vec<EdgeTypeArg>,

    /// Use this to continue walking even if walker found an error. Types of
    /// nodes to allow the walker to convert an ErrorKind::NotTraversable to
    /// a NodeData::ErrorAsData(NotTraversable)
    #[clap(long, short = 'e')]
    pub error_as_data_node_type: Vec<NodeTypeArg>,
    /// Types of edges to allow the walker to convert an ErrorKind::NotTraversable
    /// to a NodeData::ErrorAsData(NotTraversable). If empty then allow all
    /// edges for the nodes specified via error-as-data-node-type.
    #[clap(long, short = 'E')]
    pub error_as_data_edge_type: Vec<EdgeTypeArg>,
}

pub struct WalkerGraphParams {
    pub include_node_types: HashSet<NodeType>,
    pub include_edge_types: HashSet<EdgeType>,
    pub error_as_data_node_types: HashSet<NodeType>,
    pub error_as_data_edge_types: HashSet<EdgeType>,
}

impl WalkerGraphArgs {
    pub fn parse_args(&self) -> Result<WalkerGraphParams, Error> {
        let include_node_types =
            NodeTypeArg::filter(&self.include_node_type, &self.exclude_node_type);

        let include_edge_types =
            EdgeTypeArg::filter(&self.include_edge_type, &self.exclude_edge_type);

        let error_as_data_node_types = NodeTypeArg::parse_args(&self.error_as_data_node_type);
        let error_as_data_edge_types =
            EdgeTypeArg::filter(&self.error_as_data_edge_type, &self.exclude_edge_type);

        Ok(WalkerGraphParams {
            include_node_types,
            include_edge_types,
            error_as_data_node_types,
            error_as_data_edge_types,
        })
    }
}


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

use anyhow::Error;
use bookmarks::BookmarkKey;
use clap::Args;

use crate::detail::graph::EdgeType;
use crate::detail::graph::Node;
use crate::detail::parse_node::parse_node;
use crate::detail::walk::OutgoingEdge;

#[derive(Args, Debug)]
pub struct WalkRootArgs {
    /// Bookmark(s) to start traversal from.
    #[clap(long, short = 'b')]
    pub bookmark: Vec<BookmarkKey>,
    /// Root(s) to start traversal from in format <NodeType>:<node_key>, e.g.
    /// Bookmark:master or HgChangeset:7712b62acdc858689504945ac8965a303ded6626
    #[clap(long, short = 'r')]
    pub walk_root: Vec<String>,
}

impl WalkRootArgs {
    pub fn parse_args(&self) -> Result<Vec<OutgoingEdge>, Error> {
        let mut walk_roots: Vec<OutgoingEdge> = vec![];

        let mut bookmarks = self
            .bookmark
            .iter()
            .map(|b| OutgoingEdge::new(EdgeType::RootToBookmark, Node::Bookmark(b.clone())))
            .collect();
        walk_roots.append(&mut bookmarks);

        let roots: Result<Vec<_>, Error> =
            self.walk_root.iter().map(|root| parse_node(root)).collect();
        let mut roots = roots?
            .into_iter()
            .filter_map(|node| {
                node.get_type()
                    .root_edge_type()
                    .map(|et| OutgoingEdge::new(et, node))
            })
            .collect();
        walk_roots.append(&mut roots);

        Ok(walk_roots)
    }
}


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

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use clap::Parser;
use executor_lib::RepoShardedProcess;
use executor_lib::RepoShardedProcessExecutor;
use executor_lib::ShardedProcessExecutor;
use fbinit::FacebookInit;
use mononoke_app::args::MultiRepoArgs;
use mononoke_app::MononokeApp;
use once_cell::sync::OnceCell;
use sharding_ext::RepoShard;
use slog::info;
use slog::Logger;

use crate::args::SamplingArgs;
use crate::args::WalkerCommonArgs;
use crate::commands::JobParams;
use crate::commands::COMPRESSION_BENEFIT;
use crate::detail::graph::Node;
use crate::detail::sampling::WalkSampleMapping;
use crate::detail::sizing::compression_benefit;
use crate::detail::sizing::SizingCommand;
use crate::detail::sizing::SizingSample;
use crate::setup::setup_common;
use crate::WalkerArgs;

const SM_SERVICE_SCOPE: &str = "global";
const SM_CLEANUP_TIMEOUT_SECS: u64 = 120;

/// Estimate compression benefit.
#[derive(Parser)]
pub struct CommandArgs {
    /// Zstd compression level to use.
    #[clap(long, default_value = "3")]
    pub compression_level: i32,

    #[clap(flatten, next_help_heading = "SAMPLING OPTIONS")]
    pub sampling: SamplingArgs,

    #[clap(flatten)]
    pub common_args: WalkerCommonArgs,
}

/// Struct representing the Walker Sizing BP.
pub struct WalkerSizingProcess {
    app: MononokeApp,
    args: CommandArgs,
}

impl WalkerSizingProcess {
    fn new(app: MononokeApp, args: CommandArgs) -> Self {
        Self { app, args }
    }
}

#[async_trait]
impl RepoShardedProcess for WalkerSizingProcess {
    async fn setup(&self, repo: &RepoShard) -> anyhow::Result<Arc<dyn RepoShardedProcessExecutor>> {
        let repo_name = repo.repo_name.as_str();
        let logger = self.app.repo_logger(repo_name);
        info!(&logger, "Setting up walker sizing for repo {}", repo_name);
        let repos = MultiRepoArgs {
            repo_name: vec![repo_name.to_string()],
            repo_id: vec![],
        };
        let (job_params, command) = setup_sizing(&repos, &self.app, &self.args)
            .await
            .with_context(|| {
                format!(
                    "Failure in setting up walker sizing for repo {}",
                    &repo_name
                )
            })?;
        info!(
            &logger,
            "Completed walker sizing setup for repo {}", repo_name
        );
        Ok(Arc::new(WalkerSizingProcessExecutor::new(
            self.app.fb,
            logger,
            job_params,
            command,
            repo_name.to_string(),
        )))
    }
}

/// Struct representing the execution of Walker Sizing
/// BP over the context of a provided repo.
pub struct WalkerSizingProcessExecutor {
    fb: FacebookInit,
    logger: Logger,
    job_params: JobParams,
    command: SizingCommand,
    cancellation_requested: Arc<AtomicBool>,
    repo_name: String,
}

impl WalkerSizingProcessExecutor {
    fn new(
        fb: FacebookInit,
        logger: Logger,
        job_params: JobParams,
        command: SizingCommand,
        repo_name: String,
    ) -> Self {
        Self {
            cancellation_requested: Arc::new(AtomicBool::new(false)),
            fb,
            logger,
            job_params,
            command,
            repo_name,
        }
    }
}

#[async_trait]
impl RepoShardedProcessExecutor for WalkerSizingProcessExecutor {
    async fn execute(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Initiating walker sizing execution for repo {}", &self.repo_name,
        );
        compression_benefit(
            self.fb,
            self.job_params.clone(),
            self.command.clone(),
            Arc::clone(&self.cancellation_requested),
        )
        .await
        .with_context(|| {
            format!(
                "Error while executing walker sizing for repo {}",
                &self.repo_name
            )
        })
    }

    async fn stop(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Terminating walker sizing execution for repo {}", &self.repo_name,
        );
        self.cancellation_requested.store(true, Ordering::Relaxed);
        Ok(())
    }
}

async fn setup_sizing(
    repos: &MultiRepoArgs,
    app: &MononokeApp,
    args: &CommandArgs,
) -> Result<(JobParams, SizingCommand), Error> {
    let CommandArgs {
        compression_level,
        sampling,
        common_args,
    } = args;

    let sampler = Arc::new(WalkSampleMapping::<Node, SizingSample>::new());
    let job_params = setup_common(
        COMPRESSION_BENEFIT,
        app,
        repos,
        common_args,
        Some(sampler.clone()), // blobstore sampler
        None,                  // blobstore component sampler
    )
    .await?;

    let command = SizingCommand {
        compression_level: *compression_level,
        progress_options: common_args.progress.parse_args(),
        sampling_options: sampling.parse_args(100 /* default_sample_rate */)?,
        sampler,
    };
    Ok((job_params, command))
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<(), Error> {
    let walker_args = &app.args::<WalkerArgs>()?;
    match &walker_args.sharded_service_name {
        Some(service_name) => run_sharded(app, args, service_name.to_string()).await,
        None => run_unsharded(&walker_args.repos, app, args).await,
    }
}

/// The run variant for sharded execution of walker sizing.
pub async fn run_sharded(
    app: MononokeApp,
    args: CommandArgs,
    service_name: String,
) -> Result<(), Error> {
    let sizing_process = WalkerSizingProcess::new(app, args);
    let logger = sizing_process.app.logger().clone();
    // The service name needs to be 'static to satisfy SM contract
    static SM_SERVICE_NAME: OnceCell<String> = OnceCell::new();
    let mut executor = ShardedProcessExecutor::new(
        sizing_process.app.fb,
        sizing_process.app.runtime().clone(),
        &logger,
        SM_SERVICE_NAME.get_or_init(|| service_name),
        SM_SERVICE_SCOPE,
        SM_CLEANUP_TIMEOUT_SECS,
        Arc::new(sizing_process),
        true, // enable shard (repo) level healing
    )?;
    executor
        .block_and_execute(&logger, Arc::new(AtomicBool::new(false)))
        .await
}

pub async fn run_unsharded(
    repos: &MultiRepoArgs,
    app: MononokeApp,
    args: CommandArgs,
) -> Result<(), Error> {
    let (job_params, command) = setup_sizing(repos, &app, &args).await?;
    // When running in unsharded setting, walker sizing doesn't need to
    // be cancelled midway.
    compression_benefit(
        app.fb,
        job_params,
        command,
        Arc::new(AtomicBool::new(false)),
    )
    .await
}


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

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use clap::Parser;
use executor_lib::RepoShardedProcess;
use executor_lib::RepoShardedProcessExecutor;
use executor_lib::ShardedProcessExecutor;
use fbinit::FacebookInit;
use mononoke_app::args::MultiRepoArgs;
use mononoke_app::MononokeApp;
use once_cell::sync::OnceCell;
use sharding_ext::RepoShard;
use slog::info;
use slog::Logger;

use crate::args::SamplingArgs;
use crate::args::WalkerCommonArgs;
use crate::commands::JobParams;
use crate::commands::CORPUS;
use crate::detail::corpus::corpus;
use crate::detail::corpus::CorpusCommand;
use crate::detail::corpus::CorpusSample;
use crate::detail::corpus::CorpusSamplingHandler;
use crate::setup::setup_common;
use crate::WalkerArgs;

const SM_SERVICE_SCOPE: &str = "global";
const SM_CLEANUP_TIMEOUT_SECS: u64 = 120;

/// Dump a sampled corpus of blobstore data.
#[derive(Parser)]
pub struct CommandArgs {
    /// Where to write the output corpus. Default is to to a dry run with no output.
    #[clap(long)]
    pub output_dir: Option<String>,

    #[clap(flatten, next_help_heading = "SAMPLING OPTIONS")]
    pub sampling: SamplingArgs,

    #[clap(flatten)]
    pub common_args: WalkerCommonArgs,
}

/// Struct representing the Walker Corpus BP.
pub struct WalkerCorpusProcess {
    app: MononokeApp,
    args: CommandArgs,
}

impl WalkerCorpusProcess {
    fn new(app: MononokeApp, args: CommandArgs) -> Self {
        Self { app, args }
    }
}

#[async_trait]
impl RepoShardedProcess for WalkerCorpusProcess {
    async fn setup(&self, repo: &RepoShard) -> anyhow::Result<Arc<dyn RepoShardedProcessExecutor>> {
        let repo_name = repo.repo_name.as_str();
        let logger = self.app.repo_logger(repo_name);
        info!(&logger, "Setting up walker corpus for repo {}", repo_name);
        let repos = MultiRepoArgs {
            repo_name: vec![repo_name.to_string()],
            repo_id: vec![],
        };
        let (job_params, command) = setup_corpus(&repos, &self.app, &self.args)
            .await
            .with_context(|| {
                format!(
                    "Failure in setting up walker corpus for repo {}",
                    &repo_name
                )
            })?;
        info!(
            &logger,
            "Completed walker corpus setup for repo {}", repo_name
        );
        Ok(Arc::new(WalkerCorpusProcessExecutor::new(
            self.app.fb,
            logger,
            job_params,
            command,
            repo_name.to_string(),
        )))
    }
}

/// Struct representing the execution of Walker Corpus
/// BP over the context of a provided repo.
pub struct WalkerCorpusProcessExecutor {
    fb: FacebookInit,
    logger: Logger,
    job_params: JobParams,
    command: CorpusCommand,
    cancellation_requested: Arc<AtomicBool>,
    repo_name: String,
}

impl WalkerCorpusProcessExecutor {
    fn new(
        fb: FacebookInit,
        logger: Logger,
        job_params: JobParams,
        command: CorpusCommand,
        repo_name: String,
    ) -> Self {
        Self {
            cancellation_requested: Arc::new(AtomicBool::new(false)),
            fb,
            logger,
            job_params,
            command,
            repo_name,
        }
    }
}

#[async_trait]
impl RepoShardedProcessExecutor for WalkerCorpusProcessExecutor {
    async fn execute(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Initiating walker corpus execution for repo {}", &self.repo_name,
        );
        corpus(
            self.fb,
            self.job_params.clone(),
            self.command.clone(),
            Arc::clone(&self.cancellation_requested),
        )
        .await
        .with_context(|| {
            format!(
                "Error while executing walker corpus execution for repo {}",
                &self.repo_name
            )
        })
    }

    async fn stop(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Terminating walker corpus execution for repo {}", &self.repo_name,
        );
        self.cancellation_requested.store(true, Ordering::Relaxed);
        Ok(())
    }
}

async fn setup_corpus(
    repos: &MultiRepoArgs,
    app: &MononokeApp,
    args: &CommandArgs,
) -> Result<(JobParams, CorpusCommand), Error> {
    let CommandArgs {
        output_dir,
        sampling,
        common_args,
    } = args;

    let sampler = Arc::new(CorpusSamplingHandler::<CorpusSample>::new(
        output_dir.clone(),
    ));
    let job_params = setup_common(
        CORPUS,
        app,
        repos,
        common_args,
        Some(sampler.clone()), // blobstore sampler
        None,                  // blobstore component sampler
    )
    .await?;

    if let Some(output_dir) = &output_dir {
        if !std::path::Path::new(output_dir).exists() {
            std::fs::create_dir(output_dir).map_err(Error::from)?
        }
    }

    let command = CorpusCommand {
        output_dir: output_dir.clone(),
        progress_options: common_args.progress.parse_args(),
        sampling_options: sampling.parse_args(100 /* default_sample_rate */)?,
        sampling_path_regex: sampling.sample_path_regex.clone(),
        sampler,
    };

    Ok((job_params, command))
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<(), Error> {
    let walker_args = &app.args::<WalkerArgs>()?;
    match &walker_args.sharded_service_name {
        Some(service_name) => run_sharded(app, args, service_name.to_string()).await,
        None => run_unsharded(&walker_args.repos, app, args).await,
    }
}

/// The run variant for sharded execution of walker corpus.
pub async fn run_sharded(
    app: MononokeApp,
    args: CommandArgs,
    service_name: String,
) -> Result<(), Error> {
    let corpus_process = WalkerCorpusProcess::new(app, args);
    let logger = corpus_process.app.logger().clone();
    // The service name needs to be 'static to satisfy SM contract
    static SM_SERVICE_NAME: OnceCell<String> = OnceCell::new();
    let mut executor = ShardedProcessExecutor::new(
        corpus_process.app.fb,
        corpus_process.app.runtime().clone(),
        &logger,
        SM_SERVICE_NAME.get_or_init(|| service_name),
        SM_SERVICE_SCOPE,
        SM_CLEANUP_TIMEOUT_SECS,
        Arc::new(corpus_process),
        true, // enable shard (repo) level healing
    )?;
    executor
        .block_and_execute(&logger, Arc::new(AtomicBool::new(false)))
        .await
}

pub async fn run_unsharded(
    repos: &MultiRepoArgs,
    app: MononokeApp,
    args: CommandArgs,
) -> Result<(), Error> {
    let (job_params, command) = setup_corpus(repos, &app, &args).await?;
    // When running in unsharded setting, walker corpus doesn't need to
    // be cancelled midway.
    corpus(
        app.fb,
        job_params,
        command,
        Arc::new(AtomicBool::new(false)),
    )
    .await
}


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

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use clap::Parser;
use executor_lib::RepoShardedProcess;
use executor_lib::RepoShardedProcessExecutor;
use executor_lib::ShardedProcessExecutor;
use fbinit::FacebookInit;
use mononoke_app::args::MultiRepoArgs;
use mononoke_app::MononokeApp;
use once_cell::sync::OnceCell;
use sharding_ext::RepoShard;
use slog::info;
use slog::Logger;

use crate::args::OutputFormat;
use crate::args::SamplingArgs;
use crate::args::ScrubOutputNodeArgs;
use crate::args::ScrubPackLogArgs;
use crate::args::WalkerCommonArgs;
use crate::commands::JobParams;
use crate::commands::SCRUB;
use crate::detail::graph::Node;
use crate::detail::sampling::WalkSampleMapping;
use crate::detail::scrub::scrub_objects;
use crate::detail::scrub::ScrubCommand;
use crate::detail::scrub::ScrubSample;
use crate::setup::setup_common;
use crate::WalkerArgs;

const SM_SERVICE_SCOPE: &str = "global";
const SM_CLEANUP_TIMEOUT_SECS: u64 = 120;

/// Checks the data is present by reading it and counting it.
/// Combine with --enable-scrub-blobstore to check across the multiplex.
#[derive(Parser)]
pub struct CommandArgs {
    /// Set the output format
    #[clap(long, short = 'F', default_value = "PrettyDebug")]
    pub output_format: OutputFormat,

    #[clap(flatten)]
    pub output_nodes: ScrubOutputNodeArgs,

    #[clap(flatten)]
    pub pack_log_info: ScrubPackLogArgs,

    #[clap(flatten, next_help_heading = "SAMPLING OPTIONS")]
    pub sampling: SamplingArgs,

    #[clap(flatten)]
    pub common_args: WalkerCommonArgs,
}

/// Struct representing the Walker Scrub BP.
pub struct WalkerScrubProcess {
    app: MononokeApp,
    args: CommandArgs,
}

impl WalkerScrubProcess {
    fn new(app: MononokeApp, args: CommandArgs) -> Self {
        Self { app, args }
    }
}

#[async_trait]
impl RepoShardedProcess for WalkerScrubProcess {
    async fn setup(&self, repo: &RepoShard) -> anyhow::Result<Arc<dyn RepoShardedProcessExecutor>> {
        let repo_name = repo.repo_name.as_str();
        let logger = self.app.repo_logger(repo_name);
        info!(&logger, "Setting up walker scrub for repo {}", repo_name);
        let repos = MultiRepoArgs {
            repo_name: vec![repo_name.to_string()],
            repo_id: vec![],
        };
        let (job_params, command) = setup_scrub(&repos, &self.app, &self.args)
            .await
            .with_context(|| {
                format!("Failure in setting up walker scrub for repo {}", &repo_name)
            })?;
        info!(
            &logger,
            "Completed walker scrub setup for repo {}", repo_name
        );
        Ok(Arc::new(WalkerScrubProcessExecutor::new(
            self.app.fb,
            logger,
            job_params,
            command,
            repo_name.to_string(),
        )))
    }
}

/// Struct representing the execution of the Walker Scrub
/// BP over the context of a provided repo.
pub struct WalkerScrubProcessExecutor {
    fb: FacebookInit,
    logger: Logger,
    job_params: JobParams,
    command: ScrubCommand,
    cancellation_requested: Arc<AtomicBool>,
    repo_name: String,
}

impl WalkerScrubProcessExecutor {
    fn new(
        fb: FacebookInit,
        logger: Logger,
        job_params: JobParams,
        command: ScrubCommand,
        repo_name: String,
    ) -> Self {
        Self {
            cancellation_requested: Arc::new(AtomicBool::new(false)),
            fb,
            logger,
            job_params,
            command,
            repo_name,
        }
    }
}

#[async_trait]
impl RepoShardedProcessExecutor for WalkerScrubProcessExecutor {
    async fn execute(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Initiating walker scrub execution for repo {}", &self.repo_name,
        );
        scrub_objects(
            self.fb,
            self.job_params.clone(),
            self.command.clone(),
            Arc::clone(&self.cancellation_requested),
        )
        .await
        .with_context(|| {
            format!(
                "Error while executing walker scrub execution for repo {}",
                &self.repo_name
            )
        })
    }

    async fn stop(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Terminating walker scrub execution for repo {}", &self.repo_name,
        );
        self.cancellation_requested.store(true, Ordering::Relaxed);
        Ok(())
    }
}

async fn setup_scrub(
    repos: &MultiRepoArgs,
    app: &MononokeApp,
    args: &CommandArgs,
) -> Result<(JobParams, ScrubCommand), Error> {
    let component_sampler = Arc::new(WalkSampleMapping::<Node, ScrubSample>::new());
    let job_params = setup_common(
        SCRUB,
        app,
        repos,
        &args.common_args,
        None,
        Some(component_sampler.clone()),
    )
    .await?;

    let CommandArgs {
        output_format,
        output_nodes,
        pack_log_info,
        sampling,
        common_args,
    } = args;
    let command = ScrubCommand {
        limit_data_fetch: common_args.limit_data_fetch,
        output_format: output_format.clone(),
        output_node_types: output_nodes.parse_args(),
        progress_options: common_args.progress.parse_args(),
        sampling_options: sampling.parse_args(1)?,
        pack_info_log_options: pack_log_info.parse_args(app.fb)?,
        sampler: component_sampler,
    };
    Ok((job_params, command))
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<(), Error> {
    let walker_args = &app.args::<WalkerArgs>()?;
    match &walker_args.sharded_service_name {
        Some(service_name) => run_sharded(app, args, service_name.to_string()).await,
        None => run_unsharded(&walker_args.repos, app, args).await,
    }
}

/// The run variant for sharded execution of walker scrub.
pub async fn run_sharded(
    app: MononokeApp,
    args: CommandArgs,
    service_name: String,
) -> Result<(), Error> {
    let scrub_process = WalkerScrubProcess::new(app, args);
    let logger = scrub_process.app.logger().clone();
    // The service name needs to be 'static to satisfy SM contract
    static SM_SERVICE_NAME: OnceCell<String> = OnceCell::new();
    let mut executor = ShardedProcessExecutor::new(
        scrub_process.app.fb,
        scrub_process.app.runtime().clone(),
        &logger,
        SM_SERVICE_NAME.get_or_init(|| service_name),
        SM_SERVICE_SCOPE,
        SM_CLEANUP_TIMEOUT_SECS,
        Arc::new(scrub_process),
        true, // enable shard (repo) level healing
    )?;
    executor
        .block_and_execute(&logger, Arc::new(AtomicBool::new(false)))
        .await
}

pub async fn run_unsharded(
    repos: &MultiRepoArgs,
    app: MononokeApp,
    args: CommandArgs,
) -> Result<(), Error> {
    let (job_params, command) = setup_scrub(repos, &app, &args).await?;
    // When running in unsharded setting, walker scrub doesn't have a need to
    // be cancelled midway.
    scrub_objects(
        app.fb,
        job_params,
        command,
        Arc::new(AtomicBool::new(false)),
    )
    .await
}


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

use std::sync::atomic::AtomicBool;
use std::sync::atomic::Ordering;
use std::sync::Arc;

use anyhow::Context;
use anyhow::Error;
use async_trait::async_trait;
use clap::Parser;
use executor_lib::RepoShardedProcess;
use executor_lib::RepoShardedProcessExecutor;
use executor_lib::ShardedProcessExecutor;
use fbinit::FacebookInit;
use mononoke_app::args::MultiRepoArgs;
use mononoke_app::MononokeApp;
use once_cell::sync::OnceCell;
use sharding_ext::RepoShard;
use slog::info;
use slog::Logger;

use crate::args::ValidateCheckTypeArgs;
use crate::args::WalkerCommonArgs;
use crate::commands::JobParams;
use crate::commands::VALIDATE;
use crate::detail::validate::validate;
use crate::detail::validate::ValidateCommand;
use crate::setup::setup_common;
use crate::WalkerArgs;

const SM_SERVICE_SCOPE: &str = "global";
const SM_CLEANUP_TIMEOUT_SECS: u64 = 120;

/// Walk the graph and perform checks on it.
#[derive(Parser)]
pub struct CommandArgs {
    #[clap(flatten)]
    pub check_types: ValidateCheckTypeArgs,

    #[clap(flatten)]
    pub common_args: WalkerCommonArgs,
}

/// Struct representing the Walker Validate BP.
pub struct WalkerValidateProcess {
    app: MononokeApp,
    args: CommandArgs,
}

impl WalkerValidateProcess {
    fn new(app: MononokeApp, args: CommandArgs) -> Self {
        Self { app, args }
    }
}

#[async_trait]
impl RepoShardedProcess for WalkerValidateProcess {
    async fn setup(&self, repo: &RepoShard) -> anyhow::Result<Arc<dyn RepoShardedProcessExecutor>> {
        let repo_name = repo.repo_name.as_str();
        let logger = self.app.repo_logger(repo_name);
        info!(&logger, "Setting up walker validate for repo {}", repo_name);
        let repos = MultiRepoArgs {
            repo_name: vec![repo_name.to_string()],
            repo_id: vec![],
        };
        let (job_params, command) = setup_validate(&repos, &self.app, &self.args)
            .await
            .with_context(|| {
                format!(
                    "Failure in setting up walker validate for repo {}",
                    &repo_name
                )
            })?;
        info!(
            &logger,
            "Completed walker validate setup for repo {}", repo_name
        );
        Ok(Arc::new(WalkerValidateProcessExecutor::new(
            self.app.fb,
            logger,
            job_params,
            command,
            repo_name.to_string(),
        )))
    }
}

/// Struct representing the execution of Walker Validate
/// BP over the context of a provided repo.
pub struct WalkerValidateProcessExecutor {
    fb: FacebookInit,
    logger: Logger,
    job_params: JobParams,
    command: ValidateCommand,
    cancellation_requested: Arc<AtomicBool>,
    repo_name: String,
}

impl WalkerValidateProcessExecutor {
    fn new(
        fb: FacebookInit,
        logger: Logger,
        job_params: JobParams,
        command: ValidateCommand,
        repo_name: String,
    ) -> Self {
        Self {
            cancellation_requested: Arc::new(AtomicBool::new(false)),
            fb,
            logger,
            job_params,
            command,
            repo_name,
        }
    }
}

#[async_trait]
impl RepoShardedProcessExecutor for WalkerValidateProcessExecutor {
    async fn execute(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Initiating walker validate execution for repo {}", &self.repo_name,
        );
        validate(
            self.fb,
            self.job_params.clone(),
            self.command.clone(),
            Arc::clone(&self.cancellation_requested),
        )
        .await
        .with_context(|| {
            format!(
                "Error while executing walker validate execution for repo {}",
                &self.repo_name
            )
        })
    }

    async fn stop(&self) -> anyhow::Result<()> {
        info!(
            self.logger,
            "Terminating walker validate execution for repo {}", &self.repo_name,
        );
        self.cancellation_requested.store(true, Ordering::Relaxed);
        Ok(())
    }
}

async fn setup_validate(
    repos: &MultiRepoArgs,
    app: &MononokeApp,
    args: &CommandArgs,
) -> Result<(JobParams, ValidateCommand), Error> {
    let CommandArgs {
        check_types,
        common_args,
    } = args;
    let job_params = setup_common(
        VALIDATE,
        app,
        repos,
        common_args,
        None, // blobstore sampler
        None, // blobstore component sampler
    )
    .await?;

    let command = ValidateCommand {
        include_check_types: check_types.parse_args(),
        progress_options: common_args.progress.parse_args(),
    };
    Ok((job_params, command))
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<(), Error> {
    let walker_args = &app.args::<WalkerArgs>()?;
    match &walker_args.sharded_service_name {
        Some(service_name) => run_sharded(app, args, service_name.to_string()).await,
        None => run_unsharded(&walker_args.repos, app, args).await,
    }
}

/// The run variant for sharded execution of walker validatet.
pub async fn run_sharded(
    app: MononokeApp,
    args: CommandArgs,
    service_name: String,
) -> Result<(), Error> {
    let validate_process = WalkerValidateProcess::new(app, args);
    let logger = validate_process.app.logger().clone();
    // The service name needs to be 'static to satisfy SM contract
    static SM_SERVICE_NAME: OnceCell<String> = OnceCell::new();
    let mut executor = ShardedProcessExecutor::new(
        validate_process.app.fb,
        validate_process.app.runtime().clone(),
        &logger,
        SM_SERVICE_NAME.get_or_init(|| service_name),
        SM_SERVICE_SCOPE,
        SM_CLEANUP_TIMEOUT_SECS,
        Arc::new(validate_process),
        true, // enable shard (repo) level healing
    )?;
    executor
        .block_and_execute(&logger, Arc::new(AtomicBool::new(false)))
        .await
}

pub async fn run_unsharded(
    repos: &MultiRepoArgs,
    app: MononokeApp,
    args: CommandArgs,
) -> Result<(), Error> {
    let (job_params, command) = setup_validate(repos, &app, &args).await?;
    // When running in unsharded setting, walker validate doesn't need to
    // be cancelled midway.
    validate(
        app.fb,
        job_params,
        command,
        Arc::new(AtomicBool::new(false)),
    )
    .await
}


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

use std::collections::HashMap;
use std::fmt;
use std::str::FromStr;

use anyhow::anyhow;
use anyhow::Error;
use blobstore::BlobstoreMetadata;
use context::CoreContext;
use metaconfig_types::BlobConfig;
use metaconfig_types::BlobstoreId;
use mononoke_types::repo::REPO_PREFIX_REGEX;
use mononoke_types::RepositoryId;
use multiplexedblob::LoggingScrubHandler;
use multiplexedblob::ScrubHandler;
use scuba::value::NullScubaValue;
use scuba::value::ScubaValue;
use scuba_ext::MononokeScubaSampleBuilder;
use stats::prelude::*;

use crate::detail::pack::CTIME;
use crate::detail::validate::CHECK_FAIL;
use crate::detail::validate::CHECK_TYPE;
use crate::detail::validate::ERROR_MSG;
use crate::detail::validate::NODE_KEY;
use crate::detail::validate::REPO;

define_stats! {
    prefix = "mononoke.walker";
    scrub_repaired: dynamic_timeseries("{}.blobstore.{}.{}.repaired", (subcommand: &'static str, blobstore_id: String, repo: String); Rate, Sum),
    scrub_repair_required: dynamic_timeseries("{}.blobstore.{}.{}.repair_required", (subcommand: &'static str, blobstore_id: String, repo: String); Rate, Sum),
}

pub const BLOBSTORE_ID: &str = "blobstore_id";

pub struct StatsScrubHandler {
    scuba: MononokeScubaSampleBuilder,
    subcommand_stats_key: &'static str,
    repo_id_to_name: HashMap<RepositoryId, String>,
    inner: LoggingScrubHandler,
}

impl StatsScrubHandler {
    pub fn new(
        quiet: bool,
        scuba: MononokeScubaSampleBuilder,
        subcommand_stats_key: &'static str,
        repo_id_to_name: HashMap<RepositoryId, String>,
    ) -> Self {
        Self {
            scuba,
            subcommand_stats_key,
            inner: LoggingScrubHandler::new(quiet),
            repo_id_to_name,
        }
    }
}

impl fmt::Debug for StatsScrubHandler {
    fn fmt(&self, f: &mut fmt::Formatter<'_>) -> fmt::Result {
        f.debug_struct("StatsScrubHandler")
            .field("subcommand_stats_key", &self.subcommand_stats_key)
            .field("inner", &self.inner)
            .finish()
    }
}

pub fn get_repo_id_from_key(key: &str) -> Result<Option<RepositoryId>, Error> {
    REPO_PREFIX_REGEX
        .captures(key)
        .and_then(|m| m.get(1).map(|m| RepositoryId::from_str(m.as_str())))
        .transpose()
}

impl ScrubHandler for StatsScrubHandler {
    fn on_repair(
        &self,
        ctx: &CoreContext,
        blobstore_id: BlobstoreId,
        key: &str,
        is_repaired: bool,
        meta: &BlobstoreMetadata,
    ) {
        self.inner
            .on_repair(ctx, blobstore_id, key, is_repaired, meta);

        let ctime = match meta.ctime() {
            Some(ctime) => ScubaValue::from(ctime),
            None => ScubaValue::Null(NullScubaValue::Int),
        };

        let mut scuba = self.scuba.clone();

        let repo_id = get_repo_id_from_key(key);
        match repo_id {
            Ok(Some(repo_id)) => {
                if let Some(repo_name) = self.repo_id_to_name.get(&repo_id) {
                    scuba.add(REPO, repo_name.clone());
                    if is_repaired {
                        STATS::scrub_repaired.add_value(
                            1,
                            (
                                self.subcommand_stats_key,
                                blobstore_id.to_string(),
                                repo_name.clone(),
                            ),
                        );
                    } else {
                        STATS::scrub_repair_required.add_value(
                            1,
                            (
                                self.subcommand_stats_key,
                                blobstore_id.to_string(),
                                repo_name.clone(),
                            ),
                        );
                    }
                }
            }
            Ok(_) => {}
            Err(e) => {
                scuba.add(ERROR_MSG, format!("{:?}", e));
            }
        }

        scuba
            .add(BLOBSTORE_ID, blobstore_id)
            // TODO parse out NodeType from string key prefix if we can. Or better, make blobstore keys typed?
            .add(NODE_KEY, key)
            .add(CHECK_TYPE, "scrub_repair")
            .add(CHECK_FAIL, if is_repaired { 0 } else { 1 })
            .add("session", ctx.session().metadata().session_id().to_string())
            .add(CTIME, ctime)
            .log();
    }
}

pub fn replace_blobconfig(
    blob_config: &mut BlobConfig,
    inner_blobstore_id: Option<u64>,
    repo_name: &str,
    walk_stats_key: &'static str,
    is_scrubbing: bool,
) -> Result<(), Error> {
    match blob_config {
        BlobConfig::MultiplexedWal { ref blobstores, .. } => {
            if is_scrubbing {
                // Make sure the repair stats are set to zero for each store.
                // Without this the new stats only show up when a repair is
                // needed (i.e. as they get incremented), which makes them
                // harder to monitor on (no datapoints rather than a zero
                // datapoint at start).
                for s in &[STATS::scrub_repaired, STATS::scrub_repair_required] {
                    for (id, _ty, _config) in blobstores {
                        s.add_value(0, (walk_stats_key, id.to_string(), repo_name.to_string()));
                    }
                }
            }
            if let Some(inner_blobstore_id) = inner_blobstore_id {
                let sought_id = BlobstoreId::new(inner_blobstore_id);
                let inner_blob_config = blobstores
                    .iter()
                    .find_map(|(blobstore_id, _, blobstore)| {
                        if blobstore_id == &sought_id {
                            Some(blobstore.clone())
                        } else {
                            None
                        }
                    })
                    .ok_or_else(|| {
                        anyhow!("could not find a blobstore with id {}", inner_blobstore_id)
                    })?;
                *blob_config = inner_blob_config;
            }
        }
        _ => {
            if inner_blobstore_id.is_some() {
                return Err(anyhow!(
                    "inner-blobstore-id supplied but blobstore is not multiplexed"
                ));
            }
        }
    }
    Ok(())
}