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

use blobstore_factory::BlobstoreOptions;
use blobstore_factory::ReadOnlyStorage;
use cached_config::ConfigStore;
use clap::ValueEnum;
use derived_data_remote::RemoteDerivationOptions;
use fbinit::FacebookInit;
use megarepo_config::MononokeMegarepoConfigsOptions;
use observability::ObservabilityContext;
use permission_checker::AclProvider;
use rendezvous::RendezVousOptions;
use scuba_ext::MononokeScubaSampleBuilder;
use slog::Logger;
use sql_ext::facebook::MysqlOptions;
use strum::EnumString;
use tokio::runtime::Handle;

#[derive(Copy, Clone, PartialEq, Eq)]
pub struct LocalCacheConfig {
    /// Number of shards in the local blobstore cache
    pub blobstore_cache_shards: usize,
}

#[derive(Copy, Clone, PartialEq, Eq)]
pub enum Caching {
    /// Caching is fully enabled, with both local and shared caches.
    Enabled(LocalCacheConfig),

    /// Caching is only enabled locally - the shared cache is disabled.
    LocalOnly(LocalCacheConfig),

    /// Caching is not enabled.
    Disabled,
}

#[derive(Copy, Clone, Debug, ValueEnum, EnumString, strum::Display)]
pub enum WarmBookmarksCacheDerivedData {
    HgOnly,
    AllKinds,
    NoDerivation,
}

/// Struct representing the configuration associated with a MononokeApp instance which
/// is immutable post the point of app construction.
pub struct MononokeEnvironment {
    pub fb: FacebookInit,
    pub logger: Logger,
    pub scuba_sample_builder: MononokeScubaSampleBuilder,
    pub warm_bookmarks_cache_scuba_sample_builder: MononokeScubaSampleBuilder,
    pub config_store: ConfigStore,
    pub caching: Caching,
    pub observability_context: ObservabilityContext,
    pub runtime: Handle,
    pub mysql_options: MysqlOptions,
    pub blobstore_options: BlobstoreOptions,
    pub readonly_storage: ReadOnlyStorage,
    pub rendezvous_options: RendezVousOptions,
    pub megarepo_configs_options: MononokeMegarepoConfigsOptions,
    pub remote_derivation_options: RemoteDerivationOptions,
    pub disabled_hooks: HashMap<String, HashSet<String>>,
    pub acl_provider: Arc<dyn AclProvider>,
    pub warm_bookmarks_cache_derived_data: Option<WarmBookmarksCacheDerivedData>,
    /// Function determining whether given repo (identified by name) should be loaded
    pub filter_repos: Option<Arc<dyn Fn(&str) -> bool + Send + Sync + 'static>>,
}


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

/// Trait allows to configure defaults for the clap::Args, which can be
/// different for different Mononoke binaries.
pub trait ArgDefaults {
    // TODO: We may want to add the Args type defaults are set for and
    // extend the trait with more methods.
    // type Args: clap::Args;

    /// Obtain default values for these arguments.
    fn arg_defaults(&self) -> Vec<(&'static str, String)>;
}


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

use anyhow::Result;
use fbinit::FacebookInit;
use slog::Level;

/// Sets the log level used by glog library to match the level use we use in slog.
/// It can be overridden by setting the GLOG_minloglevel env variable.
pub fn set_glog_log_level(fb: FacebookInit, level: Level) -> Result<()> {
    if std::env::var("GLOG_minloglevel").is_err() {
        let level_glog_num = match level {
            Level::Critical => 4,
            Level::Error => 3,
            Level::Warning => 2,
            // Normally glog 1 corresponds to our INFO but since we normally
            // care very little about things logged by our deps it might be
            // easier to delegate those to DEBUG output.
            Level::Info => 2,
            _ => 0,
        };
        gflags::set_gflag_value(fb, "minloglevel", gflags::GflagValue::U32(level_glog_num))?;
    }
    Ok(())
}


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

#[cfg(fbcode_build)]
pub mod glog;
pub mod log;
mod logging_args;
mod scribe;
mod scuba;

pub use logging_args::LoggingArgs;
pub use scribe::ScribeLoggingArgs;
pub use scuba::ScubaLoggingArgs;


/*
 * 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 env_logger::filter::Builder;
use env_logger::filter::Filter;
use slog::BorrowedKV;
use slog::Level;
use slog::Logger;
use slog::SingleKV;

// NOTE: The following 2 methods and parts of the implementation of LinkedLogger::log are borrowed
// from the slog-stdlog crate (MIT-licensed).

fn log_to_slog_level(level: log::Level) -> Level {
    match level {
        log::Level::Trace => Level::Trace,
        log::Level::Debug => Level::Debug,
        log::Level::Info => Level::Info,
        log::Level::Warn => Level::Warning,
        log::Level::Error => Level::Error,
    }
}

fn record_as_location(r: &log::Record) -> slog::RecordLocation {
    let module = r.module_path_static().unwrap_or("<unknown>");
    let file = r.file_static().unwrap_or("<unknown>");
    let line = r.line().unwrap_or_default();

    slog::RecordLocation {
        file,
        line,
        column: 0,
        function: "",
        module,
    }
}

struct LinkedLogger {
    logger: Logger,
    filter: Filter,
}

impl log::Log for LinkedLogger {
    fn enabled(&self, m: &log::Metadata) -> bool {
        self.filter.enabled(m)
    }

    fn log(&self, r: &log::Record) {
        if !self.filter.matches(r) {
            return;
        }

        let level = log_to_slog_level(r.metadata().level());

        let args = r.args();
        let target = r.target();
        let location = &record_as_location(r);
        let s = slog::RecordStatic {
            location,
            level,
            tag: target,
        };

        // NOTE: Normally, we'd want to use the recommended b! macro from Slog here... but it turns
        // out that expand this macro gives us something that looks like ::slog::BorrowedKV(&FOO);,
        // which means FOO gets dropped immediately and we get an error telling us to assign it to
        // a variable (FOO is the SingleKV here). This does that.
        let k = SingleKV::from(("target", target));
        let x = BorrowedKV(&k);

        let record = slog::Record::new(&s, args, x);

        self.logger.log(&record);
    }

    fn flush(&self) {}
}

/// Wire up a slog Logger as the destination for std logs as per an env_logger filter spec. This
/// sets the global logger, so it'll error if called more than once.
pub fn init_stdlog_once(logger: Logger, var: &str) -> Result<log::LevelFilter, Error> {
    // NOTE: The default level is ERROR, which should be fairly reasonable.
    let filter = Builder::from_env(var).build();
    let level = filter.filter();

    log::set_boxed_logger(Box::new(LinkedLogger { logger, filter }))?;

    // set_max_level ensures we don't produce logs that won't pass any filter at all.
    log::set_max_level(level);

    Ok(level)
}


/*
 * 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::str::FromStr;
use std::sync::Arc;

#[cfg(fbcode_build)]
use anyhow::format_err;
use anyhow::Context;
use anyhow::Result;
use cached_config::ConfigStore;
use clap::builder::PossibleValuesParser;
use clap::ArgAction;
use clap::Args;
use clap::ValueEnum;
use fbinit::FacebookInit;
use observability::ObservabilityContext;
use panichandler::Fate;
use slog::debug;
use slog::o;
use slog::Drain;
use slog::Level;
use slog::Logger;
use slog::Never;
use slog::SendSyncRefUnwindSafeDrain;
use slog_ext::make_tag_filter_drain;
use slog_glog_fmt::kv_categorizer::FacebookCategorizer;
use slog_glog_fmt::kv_defaults::FacebookKV;
use slog_glog_fmt::GlogFormat;
use slog_term::TermDecorator;

/// Command line arguments for spawning slog Logger
#[derive(Args, Debug)]
pub struct LoggingArgs {
    /// Print debug output
    ///
    /// Equivalent to --log-level=DEBUG.
    #[clap(long)]
    pub debug: bool,

    /// Log level to use
    #[clap(long, conflicts_with = "debug", value_parser = PossibleValuesParser::new(&slog::LOG_LEVEL_NAMES))]
    pub log_level: Option<String>,

    /// Include only log messages with these slog::Record::tags() or
    /// log::Record::targets
    #[clap(long, short = 'l')]
    pub log_include_tag: Vec<String>,

    /// Exclude log messages with these slog::Record::tags() or
    /// log::Record::targets
    #[clap(long, short = 'L')]
    pub log_exclude_tag: Vec<String>,

    /// Logview category to log to. Logview is not used if not set
    #[clap(long)]
    pub logview_category: Option<String>,

    /// Logview level to filter
    ///
    /// If logview-category is not set then this is ignored.
    ///
    /// Note that this level is applied AFTER --log-level/--debug was applied,
    /// so it doesn't make sense to set this parameter to a lower level than
    /// --log-level.
    #[clap(long, requires = "logview_category", value_parser = PossibleValuesParser::new(&slog::LOG_LEVEL_NAMES))]
    pub logview_additional_level_filter: Option<String>,

    /// Fate of the process when a panic happens
    #[clap(long, value_enum, default_value_t=PanicFate::Abort)]
    pub panic_fate: PanicFate,

    /// Whether to instantiate ObservabilityContext::Dynamic, which reads
    /// logging levels from configerator. Overwrites --log-level or --debug
    // For compatibility with existing usage, this arg takes value,
    // for example `--with-dynamic-observability=true`.
    #[clap(long, default_value_t = false, value_name = "BOOL", action = ArgAction::Set)]
    pub with_dynamic_observability: bool,
}

#[derive(ValueEnum, Clone, Copy, Debug)]
#[clap(rename_all = "lower")]
pub enum PanicFate {
    None,
    Continue,
    Exit,
    Abort,
}

impl LoggingArgs {
    pub fn create_log_level(&self) -> Level {
        if self.debug {
            Level::Debug
        } else {
            match &self.log_level {
                Some(log_level_str) => Level::from_str(log_level_str)
                    .unwrap_or_else(|_| panic!("Unknown log level: {}", log_level_str)),
                None => Level::Info,
            }
        }
    }

    // Logic copied from: https://fburl.com/code/ygj4muxz
    pub fn create_root_log_drain(
        &self,
        fb: FacebookInit,
        log_level: Level,
    ) -> Result<impl Drain<Ok = (), Err = Never> + Clone> {
        // Set the panic handler up here. Not really relevent to logger other than it emits output
        // when things go wrong. This writes directly to stderr as coredumper expects.
        // TODO: separate the panic handler out from logging
        let fate = match self.panic_fate {
            PanicFate::None => None,
            PanicFate::Continue => Some(Fate::Continue),
            PanicFate::Exit => Some(Fate::Exit(101)),
            PanicFate::Abort => Some(Fate::Abort),
        };
        if let Some(fate) = fate {
            panichandler::set_panichandler(fate);
        }

        let stdlog_env = "RUST_LOG";

        let glog_drain = make_tag_filter_drain(
            glog_drain(),
            self.log_include_tag.iter().cloned().collect(),
            self.log_exclude_tag.iter().cloned().collect(),
            true, // Log messages which have no tags
        )?;

        let root_log_drain: Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>> = match &self
            .logview_category
        {
            Some(category) => {
                #[cfg(fbcode_build)]
                {
                    // Sometimes scribe writes can fail due to backpressure - it's OK to drop these
                    // since logview is sampled anyway.
                    let logview_drain =
                        ::slog_logview::LogViewDrain::new(fb, category).ignore_res();
                    match &self.logview_additional_level_filter {
                        Some(log_level_str) => {
                            let logview_level = Level::from_str(log_level_str)
                                .map_err(|_| format_err!("Unknown log level: {}", log_level_str))?;

                            let drain = slog::Duplicate::new(
                                glog_drain,
                                logview_drain.filter_level(logview_level).ignore_res(),
                            );
                            Arc::new(drain.ignore_res())
                        }
                        None => {
                            let drain = slog::Duplicate::new(glog_drain, logview_drain);
                            Arc::new(drain.ignore_res())
                        }
                    }
                }
                #[cfg(not(fbcode_build))]
                {
                    let _ = (fb, category);
                    unimplemented!(
                        "Passed --logview-category, but it is supported only for fbcode builds",
                    )
                }
            }
            None => Arc::new(glog_drain),
        };

        // NOTE: We pass an unfiltered Logger to init_stdlog_once. That's because we do the filtering
        // at the stdlog level there.
        let stdlog_logger = Logger::root(root_log_drain.clone(), o![]);
        let stdlog_level = crate::log::init_stdlog_once(stdlog_logger, stdlog_env)?;

        // Note what level we enabled stdlog at, so that if someone is trying to debug they get
        // informed of potentially needing to set RUST_LOG.
        debug!(
            Logger::root(
                root_log_drain.clone().filter_level(log_level).ignore_res(),
                o![]
            ),
            "enabled stdlog with level: {:?} (set {} to configure)", stdlog_level, stdlog_env
        );

        Ok(root_log_drain)
    }

    pub fn create_logger(
        &self,
        root_log_drain: Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>,
    ) -> Result<Logger> {
        let kv = FacebookKV::new().context("Failed to initialize FacebookKV")?;
        Ok(Logger::root(root_log_drain, o![kv]))
    }

    pub fn create_observability_context(
        &self,
        config_store: &ConfigStore,
        log_level: slog::Level,
    ) -> Result<ObservabilityContext> {
        if self.with_dynamic_observability {
            Ok(ObservabilityContext::new(config_store)?)
        } else {
            Ok(ObservabilityContext::new_static(log_level))
        }
    }
}

/// Create a default root logger for Facebook services
fn glog_drain() -> impl Drain<Ok = (), Err = Never> {
    let decorator = TermDecorator::new().build();
    // FacebookCategorizer is used for slog KV arguments.
    // At the time of writing this code FacebookCategorizer and FacebookKV
    // that was added below was mainly useful for logview logging and had no effect on GlogFormat
    let drain = GlogFormat::new(decorator, FacebookCategorizer).ignore_res();
    ::std::sync::Mutex::new(drain).ignore_res()
}


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

use std::path::PathBuf;

use anyhow::Result;
use clap::Args;
use fbinit::FacebookInit;
use scribe_ext::Scribe;

/// Command line argument that affect scribe logging
#[derive(Args, Debug)]
pub struct ScribeLoggingArgs {
    /// Filesystem directory where to log all scribe writes
    #[clap(long)]
    pub scribe_logging_directory: Option<String>,
}

impl ScribeLoggingArgs {
    pub fn get_scribe(&self, fb: FacebookInit) -> Result<Scribe> {
        match &self.scribe_logging_directory {
            Some(dir) => Ok(Scribe::new_to_file(PathBuf::from(dir))),
            None => Ok(Scribe::new(fb)),
        }
    }
}


/*
 * 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::hash_map::DefaultHasher;
use std::hash::Hash;
use std::hash::Hasher;

use anyhow::Error;
use anyhow::Result;
use clap::Args;
use fbinit::FacebookInit;
use observability::ObservabilityContext;
use scuba_ext::MononokeScubaSampleBuilder;
use tunables::tunables;

/// Command line arguments that control scuba logging
#[derive(Args, Debug)]
pub struct ScubaLoggingArgs {
    /// The name of the scuba dataset to log to
    /// Prefix `file://` will direct logging to local file
    #[clap(long)]
    pub scuba_dataset: Option<String>,
    /// Do not use the default scuba dataset for this app
    #[clap(long)]
    pub no_default_scuba_dataset: bool,
    /// Special dataset to be used by warm bookmark cache.  If a binary doesn't
    /// use warm bookmark cache then this parameter is ignored
    #[clap(long)]
    pub warm_bookmark_cache_scuba_dataset: Option<String>,
}

impl ScubaLoggingArgs {
    pub fn create_scuba_sample_builder(
        &self,
        fb: FacebookInit,
        observability_context: &ObservabilityContext,
        default_scuba_set: &Option<String>,
    ) -> Result<MononokeScubaSampleBuilder> {
        let scuba_logger = if let Some(scuba_dataset) = &self.scuba_dataset {
            MononokeScubaSampleBuilder::new(fb, scuba_dataset.as_str())?
        } else if let Some(default_scuba_dataset) = default_scuba_set {
            if self.no_default_scuba_dataset {
                MononokeScubaSampleBuilder::with_discard()
            } else {
                MononokeScubaSampleBuilder::new(fb, default_scuba_dataset)?
            }
        } else {
            MononokeScubaSampleBuilder::with_discard()
        };
        let mut scuba_logger = scuba_logger
            .with_observability_context(observability_context.clone())
            .with_seq("seq");

        scuba_logger.add_common_server_data();

        Ok(scuba_logger)
    }

    pub fn create_warm_bookmark_cache_scuba_sample_builder(
        &self,
        fb: FacebookInit,
    ) -> Result<MononokeScubaSampleBuilder, Error> {
        let maybe_scuba = match self.warm_bookmark_cache_scuba_dataset.clone() {
            Some(scuba) => {
                let hostname = hostname::get_hostname()?;
                let sampling_pct = tunables()
                    .warm_bookmark_cache_logging_sampling_pct()
                    .unwrap_or_default() as u64;
                let mut hasher = DefaultHasher::new();
                hostname.hash(&mut hasher);

                if hasher.finish() % 100 < sampling_pct {
                    Some(scuba)
                } else {
                    None
                }
            }
            None => None,
        };

        MononokeScubaSampleBuilder::with_opt_table(fb, maybe_scuba)
    }
}


/*
 * 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::Result;
use clap::Args;

/// Command line arguments for shutdown timeout
#[derive(Args, Debug)]
pub struct ShutdownTimeoutArgs {
    /// Number of seconds to wait after receiving a shutdown signal before shutting down
    #[clap(long, default_value = "0", value_parser = duration_secs_from_str)]
    pub shutdown_grace_period: Duration,

    /// Number of seconds to wait for requests to complete during shutdown
    #[clap(long, default_value = "10", value_parser = duration_secs_from_str)]
    pub shutdown_timeout: Duration,
}

fn duration_secs_from_str(s: &str) -> Result<Duration> {
    Ok(Duration::from_secs(s.parse::<u64>()?))
}


/*
 * 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::any::TypeId;
use std::collections::HashMap;
use std::collections::HashSet;
use std::fs;
use std::future::Future;
use std::path::Path;
use std::sync::Arc;
use std::time::Duration;
use std::time::Instant;

use anyhow::anyhow;
use anyhow::bail;
use anyhow::Context;
use anyhow::Result;
use base_app::BaseApp;
use blobstore::Blobstore;
use blobstore_factory::BlobstoreOptions;
use blobstore_factory::ReadOnlyStorage;
use cached_config::ConfigStore;
use clap::ArgMatches;
use clap::Error as ClapError;
use clap::FromArgMatches;
use context::CoreContext;
use environment::MononokeEnvironment;
use facet::AsyncBuildable;
use fbinit::FacebookInit;
use futures::stream;
use futures::stream::StreamExt;
use futures::stream::TryStreamExt;
use futures_util::try_join;
use metaconfig_parser::RepoConfigs;
use metaconfig_parser::StorageConfigs;
use metaconfig_types::BlobConfig;
use metaconfig_types::BlobstoreId;
use metaconfig_types::Redaction;
use metaconfig_types::RepoConfig;
use metaconfig_types::ShardedService;
use mononoke_configs::MononokeConfigs;
use mononoke_types::RepositoryId;
use prefixblob::PrefixBlobstore;
use redactedblobstore::RedactedBlobstore;
use redactedblobstore::RedactedBlobstoreConfig;
use redactedblobstore::RedactionConfigBlobstore;
use repo_factory::RepoFactory;
use repo_factory::RepoFactoryBuilder;
use running::run_until_terminated;
use scuba_ext::MononokeScubaSampleBuilder;
use services::Fb303Service;
use slog::error;
use slog::info;
use slog::o;
use slog::Logger;
use sql_ext::facebook::MysqlOptions;
use stats::prelude::*;
#[cfg(not(test))]
use stats::schedule_stats_aggregation_preview;
use tokio::runtime::Handle;
use tokio::runtime::Runtime;
use tokio::sync::oneshot;

use crate::args::AsRepoArg;
use crate::args::ConfigArgs;
use crate::args::ConfigMode;
use crate::args::MultiRepoArgs;
use crate::args::RepoArg;
use crate::args::RepoBlobstoreArgs;
use crate::args::SourceAndTargetRepoArgs;
use crate::extension::AppExtension;
use crate::extension::AppExtensionArgsBox;
use crate::extension::BoxedAppExtensionArgs;
use crate::fb303::Fb303AppExtension;
use crate::repos_manager::MononokeReposManager;

define_stats! {
    prefix = "mononoke.app";
    completion_duration_secs: timeseries(Average, Sum, Count),
}

pub struct MononokeApp {
    pub fb: FacebookInit,
    config_mode: ConfigMode,
    args: ArgMatches,
    runtime: Option<Runtime>,
    env: Arc<MononokeEnvironment>,
    extension_args: HashMap<TypeId, Box<dyn BoxedAppExtensionArgs>>,
    configs: Arc<MononokeConfigs>,
    repo_factory: Arc<RepoFactory>,
}

impl BaseApp for MononokeApp {
    fn subcommand(&self) -> Option<(&str, &ArgMatches)> {
        self.args.subcommand()
    }
}

impl MononokeApp {
    pub(crate) fn new(
        fb: FacebookInit,
        config_mode: ConfigMode,
        args: ArgMatches,
        runtime: Runtime,
        env: MononokeEnvironment,
        extension_args: HashMap<TypeId, Box<dyn BoxedAppExtensionArgs>>,
    ) -> Result<Self> {
        let env = Arc::new(env);
        let config_path = ConfigArgs::from_arg_matches(&args)?.config_path();
        let config_store = &env.as_ref().config_store;
        let configs = Arc::new(MononokeConfigs::new(
            config_path,
            config_store,
            runtime.handle().clone(),
            env.logger.clone(),
        )?);

        let repo_factory = Arc::new(RepoFactory::new(env.clone()));

        Ok(MononokeApp {
            fb,
            config_mode,
            args,
            runtime: Some(runtime),
            env,
            extension_args,
            configs,
            repo_factory,
        })
    }

    pub fn extension_args<Ext>(&self) -> Result<&Ext::Args>
    where
        Ext: AppExtension + 'static,
    {
        if let Some(ext) = self.extension_args.get(&TypeId::of::<Ext>()) {
            if let Some(ext) = ext.as_any().downcast_ref::<AppExtensionArgsBox<Ext>>() {
                return Ok(ext.args());
            }
        }
        Err(anyhow!(
            "Extension {} arguments not found (was it registered with MononokeApp?)",
            std::any::type_name::<Ext>(),
        ))
    }

    /// Start the FB303 monitoring server for the provided service.
    pub fn start_monitoring<Service>(&self, app_name: &str, service: Service) -> Result<()>
    where
        Service: Fb303Service + Sync + Send + 'static,
    {
        let fb303_args = self.extension_args::<Fb303AppExtension>()?;
        fb303_args.start_fb303_server(self.fb, app_name, self.logger(), service)?;
        Ok(())
    }

    /// Start the background stats aggregation thread.
    pub fn start_stats_aggregation(&self) -> Result<()> {
        #[cfg(not(test))]
        {
            self.env.runtime.block_on(async move {
                let stats_aggregation = schedule_stats_aggregation_preview()
                    .map_err(|_| anyhow!("Failed to create stats aggregation worker"))?;
                tokio::task::spawn(stats_aggregation);
                anyhow::Ok(())
            })?;
        }
        Ok(())
    }

    /// Execute a future on this app's runtime.
    ///
    /// If you are looking for a replacement for `cmdlib::helpers::block_execute`, prefer
    /// `run_with_monitoring_and_logging`.
    pub fn run_basic<F, Fut>(mut self, main: F) -> Result<()>
    where
        F: Fn(MononokeApp) -> Fut,
        Fut: Future<Output = Result<()>> + Send + 'static,
    {
        let runtime = self
            .runtime
            .take()
            .ok_or_else(|| anyhow!("MononokeApp already started"))?;
        runtime.block_on(async move { tokio::spawn(main(self)).await? })
    }

    /// Execute a future on this app's runtime.
    ///
    /// This future will run with monitoring enabled, and errors will be logged to glog.
    pub fn run_with_monitoring_and_logging<F, Fut, Service>(
        mut self,
        main: F,
        app_name: &str,
        service: Service,
    ) -> Result<()>
    where
        F: Fn(MononokeApp) -> Fut,
        Fut: Future<Output = Result<()>>,
        Service: Fb303Service + Sync + Send + 'static,
    {
        self.start_monitoring(app_name, service)?;
        self.start_stats_aggregation()?;

        let logger = self.logger().clone();
        let runtime = self
            .runtime
            .take()
            .ok_or_else(|| anyhow!("MononokeApp already started"))?;
        let result = runtime.block_on(main(self));

        if let Err(e) = result {
            // Log error in glog format
            error!(&logger, "Execution error: {:?}", e);

            // Replace the error with a simple error so it isn't logged twice.
            return Err(anyhow!("Execution failed"));
        }

        Ok(())
    }

    /// Run a server future, and wait until a termination signal is received.
    ///
    /// When the termination signal is received, the `quiesce` callback is
    /// called.  This should perform any steps required to quiesce the server,
    /// for example by removing this instance from routing configuration, or
    /// asking the load balancer to stop sending requests to this instance.
    /// Requests that do arrive should still be accepted.
    ///
    /// After the `shutdown_grace_period`, the `shutdown` future is awaited.
    /// This should do any additional work to stop accepting connections and wait
    /// until all outstanding requests have been handled. The `server` future
    /// continues to run while `shutdown` is being awaited.
    ///
    /// Once both `shutdown` and `server` have completed, the process
    /// exits. If `shutdown_timeout` is exceeded, the server future is canceled
    /// and an error is returned.
    pub fn run_until_terminated<ServerFn, ServerFut, QuiesceFn, ShutdownFut>(
        mut self,
        server: ServerFn,
        quiesce: QuiesceFn,
        shutdown_grace_period: Duration,
        shutdown: ShutdownFut,
        shutdown_timeout: Duration,
    ) -> Result<()>
    where
        ServerFn: FnOnce(MononokeApp) -> ServerFut + Send + 'static,
        ServerFut: Future<Output = Result<()>> + Send + 'static,
        QuiesceFn: FnOnce(),
        ShutdownFut: Future<Output = ()>,
    {
        let logger = self.logger().clone();
        let runtime = self
            .runtime
            .take()
            .ok_or_else(|| anyhow!("MononokeApp already started"))?;
        let server = async move { server(self).await };
        runtime.block_on(run_until_terminated(
            server,
            &logger,
            quiesce,
            shutdown_grace_period,
            shutdown,
            shutdown_timeout,
        ))
    }

    /// Wait until a termination signal is received.
    ///
    /// This method does not have a server future, and so is useful when all
    /// serving listeners are running on another executor (e.g. a C++
    /// executor for a thrift service).
    ///
    /// When the termination signal is received, the same quiesce-shutdown
    /// procedure as for `run_until_terminated` is followed.
    pub fn wait_until_terminated<QuiesceFn, ShutdownFut>(
        self,
        quiesce: QuiesceFn,
        shutdown_grace_period: Duration,
        shutdown: ShutdownFut,
        shutdown_timeout: Duration,
    ) -> Result<()>
    where
        QuiesceFn: FnOnce(),
        ShutdownFut: Future<Output = ()>,
    {
        let (exit_tx, exit_rx) = oneshot::channel();
        let server = move |app| async move {
            exit_rx.await?;
            drop(app);
            Ok(())
        };

        self.run_until_terminated(
            server,
            || {
                let _ = exit_tx.send(());
                quiesce();
            },
            shutdown_grace_period,
            shutdown,
            shutdown_timeout,
        )
    }

    /// Returns the selected subcommand of the app (if this app
    /// has subcommands).
    pub fn matches(&self) -> &ArgMatches {
        &self.args
    }

    /// Returns a parsed args struct based on the arguments provided
    /// on the command line.
    pub fn args<Args>(&self) -> Result<Args, ClapError>
    where
        Args: FromArgMatches,
    {
        Args::from_arg_matches(&self.args)
    }

    /// Returns a handle to this app's runtime.
    pub fn runtime(&self) -> &Handle {
        &self.env.runtime
    }

    /// The config store for this app.
    pub fn config_store(&self) -> &ConfigStore {
        &self.env.config_store
    }

    /// The repo configs for this app.
    pub fn repo_configs(&self) -> Arc<RepoConfigs> {
        self.configs.repo_configs()
    }

    /// The storage configs for this app.
    pub fn storage_configs(&self) -> Arc<StorageConfigs> {
        self.configs.storage_configs()
    }

    /// The logger for this app.
    pub fn logger(&self) -> &Logger {
        &self.env.logger
    }

    /// Construct a logger for a specific repo.
    pub fn repo_logger(&self, repo_name: &str) -> Logger {
        self.env.logger.new(o!("repo" => repo_name.to_string()))
    }

    /// The mysql options for this app.
    pub fn mysql_options(&self) -> &MysqlOptions {
        &self.env.mysql_options
    }

    /// The blobstore options for this app.
    pub fn blobstore_options(&self) -> &BlobstoreOptions {
        &self.env.blobstore_options
    }

    /// The readonly storage options for this app.
    pub fn readonly_storage(&self) -> &ReadOnlyStorage {
        &self.env.readonly_storage
    }

    /// Create a basic CoreContext without scuba logging.  Good choice for
    /// simple CLI tools like admin.
    ///
    /// Warning: returned context doesn't provide any scuba logging!
    pub fn new_basic_context(&self) -> CoreContext {
        CoreContext::new_with_logger(self.env.fb, self.logger().clone())
    }

    /// Return repo factory used by app.
    pub fn repo_factory(&self) -> &Arc<RepoFactory> {
        &self.repo_factory
    }

    /// Mononoke environment for this app.
    pub fn environment(&self) -> &Arc<MononokeEnvironment> {
        &self.env
    }

    /// Returns true if this is a production configuration of Mononoke
    pub fn is_production(&self) -> bool {
        self.config_mode == ConfigMode::Production
    }

    pub fn repo_config_by_name(&self, repo_name: &str) -> Result<RepoConfig> {
        self.repo_configs()
            .repos
            .get(repo_name)
            .cloned()
            .ok_or_else(|| anyhow!("unknown reponame: {:?}", repo_name))
    }

    /// Get repo config based on user-provided arguments.
    pub fn repo_config(&self, repo_arg: &RepoArg) -> Result<(String, RepoConfig)> {
        match repo_arg {
            RepoArg::Id(repo_id) => {
                let repo_configs = self.repo_configs();
                let (repo_name, repo_config) = repo_configs
                    .get_repo_config(*repo_id)
                    .ok_or_else(|| anyhow!("unknown repoid: {:?}", repo_id))?;
                Ok((repo_name.clone(), repo_config.clone()))
            }
            RepoArg::Name(repo_name) => {
                let repo_config = self.repo_config_by_name(repo_name)?;
                Ok((repo_name.to_string(), repo_config))
            }
        }
    }

    /// Get repo configs based on user-provided arguments.
    pub fn multi_repo_configs(&self, repo_args: Vec<RepoArg>) -> Result<Vec<(String, RepoConfig)>> {
        let mut repos = vec![];
        let mut unique_repos = HashSet::new();
        for repo in repo_args {
            let (name, repo_conf) = self.repo_config(&repo)?;
            if unique_repos.insert(name.clone()) {
                repos.push((name, repo_conf));
            }
        }

        Ok(repos)
    }

    /// Open repositories based on user-provided arguments.
    pub async fn open_repos<Repo>(&self, repos_args: &MultiRepoArgs) -> Result<Vec<Repo>>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>,
    {
        let args = repos_args.ids_or_names()?;
        let mut repos = vec![];
        for arg in args {
            repos.push(self.repo_config(&arg)?);
        }

        let repos: HashMap<_, _> = repos.into_iter().collect();
        let common_config = self.repo_configs().common.clone();
        let repos: Vec<_> = stream::iter(repos)
            .map(|(repo_name, repo_config)| {
                let repo_factory = self.repo_factory.clone();
                let common_config = common_config.clone();
                async move {
                    repo_factory
                        .build(repo_name, repo_config, common_config)
                        .await
                }
            })
            .buffered(100)
            .try_collect()
            .await?;

        Ok(repos)
    }

    /// Open a repository based on user-provided arguments.
    pub async fn open_repo<Repo>(&self, repo_args: &impl AsRepoArg) -> Result<Repo>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>,
    {
        let repo_arg = repo_args.as_repo_arg();
        let (repo_name, repo_config) = self.repo_config(repo_arg)?;
        let common_config = self.repo_configs().common.clone();
        let repo = self
            .repo_factory
            .build(repo_name, repo_config, common_config)
            .await?;
        Ok(repo)
    }

    /// Open an existing repo object
    /// Make sure that the opened repo has redaction DISABLED
    pub async fn open_repo_unredacted<Repo>(&self, repo_args: &impl AsRepoArg) -> Result<Repo>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>,
    {
        let repo_arg = repo_args.as_repo_arg();
        let (repo_name, mut repo_config) = self.repo_config(repo_arg)?;
        let common_config = self.repo_configs().common.clone();
        repo_config.redaction = Redaction::Disabled;
        let repo = self
            .repo_factory
            .build(repo_name, repo_config, common_config)
            .await?;
        Ok(repo)
    }

    /// Create a new repo object -- for local instances, expect its contents to be empty.
    /// Makes sure that the opened repo has redaction DISABLED
    pub async fn create_repo_unredacted<Repo>(&self, repo_arg: &impl AsRepoArg) -> Result<Repo>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>,
    {
        let (repo_name, mut repo_config) = self.repo_config(repo_arg.as_repo_arg())?;
        let common_config = self.repo_configs().common.clone();

        match &repo_config.storage_config.blobstore {
            BlobConfig::Files { path } | BlobConfig::Sqlite { path } => {
                setup_repo_dir(path, CreateStorage::ExistingOrCreate)?;
            }
            _ => {}
        }
        repo_config.redaction = Redaction::Disabled;
        let repo = self
            .repo_factory
            .build(repo_name, repo_config, common_config)
            .await?;
        Ok(repo)
    }

    /// Open a source and target repos based on user-provided arguments.
    pub async fn open_source_and_target_repos<Repo>(
        &self,
        repo_args: &SourceAndTargetRepoArgs,
    ) -> Result<(Repo, Repo)>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>,
    {
        let (source_repo_name, source_repo_config) =
            self.repo_config(repo_args.source_repo.as_repo_arg())?;
        let (target_repo_name, target_repo_config) =
            self.repo_config(repo_args.target_repo.as_repo_arg())?;
        let common_config = self.repo_configs().common.clone();
        let source_repo_fut =
            self.repo_factory
                .build(source_repo_name, source_repo_config, common_config.clone());
        let target_repo_fut =
            self.repo_factory
                .build(target_repo_name, target_repo_config, common_config);

        let (source_repo, target_repo) = try_join!(source_repo_fut, target_repo_fut)?;
        Ok((source_repo, target_repo))
    }

    /// Create a manager for all configured repos based on deep-sharding status, excluding
    /// those filtered by `repo_filter_from` in `MononokeEnvironment`.
    pub async fn open_managed_repos<Repo>(
        &self,
        service: Option<ShardedService>,
    ) -> Result<MononokeReposManager<Repo>>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let redaction_disabled = false;
        self.open_managed_repos_with_redaction_disabled(service, redaction_disabled)
            .await
    }

    /// Create a manager for all configured repos based on deep-sharding status, excluding
    /// those filtered by `repo_filter_from` in `MononokeEnvironment`.
    /// When reloading the config, make sure that the opened repo has redaction DISABLED
    pub async fn open_managed_repos_unredacted<Repo>(
        &self,
        service: Option<ShardedService>,
    ) -> Result<MononokeReposManager<Repo>>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let redaction_disabled = true;
        self.open_managed_repos_with_redaction_disabled(service, redaction_disabled)
            .await
    }

    /// Create a manager for all configured repos based on deep-sharding status, excluding
    /// those filtered by `repo_filter_from` in `MononokeEnvironment`.
    async fn open_managed_repos_with_redaction_disabled<Repo>(
        &self,
        service: Option<ShardedService>,
        redaction_disabled: bool,
    ) -> Result<MononokeReposManager<Repo>>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let repo_filter = self.environment().filter_repos.clone();
        let service_name = service.clone();
        let repo_names =
            self.repo_configs()
                .repos
                .clone()
                .into_iter()
                .filter_map(|(name, config)| {
                    let is_matching_filter =
                        repo_filter.as_ref().map_or(true, |filter| filter(&name));
                    let is_deep_sharded = service
                        .as_ref()
                        .and_then(|service| {
                            config
                                .deep_sharding_config
                                .and_then(|c| c.status.get(service).copied())
                        })
                        .unwrap_or(false);
                    // Initialize repos that are enabled and not deep-sharded (i.e. need to exist
                    // at service startup)
                    if config.enabled && !is_deep_sharded && is_matching_filter {
                        Some(name)
                    } else {
                        None
                    }
                });
        self.open_named_managed_repos_with_redaction_disabled(
            repo_names,
            service_name,
            redaction_disabled,
        )
        .await
    }

    /// Create a manager for a set of named managed repos.  These repos must
    /// be configured in the config.
    pub async fn open_named_managed_repos<Repo, Names>(
        &self,
        repo_names: Names,
        service: Option<ShardedService>,
    ) -> Result<MononokeReposManager<Repo>>
    where
        Names: IntoIterator<Item = String>,
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let redaction_disabled = false;
        self.open_named_managed_repos_with_redaction_disabled(
            repo_names,
            service,
            redaction_disabled,
        )
        .await
    }

    /// Create a manager for a set of named managed repos.  These repos must
    /// be configured in the config.
    /// When reloading the config, make sure that the opened repo has redaction DISABLED
    pub async fn open_named_managed_repos_unredacted<Repo, Names>(
        &self,
        repo_names: Names,
        service: Option<ShardedService>,
    ) -> Result<MononokeReposManager<Repo>>
    where
        Names: IntoIterator<Item = String>,
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let redaction_disabled = true;
        self.open_named_managed_repos_with_redaction_disabled(
            repo_names,
            service,
            redaction_disabled,
        )
        .await
    }

    async fn open_named_managed_repos_with_redaction_disabled<Repo, Names>(
        &self,
        repo_names: Names,
        service: Option<ShardedService>,
        redaction_disabled: bool,
    ) -> Result<MononokeReposManager<Repo>>
    where
        Names: IntoIterator<Item = String>,
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let logger = self.logger().clone();
        let start = Instant::now();
        let repos_mgr = MononokeReposManager::new_with_redaction_disabled(
            self.configs.clone(),
            self.repo_factory().clone(),
            self.logger().clone(),
            service,
            repo_names,
            redaction_disabled,
        )
        .await?;
        info!(
            &logger,
            "All repos initialized. It took: {} seconds",
            start.elapsed().as_secs()
        );
        STATS::completion_duration_secs
            .add_value(start.elapsed().as_secs().try_into().unwrap_or(i64::MAX));
        Ok(repos_mgr)
    }

    /// Create a manager for a single repo, specified by repo arguments.
    pub async fn open_managed_repo_arg<Repo>(
        &self,
        repo_arg: &impl AsRepoArg,
    ) -> Result<MononokeReposManager<Repo>>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let (repo_name, _) = self.repo_config(repo_arg.as_repo_arg())?;
        self.open_named_managed_repos(Some(repo_name), None).await
    }

    /// Open just the blobstore based on user-provided arguments.
    pub async fn open_blobstore(
        &self,
        repo_blobstore_args: &RepoBlobstoreArgs,
    ) -> Result<Arc<dyn Blobstore>> {
        let repo_configs = self.repo_configs();
        let storage_configs = self.storage_configs();
        let (mut repo_id, redaction, storage_config) =
            if let Some(repo_id) = repo_blobstore_args.repo_id {
                let repo_id = RepositoryId::new(repo_id);
                let (_repo_name, repo_config) = repo_configs
                    .get_repo_config(repo_id)
                    .ok_or_else(|| anyhow!("unknown repoid: {:?}", repo_id))?;
                (
                    Some(repo_id),
                    repo_config.redaction,
                    repo_config.storage_config.clone(),
                )
            } else if let Some(repo_name) = &repo_blobstore_args.repo_name {
                let repo_config = repo_configs
                    .repos
                    .get(repo_name)
                    .ok_or_else(|| anyhow!("unknown reponame: {:?}", repo_name))?;
                (
                    Some(repo_config.repoid),
                    repo_config.redaction,
                    repo_config.storage_config.clone(),
                )
            } else if let Some(storage_name) = &repo_blobstore_args.storage_name {
                let storage_config = storage_configs
                    .storage
                    .get(storage_name)
                    .ok_or_else(|| anyhow!("unknown storage name: {:?}", storage_name))?;
                (None, Redaction::Enabled, storage_config.clone())
            } else {
                return Err(anyhow!("Expected a storage argument"));
            };

        let blob_config = match repo_blobstore_args.inner_blobstore_id {
            None => storage_config.blobstore,
            Some(id) => match storage_config.blobstore {
                BlobConfig::MultiplexedWal { blobstores, .. } => {
                    let sought_id = BlobstoreId::new(id);
                    blobstores
                        .into_iter()
                        .find_map(|(blobstore_id, _, blobstore)| {
                            if blobstore_id == sought_id {
                                Some(blobstore)
                            } else {
                                None
                            }
                        })
                        .ok_or_else(|| anyhow!("could not find a blobstore with id {}", id))?
                }
                _ => {
                    return Err(anyhow!(
                        "inner-blobstore-id supplied by blobstore is not multiplexed"
                    ));
                }
            },
        };

        if repo_blobstore_args.no_prefix {
            repo_id = None;
        }

        let blobstore = blobstore_factory::make_blobstore(
            self.env.fb,
            blob_config,
            &self.env.mysql_options,
            self.env.readonly_storage,
            &self.env.blobstore_options,
            &self.env.logger,
            &self.env.config_store,
            &blobstore_factory::default_scrub_handler(),
            None,
        )
        .await?;

        let blobstore = if let Some(repo_id) = repo_id {
            PrefixBlobstore::new(blobstore, repo_id.prefix())
        } else {
            PrefixBlobstore::new(blobstore, String::new())
        };

        let blobstore = if redaction == Redaction::Enabled {
            let redacted_blobs = self
                .repo_factory
                .redacted_blobs(
                    self.new_basic_context(),
                    &storage_config.metadata,
                    &Arc::new(self.repo_configs().common.clone()),
                )
                .await?;
            RedactedBlobstore::new(
                blobstore,
                RedactedBlobstoreConfig::new(Some(redacted_blobs), self.redaction_scuba_builder()?),
            )
            .boxed()
        } else {
            Arc::new(blobstore)
        };

        Ok(blobstore)
    }

    pub async fn redaction_config_blobstore(&self) -> Result<Arc<RedactionConfigBlobstore>> {
        self.repo_factory
            .redaction_config_blobstore_from_config(
                &self.repo_configs().common.redaction_config.blobstore,
            )
            .await
    }

    pub async fn redaction_config_blobstore_for_darkstorm(
        &self,
    ) -> Result<Arc<RedactionConfigBlobstore>> {
        let blobstore_config = self
            .repo_configs()
            .common
            .redaction_config
            .darkstorm_blobstore
            .clone()
            .ok_or_else(|| anyhow!("Configuration must have darkstorm blobstore"))?;
        self.repo_factory
            .redaction_config_blobstore_from_config(&blobstore_config)
            .await
    }

    fn redaction_scuba_builder(&self) -> Result<MononokeScubaSampleBuilder> {
        let params = &self.repo_configs().common.censored_scuba_params;
        let mut builder =
            MononokeScubaSampleBuilder::with_opt_table(self.env.fb, params.table.clone())?;
        if let Some(file) = &params.local_path {
            builder = builder
                .with_log_file(file)
                .context("Failed to open scuba log file")?;
        }

        Ok(builder)
    }
}

#[derive(Clone, Debug, PartialEq, Eq, Hash)]
pub enum CreateStorage {
    ExistingOnly,
    ExistingOrCreate,
}

pub fn setup_repo_dir<P: AsRef<Path>>(data_dir: P, create: CreateStorage) -> Result<()> {
    let data_dir = data_dir.as_ref();

    if !data_dir.is_dir() {
        bail!("{:?} does not exist or is not a directory", data_dir);
    }

    // Validate directory layout
    #[allow(clippy::single_element_loop)]
    for subdir in &["blobs"] {
        let subdir = data_dir.join(subdir);

        if subdir.exists() && !subdir.is_dir() {
            bail!("{:?} already exists and is not a directory", subdir);
        }

        if !subdir.exists() {
            if CreateStorage::ExistingOnly == create {
                bail!("{:?} not found in ExistingOnly mode", subdir,);
            }
            fs::create_dir(&subdir)
                .with_context(|| format!("failed to create subdirectory {:?}", subdir))?;
        }
    }
    Ok(())
}


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

use std::any::TypeId;
use std::collections::HashMap;
use std::sync::Arc;
use std::time::Duration;

use anyhow::Context;
use anyhow::Result;
use arg_extensions::ArgDefaults;
use blobstore_factory::BlobstoreArgs;
use blobstore_factory::BlobstoreOptions;
use blobstore_factory::CachelibBlobstoreOptions;
use blobstore_factory::ChaosOptions;
use blobstore_factory::DelayOptions;
#[cfg(fbcode_build)]
use blobstore_factory::ManifoldArgs;
use blobstore_factory::PackOptions;
use blobstore_factory::ReadOnlyStorage;
use blobstore_factory::ReadOnlyStorageArgs;
use blobstore_factory::ThrottleOptions;
use cached_config::ConfigHandle;
use cached_config::ConfigStore;
use clap::Args;
use clap::Command;
use clap::CommandFactory;
use clap::FromArgMatches;
use cmdlib_caching::init_cachelib;
use cmdlib_caching::CachelibArgs;
use cmdlib_caching::CachelibSettings;
use cmdlib_logging::LoggingArgs;
use cmdlib_logging::ScubaLoggingArgs;
use derived_data_remote::RemoteDerivationArgs;
use environment::MononokeEnvironment;
use environment::WarmBookmarksCacheDerivedData;
use fbinit::FacebookInit;
use megarepo_config::MegarepoConfigsArgs;
use megarepo_config::MononokeMegarepoConfigsOptions;
use observability::DynamicLevelDrain;
use permission_checker::AclProvider;
use permission_checker::DefaultAclProvider;
use permission_checker::InternalAclProvider;
use rendezvous::RendezVousArgs;
use slog::debug;
use slog::o;
use slog::Logger;
use slog::Never;
use slog::SendSyncRefUnwindSafeDrain;
use sql_ext::facebook::MysqlOptions;
use sql_ext::facebook::PoolConfig;
use sql_ext::facebook::ReadConnectionType;
use sql_ext::facebook::SharedConnectionPool;
use tokio::runtime::Handle;
use tokio::runtime::Runtime;

use crate::app::MononokeApp;
use crate::args::parse_config_spec_to_path;
use crate::args::AclArgs;
use crate::args::ConfigArgs;
use crate::args::MysqlArgs;
use crate::args::RuntimeArgs;
use crate::args::TunablesArgs;
use crate::extension::AppExtension;
use crate::extension::AppExtensionBox;
use crate::extension::BoxedAppExtension;
use crate::extension::BoxedAppExtensionArgs;

pub struct MononokeAppBuilder {
    fb: FacebookInit,
    extensions: Vec<(TypeId, Box<dyn BoxedAppExtension>)>,
    arg_defaults: Vec<Box<dyn ArgDefaults>>,
    cachelib_settings: CachelibSettings,
    default_scuba_dataset: Option<String>,
    defaults: HashMap<&'static str, String>,
    warm_bookmarks_cache_derived_data: Option<WarmBookmarksCacheDerivedData>,
}

#[derive(Args, Debug)]
pub struct EnvironmentArgs {
    #[clap(flatten, next_help_heading = "CONFIG OPTIONS")]
    config_args: ConfigArgs,

    #[clap(flatten, next_help_heading = "LOGGING OPTIONS")]
    logging_args: LoggingArgs,

    #[clap(flatten, next_help_heading = "SCUBA LOGGING OPTIONS")]
    scuba_logging_args: ScubaLoggingArgs,

    #[clap(flatten, next_help_heading = "CACHELIB OPTIONS")]
    cachelib_args: CachelibArgs,

    #[clap(flatten, next_help_heading = "MYSQL OPTIONS")]
    mysql_args: MysqlArgs,

    #[clap(flatten, next_help_heading = "TUNABLES OPTIONS")]
    tunables_args: TunablesArgs,

    #[clap(flatten, next_help_heading = "BLOBSTORE OPTIONS")]
    blobstore_args: BlobstoreArgs,

    #[cfg(fbcode_build)]
    #[clap(flatten, next_help_heading = "MANIFOLD OPTIONS")]
    manifold_args: ManifoldArgs,

    #[clap(flatten, next_help_heading = "ACL OPTIONS")]
    acl_args: AclArgs,

    #[clap(flatten, next_help_heading = "REMOTE DERIVATION OPTIONS")]
    remote_derivation_args: RemoteDerivationArgs,

    #[clap(flatten, next_help_heading = "STORAGE OPTIONS")]
    readonly_storage_args: ReadOnlyStorageArgs,

    #[clap(flatten, next_help_heading = "RENDEZ-VOUS OPTIONS")]
    rendezvous_args: RendezVousArgs,

    #[clap(flatten, next_help_heading = "MEGAREPO OPTIONS")]
    megarepo_configs_args: MegarepoConfigsArgs,
}

impl MononokeAppBuilder {
    pub fn new(fb: FacebookInit) -> Self {
        MononokeAppBuilder {
            fb,
            extensions: Vec::new(),
            arg_defaults: Vec::new(),
            cachelib_settings: CachelibSettings::default(),
            default_scuba_dataset: None,
            defaults: HashMap::new(),
            warm_bookmarks_cache_derived_data: None,
        }
    }

    pub fn with_arg_defaults(mut self, arg_defaults: impl ArgDefaults + 'static) -> Self {
        self.arg_defaults.push(Box::new(arg_defaults));
        self
    }

    pub fn with_default_scuba_dataset(mut self, default: impl Into<String>) -> Self {
        self.default_scuba_dataset = Some(default.into());
        self
    }

    pub fn with_warm_bookmarks_cache(
        mut self,
        warm_bookmarks_cache_derived_data: WarmBookmarksCacheDerivedData,
    ) -> Self {
        self.warm_bookmarks_cache_derived_data = Some(warm_bookmarks_cache_derived_data);
        self
    }

    pub fn with_default_cachelib_settings(mut self, cachelib_settings: CachelibSettings) -> Self {
        self.cachelib_settings = cachelib_settings;
        self
    }

    pub fn with_app_extension<Ext>(mut self, ext: Ext) -> Self
    where
        Ext: AppExtension + 'static,
    {
        self.extensions
            .push((TypeId::of::<Ext>(), AppExtensionBox::new(ext)));
        self
    }

    pub fn build<AppArgs>(&mut self) -> Result<MononokeApp>
    where
        AppArgs: CommandFactory,
    {
        self.build_with_subcommands::<AppArgs>(Vec::new())
    }

    pub fn build_with_subcommands<'a, AppArgs>(
        &'a mut self,
        subcommands: Vec<Command>,
    ) -> Result<MononokeApp>
    where
        AppArgs: CommandFactory,
    {
        for (arg, default) in self.cachelib_settings.arg_defaults() {
            self.defaults.insert(arg, default);
        }

        for defaults in self.arg_defaults.iter() {
            for (arg, default) in defaults.arg_defaults() {
                self.defaults.insert(arg, default);
            }
        }

        for (_type_id, ext) in self.extensions.iter() {
            for (arg, default) in ext.arg_defaults() {
                self.defaults.insert(arg, default);
            }
        }

        let mut app = AppArgs::command();

        {
            // Save app-generated about so we can restore it.
            let about = app.get_about().cloned();
            let long_about = app.get_long_about().cloned();

            app = RuntimeArgs::augment_args_for_update(app);
            app = EnvironmentArgs::augment_args_for_update(app);
            for (_type_id, ext) in self.extensions.iter() {
                app = ext.augment_args(app);
            }

            // Adding the additional args overrode the about messages.
            // Restore them.
            if let Some(about) = about {
                app = app.about(about);
            }
            if let Some(long_about) = long_about {
                app = app.long_about(long_about);
            }
        }

        if !subcommands.is_empty() {
            app = app
                .subcommands(subcommands)
                .subcommand_required(true)
                .arg_required_else_help(true);
        }

        for (name, default) in self.defaults.iter() {
            app = app.mut_arg(name, |arg| arg.default_value(default));
        }

        let args = app.get_matches();

        let extension_args = self
            .extensions
            .iter()
            .map(|(type_id, ext)| Ok((*type_id, ext.parse_args(&args)?)))
            .collect::<Result<Vec<_>>>()?;

        let runtime_args = RuntimeArgs::from_arg_matches(&args)?;
        let runtime = create_runtime(&runtime_args)?;

        let env_args = EnvironmentArgs::from_arg_matches(&args)?;
        let config_mode = env_args.config_args.mode();
        let mut env = self.build_environment(
            &runtime,
            env_args,
            extension_args.iter().map(|(_type_id, ext)| ext.as_ref()),
        )?;

        for (_type_id, ext) in extension_args.iter() {
            ext.environment_hook(&mut env)?;
        }

        MononokeApp::new(
            self.fb,
            config_mode,
            args,
            runtime,
            env,
            extension_args.into_iter().collect(),
        )
    }

    fn build_environment<'a>(
        &self,
        runtime: &Runtime,
        env_args: EnvironmentArgs,
        extension_args: impl IntoIterator<Item = &'a dyn BoxedAppExtensionArgs> + Clone,
    ) -> Result<MononokeEnvironment> {
        let EnvironmentArgs {
            blobstore_args,
            config_args,
            logging_args,
            scuba_logging_args,
            cachelib_args,
            #[cfg(fbcode_build)]
            manifold_args,
            megarepo_configs_args,
            mysql_args,
            readonly_storage_args,
            acl_args,
            remote_derivation_args,
            rendezvous_args,
            tunables_args,
        } = env_args;

        let log_level = logging_args.create_log_level();
        #[cfg(fbcode_build)]
        cmdlib_logging::glog::set_glog_log_level(self.fb, log_level)?;
        let root_log_drain = logging_args
            .create_root_log_drain(self.fb, log_level)
            .context("Failed to create root log drain")?;

        let config_store = create_config_store(
            self.fb,
            &config_args,
            Logger::root(root_log_drain.clone(), o![]),
        )
        .context("Failed to create config store")?;

        let observability_context = logging_args
            .create_observability_context(&config_store, log_level)
            .context("Failed to initialize observability context")?;

        let mut root_log_drain: Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>> =
            Arc::new(DynamicLevelDrain::new(
                root_log_drain,
                observability_context.clone(),
            ));
        for ext in extension_args {
            root_log_drain = ext.log_drain_hook(root_log_drain)?;
        }

        let logger = logging_args.create_logger(root_log_drain)?;

        let scuba_sample_builder = scuba_logging_args
            .create_scuba_sample_builder(
                self.fb,
                &observability_context,
                &self.default_scuba_dataset,
            )
            .context("Failed to create scuba sample builder")?;
        let warm_bookmarks_cache_scuba_sample_builder = scuba_logging_args
            .create_warm_bookmark_cache_scuba_sample_builder(self.fb)
            .context("Failed to create warm bookmark cache scuba sample builder")?;

        let caching = init_cachelib(self.fb, &self.cachelib_settings, &cachelib_args);

        let mysql_options =
            create_mysql_options(&mysql_args, create_mysql_pool_config(&mysql_args));

        let blobstore_options = create_blobstore_options(
            &blobstore_args,
            &mysql_args,
            #[cfg(fbcode_build)]
            manifold_args,
        )
        .context("Failed to parse blobstore options")?;

        let readonly_storage = ReadOnlyStorage::from_args(&readonly_storage_args);

        let rendezvous_options = rendezvous_args.into();

        let megarepo_configs_options = MononokeMegarepoConfigsOptions::from_args(
            config_args.local_configerator_path.as_deref(),
            &megarepo_configs_args,
        );

        let remote_derivation_options = remote_derivation_args.into();

        let acl_provider =
            create_acl_provider(self.fb, &acl_args).context("Failed to create ACL provider")?;

        init_tunables_worker(
            &tunables_args,
            &config_store,
            logger.clone(),
            runtime.handle().clone(),
        )?;

        Ok(MononokeEnvironment {
            fb: self.fb,
            logger,
            scuba_sample_builder,
            warm_bookmarks_cache_scuba_sample_builder,
            config_store,
            caching,
            observability_context,
            runtime: runtime.handle().clone(),
            mysql_options,
            blobstore_options,
            readonly_storage,
            acl_provider,
            rendezvous_options,
            megarepo_configs_options,
            remote_derivation_options,
            disabled_hooks: HashMap::new(),
            warm_bookmarks_cache_derived_data: self.warm_bookmarks_cache_derived_data,
            filter_repos: None,
        })
    }
}

fn create_config_store(
    fb: FacebookInit,
    config_args: &ConfigArgs,
    logger: Logger,
) -> Result<ConfigStore> {
    const CRYPTO_PROJECT: &str = "SCM";
    const CONFIGERATOR_POLL_INTERVAL: Duration = Duration::from_secs(1);
    const CONFIGERATOR_REFRESH_TIMEOUT: Duration = Duration::from_secs(1);

    if let Some(path) = &config_args.local_configerator_path {
        Ok(ConfigStore::file(
            logger,
            path.clone(),
            String::new(),
            CONFIGERATOR_POLL_INTERVAL,
        ))
    } else {
        let crypto_regex_paths = match &config_args.crypto_path_regex {
            Some(paths) => paths.clone(),
            None => vec![
                "scm/mononoke/tunables/.*".to_string(),
                "scm/mononoke/repos/.*".to_string(),
                "scm/mononoke/redaction/.*".to_string(),
            ],
        };
        let crypto_regex = crypto_regex_paths
            .into_iter()
            .map(|path| (path, CRYPTO_PROJECT.to_string()))
            .collect();
        ConfigStore::regex_signed_configerator(
            fb,
            logger,
            crypto_regex,
            CONFIGERATOR_POLL_INTERVAL,
            CONFIGERATOR_REFRESH_TIMEOUT,
        )
    }
}

fn create_runtime(runtime_args: &RuntimeArgs) -> Result<Runtime> {
    let mut builder = tokio::runtime::Builder::new_multi_thread();
    builder.enable_all();
    builder.thread_name("tk");
    if let Some(threads) = runtime_args.runtime_threads {
        builder.worker_threads(threads);
    }
    let runtime = builder.build()?;
    Ok(runtime)
}

fn create_mysql_options(mysql_args: &MysqlArgs, pool_config: PoolConfig) -> MysqlOptions {
    let pool = SharedConnectionPool::new();
    let read_connection_type = if mysql_args.mysql_master_only {
        ReadConnectionType::Master
    } else {
        ReadConnectionType::ReplicaOnly
    };
    MysqlOptions {
        pool,
        pool_config,
        read_connection_type,
    }
}

fn create_mysql_pool_config(mysql_args: &MysqlArgs) -> PoolConfig {
    PoolConfig::new(
        mysql_args.mysql_pool_limit,
        mysql_args.mysql_pool_threads_num,
        mysql_args.mysql_pool_per_key_limit,
        mysql_args.mysql_pool_age_timeout,
        mysql_args.mysql_pool_idle_timeout,
        mysql_args.mysql_conn_open_timeout,
        Duration::from_millis(mysql_args.mysql_query_time_limit),
    )
}

fn create_mysql_sqlblob_pool_config(mysql_args: &MysqlArgs) -> PoolConfig {
    PoolConfig::new(
        mysql_args.mysql_sqlblob_pool_limit,
        mysql_args.mysql_sqlblob_pool_threads_num,
        mysql_args.mysql_sqlblob_pool_per_key_limit,
        mysql_args.mysql_sqlblob_pool_age_timeout,
        mysql_args.mysql_sqlblob_pool_idle_timeout,
        mysql_args.mysql_conn_open_timeout,
        Duration::from_millis(mysql_args.mysql_query_time_limit),
    )
}

fn create_blobstore_options(
    blobstore_args: &BlobstoreArgs,
    mysql_args: &MysqlArgs,
    #[cfg(fbcode_build)] manifold_args: ManifoldArgs,
) -> Result<BlobstoreOptions> {
    let chaos_options = ChaosOptions::new(
        blobstore_args.blobstore_read_chaos_rate,
        blobstore_args.blobstore_write_chaos_rate,
    );

    let delay_options = DelayOptions {
        get_dist: blobstore_args.get_delay_distribution()?,
        put_dist: blobstore_args.put_delay_distribution()?,
    };

    let throttle_options = ThrottleOptions {
        read_qps: blobstore_args.blobstore_read_qps,
        write_qps: blobstore_args.blobstore_write_qps,
        read_bytes: blobstore_args.blobstore_read_bytes_s,
        write_bytes: blobstore_args.blobstore_write_bytes_s,
        read_burst_bytes: blobstore_args.blobstore_read_burst_bytes_s,
        write_burst_bytes: blobstore_args.blobstore_write_burst_bytes_s,
        bytes_min_count: blobstore_args.blobstore_bytes_min_throttle,
    };

    let pack_options = PackOptions::new(blobstore_args.put_format_override()?);

    let cachelib_blobstore_options =
        CachelibBlobstoreOptions::new_lazy(Some(blobstore_args.blobstore_cachelib_attempt_zstd));

    let blobstore_put_behaviour = blobstore_args.blobstore_put_behaviour;

    let mysql_sqlblob_options =
        create_mysql_options(mysql_args, create_mysql_sqlblob_pool_config(mysql_args));

    let blobstore_options = BlobstoreOptions::new(
        chaos_options,
        delay_options,
        throttle_options,
        #[cfg(fbcode_build)]
        manifold_args.into(),
        pack_options,
        cachelib_blobstore_options,
        blobstore_put_behaviour,
        mysql_sqlblob_options,
    );

    Ok(blobstore_options)
}

fn init_tunables_worker(
    tunables_args: &TunablesArgs,
    config_store: &ConfigStore,
    logger: Logger,
    handle: Handle,
) -> Result<()> {
    if tunables_args.disable_tunables {
        debug!(logger, "Tunables are disabled");
        return Ok(());
    }

    if let Some(tunables_local_path) = &tunables_args.tunables_local_path {
        let value = std::fs::read_to_string(tunables_local_path)
            .with_context(|| format!("failed to open tunables path {}", tunables_local_path))?;
        let config_handle = ConfigHandle::from_json(&value)
            .with_context(|| format!("failed to parse tunables at path {}", tunables_local_path))?;
        return tunables::init_tunables(&logger, &config_handle);
    }

    let tunables_config = tunables_args.tunables_config_or_default();
    let config_handle =
        config_store.get_config_handle(parse_config_spec_to_path(&tunables_config)?)?;

    tunables::init_tunables_worker(logger, config_handle, handle)
}

fn create_acl_provider(fb: FacebookInit, acl_args: &AclArgs) -> Result<Arc<dyn AclProvider>> {
    let acl_provider = match &acl_args.acl_file {
        Some(acl_file) => InternalAclProvider::from_file(acl_file).with_context(|| {
            format!("Failed to load ACLs from '{}'", acl_file.to_string_lossy())
        })?,
        None => DefaultAclProvider::new(fb),
    };
    Ok(acl_provider)
}


/*
 * 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::any::Any;
use std::sync::Arc;

use anyhow::Result;
use clap::ArgMatches;
use clap::Args;
use clap::Command;
use clap::FromArgMatches;
use environment::MononokeEnvironment;
use slog::Never;
use slog::SendSyncRefUnwindSafeDrain;

/// Trait implemented by things that need to extend the app building process,
/// including adding additional arguments and modifying the environment before
/// it is used to start Mononoke.
pub trait AppExtension: Send + Sync + 'static {
    /// Argument type to extend Mononoke arguments with.
    type Args: clap::Args + Send + Sync + 'static;

    /// Obtain default values for these arguments.
    fn arg_defaults(&self) -> Vec<(&'static str, String)> {
        Vec::new()
    }

    /// Hook executed after creating the environment before initializing Mononoke.
    fn environment_hook(&self, _args: &Self::Args, _env: &mut MononokeEnvironment) -> Result<()> {
        Ok(())
    }

    /// Hook executed after creating the log drain allowing for augmenting the logging.
    fn log_drain_hook(
        &self,
        _args: &Self::Args,
        drain: Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>,
    ) -> Result<Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>> {
        Ok(drain)
    }
}

// Internal trait to hide the concrete extension type.
pub(crate) trait BoxedAppExtension: Send + Sync + 'static {
    fn augment_args(&self, app: Command) -> Command;
    fn arg_defaults(&self) -> Vec<(&'static str, String)>;
    fn parse_args(&self, args: &ArgMatches) -> Result<Box<dyn BoxedAppExtensionArgs>>;
}

// Box to store an app extension.
#[derive(Clone)]
pub(crate) struct AppExtensionBox<Ext: AppExtension> {
    ext: Arc<Ext>,
}

impl<Ext: AppExtension> AppExtensionBox<Ext> {
    pub(crate) fn new(ext: Ext) -> Box<dyn BoxedAppExtension> {
        Box::new(AppExtensionBox { ext: Arc::new(ext) })
    }
}

impl<Ext: AppExtension> BoxedAppExtension for AppExtensionBox<Ext> {
    fn augment_args(&self, command: Command) -> Command {
        Ext::Args::augment_args_for_update(command)
    }

    fn arg_defaults(&self) -> Vec<(&'static str, String)> {
        self.ext.arg_defaults()
    }

    fn parse_args(&self, args: &ArgMatches) -> Result<Box<dyn BoxedAppExtensionArgs>> {
        let args = Ext::Args::from_arg_matches(args)?;
        Ok(Box::new(AppExtensionArgsBox {
            ext: self.ext.clone(),
            args,
        }))
    }
}

pub(crate) trait Downcast: Any {
    fn as_any(&self) -> &dyn Any;
}

impl<T: Any> Downcast for T {
    fn as_any(&self) -> &dyn Any {
        self
    }
}

// Internal trait to hide the concrete extension args type.
pub(crate) trait BoxedAppExtensionArgs: Downcast + Send + Sync + 'static {
    fn environment_hook(&self, env: &mut MononokeEnvironment) -> Result<()>;
    fn log_drain_hook(
        &self,
        drain: Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>,
    ) -> Result<Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>>;
}

// Box to store an app extension and its parsed args.
pub(crate) struct AppExtensionArgsBox<Ext: AppExtension> {
    ext: Arc<Ext>,
    args: Ext::Args,
}

impl<Ext: AppExtension> AppExtensionArgsBox<Ext> {
    pub(crate) fn args(&self) -> &Ext::Args {
        &self.args
    }
}

impl<Ext: AppExtension> BoxedAppExtensionArgs for AppExtensionArgsBox<Ext> {
    fn environment_hook(&self, env: &mut MononokeEnvironment) -> Result<()> {
        self.ext.environment_hook(&self.args, env)
    }

    fn log_drain_hook(
        &self,
        drain: Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>,
    ) -> Result<Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>> {
        self.ext.log_drain_hook(&self.args, drain)
    }
}


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

use anyhow::Error;
use anyhow::Result;
use clap::Args;
use fbinit::FacebookInit;
use services::Fb303Service;
use services::FbStatus;
use slog::info;
use slog::Logger;
use slog::Never;
use slog::SendSyncRefUnwindSafeDrain;

use crate::AppExtension;

/// Command line arguments that fb303 for service
#[derive(Args, Debug)]
pub struct Fb303Args {
    /// Port for fb303 service
    // thrift_port alias is for compatibility with mononoke server
    // TODO: switch mononoke server to use the same flags as everybody.
    #[clap(long, alias = "thrift_port", value_name = "PORT")]
    pub fb303_thrift_port: Option<i32>,
}

// Re-eport AliveService for convenience so callers do not have to get the services dependency to
// get AliveService.
pub use services::AliveService;

/// A FB303 service that reports healthy once set_ready has been called.
#[derive(Clone)]
pub struct ReadyFlagService {
    ready: Arc<AtomicBool>,
}

impl ReadyFlagService {
    pub fn new() -> Self {
        Self {
            ready: Arc::new(AtomicBool::new(false)),
        }
    }

    pub fn set_ready(&self) {
        self.ready.store(true, Ordering::Relaxed);
    }
}

impl Fb303Service for ReadyFlagService {
    fn getStatus(&self) -> FbStatus {
        if self.ready.load(Ordering::Relaxed) {
            FbStatus::Alive
        } else {
            FbStatus::Starting
        }
    }
}

impl Fb303Args {
    /// This is a lower-level function that requires you to spawn the stats aggregation future
    /// yourself. This is useful if you'd like to be able to drop it in order to cancel it.
    ///
    /// Usually starting the fb303 server and stats aggregation is done by functions like
    /// `MononokeApp::run_with_monitoring_and_logging`.
    pub fn start_fb303_server<S: Fb303Service + Sync + Send + 'static>(
        &self,
        fb: FacebookInit,
        service_name: &str,
        logger: &Logger,
        service: S,
    ) -> Result<Option<()>, Error> {
        let service_name = service_name.to_string();
        self.fb303_thrift_port
            .map(|port| {
                info!(logger, "Initializing fb303 thrift server on port {}", port);

                thread::Builder::new()
                    .name("fb303_thrift_service".to_owned())
                    .spawn(move || {
                        services::run_service_framework(
                            fb,
                            service_name,
                            port,
                            0, // Disables separate status http server
                            Box::new(service),
                        )
                        .expect("failure while running thrift service framework")
                    })
                    .map_err(Error::from)?;

                Ok(())
            })
            .transpose()
    }
}

pub struct Fb303AppExtension;

impl AppExtension for Fb303AppExtension {
    type Args = Fb303Args;

    /// Hook executed after creating the log drain allowing for augmenting the logging.
    fn log_drain_hook(
        &self,
        args: &Fb303Args,
        drain: Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>,
    ) -> Result<Arc<dyn SendSyncRefUnwindSafeDrain<Ok = (), Err = Never>>> {
        if args.fb303_thrift_port.is_some() {
            Ok(Arc::new(slog_stats::StatsDrain::new(drain)))
        } else {
            Ok(drain)
        }
    }
}


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

#![feature(let_chains)]

mod app;
pub mod args;
mod builder;
mod extension;
pub mod fb303;
mod repos_manager;

pub use app::MononokeApp;
pub use builder::MononokeAppBuilder;
pub use extension::AppExtension;
pub use repos_manager::MononokeReposManager;

#[doc(hidden)]
pub mod macro_export {
    pub use base_app;
}

/// Define subcommands for a Mononoke App
///
/// This macro is a convenience macro for defining a Mononoke App with
/// subcommands, where each subcommand gets its own module.
///
/// To use this:
///
/// * Create `commands.rs` with a call to this macro.  List each module
///   that defines a subcommand, which should be placed in a `commands`
///   directory.
///
/// ```text
/// use mononoke_app::subcommands;
///
/// subcommands! {
///    mod first_command;
///    mod second_command;
/// }
/// ```
///
/// * Each command should provide two things at the root of its module:
///
///   * A `CommandArgs` type which should derive `clap::Parser` and represents
///     the options available to this subcommand.
///
///   * A function to run the command, with the signature:
///     `pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<()>`.
///
/// * This macro gathers these together and provides a `subcommands` function
///   and a `dispatch` function.
///
///   * Call `commands::subcommands()` to construct the subcommands and pass
///     them to `MononokeAppBuilder::build_with_subcommands` for your app.
///
///   * When executing the app, call `commands::dispatch(app)` to run the
///     selected subcommand.
#[macro_export]
macro_rules! subcommands {
    ( $( mod $command:ident; )* ) => {
        $crate::macro_export::base_app::subcommands!{ type App = $crate::MononokeApp; $( mod $command; )* }
    }
}


/*
 * 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::Arc;
use std::time::Instant;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Result;
use async_trait::async_trait;
use facet::AsyncBuildable;
use futures::stream;
use futures::stream::StreamExt;
use itertools::Itertools;
use metaconfig_parser::RepoConfigs;
use metaconfig_parser::StorageConfigs;
use metaconfig_types::Redaction;
use metaconfig_types::RepoConfig;
use metaconfig_types::ShardedService;
use mononoke_api::Mononoke;
use mononoke_configs::ConfigUpdateReceiver;
use mononoke_configs::MononokeConfigs;
use mononoke_repos::MononokeRepos;
use repo_factory::RepoFactory;
use repo_factory::RepoFactoryBuilder;
use slog::info;
use slog::o;
use slog::Logger;
use stats::prelude::*;

define_stats! {
    prefix = "mononoke.app";
    initialization_time_millisecs: dynamic_timeseries(
        "initialization_time_millisecs.{}",
        (reponame: String);
        Average, Sum, Count
    ),
    completion_duration_secs: timeseries(Average, Sum, Count),
}

/// A manager of a MononokeRepos collection.
///
/// This allows repos to be added or removed from the MononokeRepos
/// collection.
pub struct MononokeReposManager<Repo> {
    repos: Arc<MononokeRepos<Repo>>,
    configs: Arc<MononokeConfigs>,
    repo_factory: Arc<RepoFactory>,
    logger: Logger,
    redaction_disabled: bool,
}

impl<Repo> MononokeReposManager<Repo> {
    // Create a new `MononokeReposManager`.
    // Unlike `new_with_redaction_disabled`, we don't expose the mechanism to access redacted blobs
    // through this API.
    // This should be your goto constructor for this struct except if you have a specific reason
    // for needing to disable redaction.
    #[allow(unused)]
    pub(crate) async fn new<Names>(
        configs: Arc<MononokeConfigs>,
        repo_factory: Arc<RepoFactory>,
        logger: Logger,
        service_name: Option<ShardedService>,
        repo_names: Names,
    ) -> Result<Self>
    where
        Names: IntoIterator<Item = String>,
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        Self::new_with_redaction_disabled(
            configs,
            repo_factory,
            logger,
            service_name,
            repo_names,
            false,
        )
        .await
    }

    pub(crate) async fn new_with_redaction_disabled<Names>(
        configs: Arc<MononokeConfigs>,
        repo_factory: Arc<RepoFactory>,
        logger: Logger,
        service_name: Option<ShardedService>,
        repo_names: Names,
        redaction_disabled: bool,
    ) -> Result<Self>
    where
        Names: IntoIterator<Item = String>,
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>
            + Send
            + Sync
            + 'static,
    {
        let repos = Arc::new(MononokeRepos::new());
        let mgr = MononokeReposManager {
            repos,
            configs,
            repo_factory,
            logger,
            redaction_disabled,
        };
        mgr.populate_repos(repo_names).await?;
        let update_receiver = MononokeConfigUpdateReceiver::new(
            mgr.repos.clone(),
            mgr.repo_factory.clone(),
            mgr.logger.clone(),
            service_name,
        );
        mgr.configs
            .register_for_update(Arc::new(update_receiver) as Arc<dyn ConfigUpdateReceiver>);
        Ok(mgr)
    }

    /// The repo collection that is being managed.
    pub fn repos(&self) -> &Arc<MononokeRepos<Repo>> {
        &self.repos
    }

    /// The logger for the app.
    pub fn logger(&self) -> &Logger {
        &self.logger
    }

    /// Construct a logger for a specific repo.
    pub fn repo_logger(&self, repo_name: &str) -> Logger {
        self.logger.new(o!("repo" => repo_name.to_string()))
    }

    /// Return a repo config for a named repo.  This reads from the main
    /// configuration, so doesn't need to be a currently managed repo.
    pub fn repo_config(&self, repo_name: &str) -> Result<RepoConfig> {
        let mut repo_config = self
            .configs
            .repo_configs()
            .repos
            .get(repo_name)
            .cloned()
            .ok_or_else(|| anyhow!("unknown reponame: {:?}", repo_name))?;
        if self.redaction_disabled {
            repo_config.redaction = Redaction::Disabled;
        }
        Ok(repo_config)
    }

    /// Construct and add a new repo to the managed repo collection.
    pub async fn add_repo(&self, repo_name: &str) -> Result<Arc<Repo>>
    where
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>,
    {
        let repo_config = self.repo_config(repo_name)?;
        let repo_id = repo_config.repoid.id();
        let common_config = self.configs.repo_configs().common.clone();
        let repo = self
            .repo_factory
            .build(repo_name.to_string(), repo_config, common_config)
            .await?;
        self.repos.add(repo_name, repo_id, repo);
        self.repos
            .get_by_name(repo_name)
            .ok_or_else(|| anyhow!("Couldn't retrive added repo {}", repo_name))
    }

    /// Remove a repo from the managed repo collection.
    pub fn remove_repo(&self, repo_name: &str) {
        self.repos.remove(repo_name);
    }

    async fn populate_repos<Names>(&self, repo_names: Names) -> Result<()>
    where
        Names: IntoIterator<Item = String>,
        Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>>,
    {
        let repos_input = stream::iter(repo_names.into_iter().unique())
            .map(|repo_name| {
                let repo_factory = self.repo_factory.clone();
                let name = repo_name.clone();
                async move {
                    let start = Instant::now();
                    let logger = self.logger();
                    let repo_config = self.repo_config(&repo_name)?;
                    let common_config = self.configs.repo_configs().common.clone();
                    let repo_id = repo_config.repoid.id();
                    info!(logger, "Initializing repo: {}", &repo_name);
                    let repo = repo_factory
                        .build(name, repo_config, common_config)
                        .await
                        .with_context(|| format!("Failed to initialize repo '{}'", &repo_name))?;
                    info!(logger, "Initialized repo: {}", &repo_name);
                    STATS::initialization_time_millisecs.add_value(
                        start.elapsed().as_millis().try_into().unwrap_or(i64::MAX),
                        (repo_name.to_string(),),
                    );
                    anyhow::Ok((repo_id, repo_name, repo))
                }
            })
            // Repo construction can be heavy, 30 at a time is sufficient.
            .buffered(30)
            .collect::<Vec<_>>();
        // There are lots of deep FuturesUnordered here that have caused inefficient polling with
        // Tokio coop in the past.
        let repos_input = tokio::task::unconstrained(repos_input)
            .await
            .into_iter()
            .collect::<Result<Vec<_>>>()?;
        self.repos.populate(repos_input);
        Ok(())
    }
}

// This has a concrete type until we make `mononoke_api::Mononoke` generic
// also.
impl MononokeReposManager<mononoke_api::Repo> {
    pub fn make_mononoke_api(&self) -> Result<Mononoke> {
        let repo_names_in_tier =
            Vec::from_iter(self.configs.repo_configs().repos.iter().filter_map(
                |(name, config)| {
                    if config.enabled {
                        Some(name.to_string())
                    } else {
                        None
                    }
                },
            ));
        Mononoke::new(self.repos.clone(), repo_names_in_tier)
    }
}

/// Struct responsible for receiving updated configurations from MononokeConfigs
/// and refreshing repos (and related entities) based on the update.
pub struct MononokeConfigUpdateReceiver<Repo> {
    repos: Arc<MononokeRepos<Repo>>,
    repo_factory: Arc<RepoFactory>,
    logger: Logger,
    service_name: Option<ShardedService>,
}

impl<Repo> MononokeConfigUpdateReceiver<Repo> {
    fn new(
        repos: Arc<MononokeRepos<Repo>>,
        repo_factory: Arc<RepoFactory>,
        logger: Logger,
        service_name: Option<ShardedService>,
    ) -> Self {
        Self {
            repos,
            repo_factory,
            logger,
            service_name,
        }
    }
}

#[async_trait]
impl<Repo> ConfigUpdateReceiver for MononokeConfigUpdateReceiver<Repo>
where
    Repo: for<'builder> AsyncBuildable<'builder, RepoFactoryBuilder<'builder>> + Send + Sync,
{
    async fn apply_update(
        &self,
        repo_configs: Arc<RepoConfigs>,
        _: Arc<StorageConfigs>,
    ) -> Result<()> {
        let mut repos_to_load = Vec::new();
        for (repo_name, repo_config) in repo_configs.repos.clone().into_iter() {
            if self.repos.get_by_name(repo_name.as_str()).is_some() {
                // Repo was already present on the server. Need to reload it.
                repos_to_load.push((repo_name, repo_config))
            }
            // If the service name is known, then by default we need to reload or add all repos
            // that are in RepoConfig AND are shallow-sharded (i.e. NOT deep-sharded).
            else if repo_config.enabled && let Some(ref service_name) = self.service_name {
                if let Some(ref config) = repo_config.deep_sharding_config {
                    // Repo is shallow sharded for this service AND enabled, so should be loaded.
                    if !config.status.get(service_name).cloned().unwrap_or(false)
                    {
                        repos_to_load.push((repo_name, repo_config));
                    }
                } else {
                    // Service specific sharding config doesn't exist for repo but the repo is
                    // enabled so should be considered as shallow-sharded.
                    repos_to_load.push((repo_name, repo_config));
                }
            }
            // The repos present on the server but not part of RepoConfigs are ignored by
            // default. This situation can happen when the name of the repo changes
            // (e.g. whatsapp/server.mirror renamed to whatsapp/server) or when a repo is
            // added or removed. In such a case, reloading of the repo with the old name
            // would not be possible based on the new configs.
        }

        let repos_input = stream::iter(repos_to_load)
            .map(|(repo_name, repo_config)| {
                let repo_factory = self.repo_factory.clone();
                let name = repo_name.clone();
                let logger = self.logger.clone();
                let common_config = repo_configs.common.clone();
                async move {
                    let repo_id = repo_config.repoid.id();
                    info!(logger, "Reloading repo: {}", &repo_name);
                    let repo = repo_factory
                        .build(name, repo_config, common_config)
                        .await
                        .with_context(|| format!("Failed to reload repo '{}'", &repo_name))?;
                    info!(logger, "Reloaded repo: {}", &repo_name);

                    anyhow::Ok((repo_id, repo_name, repo))
                }
            })
            // Repo construction can be heavy, 30 at a time is sufficient.
            .buffered(30)
            .collect::<Vec<_>>();
        // There are lots of deep FuturesUnordered here that have caused inefficient polling with
        // Tokio coop in the past.
        let repos_input = tokio::task::unconstrained(repos_input)
            .await
            .into_iter()
            .collect::<Result<Vec<_>>>()?;
        self.repos.populate(repos_input);
        Ok(())
    }
}


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

use std::path::PathBuf;

use clap::Args;

/// Command line arguments for controlling Acls
#[derive(Args, Debug)]
pub struct AclArgs {
    /// Load ACLs from a JSON-formatted file.
    #[clap(long, value_parser)]
    pub acl_file: Option<PathBuf>,
}


/*
 * 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::Context;
use anyhow::Result;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use bookmarks::BookmarkKey;
use bookmarks::BookmarksRef;
use clap::ArgGroup;
use clap::Args;
use context::CoreContext;
use mercurial_types::HgChangesetId;
use mononoke_types::ChangesetId;

/// Command line arguments for specifying a changeset.
#[derive(Args, Debug)]
#[clap(group(
    ArgGroup::new("changeset")
        .required(true)
        .args(&["changeset_id", "hg_id", "bookmark"]),
))]
pub struct ChangesetArgs {
    /// Bonsai changeset id
    #[clap(long, short = 'i')]
    changeset_id: Option<ChangesetId>,

    /// Hg changeset id
    #[clap(long)]
    hg_id: Option<HgChangesetId>,

    /// Bookmark name
    #[clap(long, short = 'B')]
    bookmark: Option<BookmarkKey>,
}

impl ChangesetArgs {
    pub async fn resolve_changeset(
        &self,
        ctx: &CoreContext,
        repo: &(impl BookmarksRef + BonsaiHgMappingRef),
    ) -> Result<Option<ChangesetId>> {
        if let Some(bookmark) = &self.bookmark {
            repo.bookmarks()
                .get(ctx.clone(), bookmark)
                .await
                .with_context(|| format!("Failed to resolve bookmark '{}'", bookmark))
        } else if let Some(hg_id) = self.hg_id {
            repo.bonsai_hg_mapping()
                .get_bonsai_from_hg(ctx, hg_id)
                .await
                .with_context(|| format!("Failed to resolve hg changeset id {}", hg_id))
        } else {
            Ok(self.changeset_id)
        }
    }
}


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

use std::path::PathBuf;

use clap::ArgGroup;
use clap::Args;

#[derive(Args, Debug)]
#[clap(group(ArgGroup::new("config").args(&["config_path", "config_tier", "prod"]).required(true)))]
pub struct ConfigArgs {
    /// Path to Mononoke config
    #[clap(long, alias = "mononoke-config-path")]
    pub config_path: Option<String>,

    /// Use configerator-based configuration for a specific tier
    #[clap(long)]
    pub config_tier: Option<String>,

    /// Use configerator-based configuration for production
    #[clap(long)]
    pub prod: bool,

    /// Local path to fetch configerator configs from
    #[clap(long)]
    pub local_configerator_path: Option<PathBuf>,

    /// Regex for a Configerator path that must be covered by
    /// Mononoke's crypto project
    #[clap(long)]
    pub crypto_path_regex: Option<Vec<String>>,
}

const PRODUCTION_PREFIX: &str = "configerator://scm/mononoke/repos/tiers/";

fn configerator_config_path(tier: &str) -> String {
    format!("{}{}", PRODUCTION_PREFIX, tier)
}

impl ConfigArgs {
    pub fn config_path(&self) -> String {
        if let Some(config_path) = &self.config_path {
            config_path.clone()
        } else if self.prod {
            configerator_config_path("prod")
        } else if let Some(tier) = &self.config_tier {
            configerator_config_path(tier)
        } else {
            String::new()
        }
    }

    pub fn mode(&self) -> ConfigMode {
        if let Some(config_path) = &self.config_path {
            // Any configuration that matches the production prefix is prod.
            if config_path.starts_with(PRODUCTION_PREFIX) {
                return ConfigMode::Production;
            }
        } else {
            // Otherwise, we are prod if a prod tier is requested.
            if self.prod || self.config_tier.is_some() {
                return ConfigMode::Production;
            }
        }
        ConfigMode::Development
    }
}

#[derive(Copy, Clone, Debug, Eq, PartialEq)]
pub enum ConfigMode {
    Production,
    Development,
}


/*
 * 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 anyhow::Context;
use anyhow::Result;
use clap::Args;
use environment::MononokeEnvironment;
use slog::warn;

use crate::AppExtension;

/// Command line arguments for tweaking hooks
#[derive(Args, Debug)]
pub struct HooksArgs {
    /// Disable a hook. Pass this argument multiple times to disable multiple hooks.
    #[clap(long = "disable-hook")]
    disabled_hooks: Vec<String>,
}

pub struct HooksAppExtension;

impl AppExtension for HooksAppExtension {
    type Args = HooksArgs;

    fn environment_hook(&self, args: &Self::Args, env: &mut MononokeEnvironment) -> Result<()> {
        let mut res = HashMap::new();
        for repohook in &args.disabled_hooks {
            let repohook: Vec<_> = repohook.splitn(2, ':').collect();
            #[allow(clippy::get_first)]
            let repo = repohook.get(0);
            let hook = repohook.get(1);

            let (repo, hook) = repo
                .and_then(|repo| hook.map(|hook| (repo, hook)))
                .context("invalid format of disabled hook, should be 'REPONAME:HOOKNAME'")?;
            res.entry(repo.to_string())
                .or_insert_with(HashSet::new)
                .insert(hook.to_string());
        }

        if !res.is_empty() {
            warn!(env.logger, "The following Hooks were disabled: {:?}", res);
        }

        env.disabled_hooks = res;

        Ok(())
    }
}


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

use anyhow::Result;
use clap::Args;
use environment::MononokeEnvironment;

use crate::AppExtension;

/// Command line arguments that affect mcrouter usage
#[derive(Args, Debug)]
pub struct McrouterArgs {
    /// Use local McRouter for rate limits
    #[clap(long)]
    pub enable_mcrouter: bool,
}

pub struct McrouterAppExtension;

impl AppExtension for McrouterAppExtension {
    type Args = McrouterArgs;

    fn environment_hook(&self, args: &McrouterArgs, env: &mut MononokeEnvironment) -> Result<()> {
        if !args.enable_mcrouter {
            return Ok(());
        }

        #[cfg(fbcode_build)]
        {
            ::ratelim::use_proxy_if_available(env.fb);
            Ok(())
        }

        #[cfg(not(fbcode_build))]
        {
            let _ = env;
            unimplemented!("Passed --enable-mcrouter but it is supported only for fbcode builds",);
        }
    }
}


/*
 * 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::format_err;
use anyhow::Result;

mod acl;
mod changeset;
mod config;
mod hooks;
mod mcrouter;
mod mysql;
mod readonly;
mod repo;
mod repo_blobstore;
mod repo_filter;
mod runtime;
mod tunables;
mod wbc;

pub use acl::AclArgs;
pub use changeset::ChangesetArgs;
pub use config::ConfigArgs;
pub use config::ConfigMode;
pub use hooks::HooksAppExtension;
pub use mcrouter::McrouterAppExtension;
pub use mcrouter::McrouterArgs;
pub use mysql::MysqlArgs;
pub use readonly::ReadonlyArgs;
pub use repo::AsRepoArg;
pub use repo::MultiRepoArgs;
pub use repo::OptRepoArgs;
pub use repo::RepoArg;
pub use repo::RepoArgs;
pub use repo::SourceAndTargetRepoArgs;
pub use repo_blobstore::RepoBlobstoreArgs;
pub use repo_filter::RepoFilterAppExtension;
pub use runtime::RuntimeArgs;
pub use shutdown_timeout::ShutdownTimeoutArgs;
pub use tls::TLSArgs;
pub use wbc::WarmBookmarksCacheExtension;

pub use self::tunables::TunablesArgs;
pub use crate::fb303::Fb303Args;
pub use crate::repo_args;
pub use crate::repo_args_optional;

/// NOTE: Don't use this. "configerator:" prefix don't need to exist and is going to be removed.
/// Pass raw path instead.
pub fn parse_config_spec_to_path(source_spec: &str) -> Result<String> {
    // NOTE: This means we don't support file paths with ":" in them, but it also means we can
    // add other options after the first ":" later if we want.
    let mut iter = source_spec.split(':');

    // NOTE: We match None as the last element to make sure the input doesn't contain
    // disallowed trailing parts.
    match (iter.next(), iter.next(), iter.next()) {
        (Some("configerator"), Some(path), None) => Ok(path.to_string()),
        (Some(path), None, None) => Ok(path.to_string()),
        _ => Err(format_err!("Invalid configuration spec: {:?}", source_spec)),
    }
}