/*
 * 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::anyhow;
use anyhow::Result;
use clap::Parser;
use fbinit::FacebookInit;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;

mod commands;

/// Tools for Mononoke Tests
#[derive(Parser)]
struct TestToolArgs {}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<()> {
    let subcommands = commands::subcommands();
    let app = MononokeAppBuilder::new(fb).build_with_subcommands::<TestToolArgs>(subcommands)?;
    app.run_basic(async_main)
}

async fn async_main(app: MononokeApp) -> Result<()> {
    if app.is_production() {
        return Err(anyhow!(
            "mononoke-testtool cannot be run against production"
        ));
    }
    commands::dispatch(app).await
}


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

use std::collections::BTreeMap;
use std::fs::File;
use std::io::Read;
use std::path::PathBuf;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use blobrepo::BlobRepo;
use changesets_creation::save_changesets;
use clap::Parser;
use mercurial_derivation::MappedHgChangesetId;
use mononoke_app::args::RepoArgs;
use mononoke_app::MononokeApp;
use mononoke_types::hash::GitSha1;
use mononoke_types::BonsaiChangeset;
use mononoke_types::BonsaiChangesetMut;
use mononoke_types::ChangesetId;
use mononoke_types::DateTime;
use mononoke_types::FileChange;
use mononoke_types::MPath;
use repo_blobstore::RepoBlobstoreRef;
use repo_derived_data::RepoDerivedDataRef;
use serde_derive::Deserialize;
use sorted_vector_map::SortedVectorMap;

/// Create commits from a JSON-encoded bonsai changeset
///
/// The bonsai changeset is intentionally not checked for correctness, as this
/// may be used in tests to test handling of malformed bonsai changesets.
#[derive(Parser)]
pub struct CommandArgs {
    #[clap(flatten)]
    repo_args: RepoArgs,

    /// Path to a file containing a JSON-encoded bonsai changeset
    bonsai_file: PathBuf,
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<()> {
    let ctx = app.new_basic_context();

    let mut content = String::new();
    File::open(&args.bonsai_file)
        .with_context(|| {
            format!(
                "Failed to open bonsai changeset file '{}'",
                args.bonsai_file.to_string_lossy()
            )
        })?
        .read_to_string(&mut content)
        .context("Failed to read bonsai changeset file")?;

    let bcs: BonsaiChangeset = serde_json::from_str::<DeserializableBonsaiChangeset>(&content)
        .context("Failed to parse bonsai changeset data")?
        .into_bonsai()?
        .freeze()?;

    let repo: BlobRepo = app
        .open_repo(&args.repo_args)
        .await
        .context("Failed to open repo")?;

    for (_, change) in bcs.simplified_file_changes() {
        match change {
            Some(tc) => {
                if filestore::get_metadata(repo.repo_blobstore(), &ctx, &tc.content_id().into())
                    .await?
                    .is_none()
                {
                    return Err(anyhow!(
                        "file content {} is not found in the filestore",
                        &tc.content_id()
                    ));
                }
            }
            None => {}
        }
    }
    let bcs_id = bcs.get_changeset_id();
    save_changesets(&ctx, &repo, vec![bcs])
        .await
        .context("Failed to save changeset")?;
    let hg_cs = repo
        .repo_derived_data()
        .derive::<MappedHgChangesetId>(&ctx, bcs_id)
        .await
        .context("Failed to derive Mercurial changeset")?
        .hg_changeset_id();
    println!(
        "Created bonsai changeset {} for Hg changeset {}",
        bcs_id, hg_cs
    );
    Ok(())
}

#[derive(Debug, Deserialize)]
pub struct DeserializableBonsaiChangeset {
    pub parents: Vec<ChangesetId>,
    pub author: String,
    pub author_date: DateTime,
    pub committer: Option<String>,
    pub committer_date: Option<DateTime>,
    pub message: String,
    pub hg_extra: BTreeMap<String, Vec<u8>>,
    pub git_extra_headers: Option<BTreeMap<Vec<u8>, Vec<u8>>>,
    pub git_tree_hash: Option<GitSha1>,
    pub file_changes: BTreeMap<String, FileChange>,
}

impl DeserializableBonsaiChangeset {
    pub fn into_bonsai(self) -> Result<BonsaiChangesetMut, Error> {
        let files = self
            .file_changes
            .into_iter()
            .map::<Result<_, Error>, _>(|(path, changes)| {
                Ok((MPath::new(path.as_bytes())?, changes))
            })
            .collect::<Result<SortedVectorMap<_, _>, _>>()?;
        Ok(BonsaiChangesetMut {
            parents: self.parents,
            author: self.author,
            author_date: self.author_date,
            committer: self.committer,
            committer_date: self.committer_date,
            message: self.message,
            hg_extra: self.hg_extra.into(),
            git_extra_headers: self.git_extra_headers.map(|extra| {
                extra
                    .into_iter()
                    .map(|(k, v)| (smallvec::SmallVec::from(k), bytes::Bytes::from(v)))
                    .collect()
            }),
            file_changes: files,
            is_snapshot: false,
            git_annotated_tag: None,
            git_tree_hash: self.git_tree_hash,
        })
    }
}


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

//! DrawDAG for Integration Tests
//!
//! A DrawDAG specification consists of an ASCII graph (either left-to-right
//! or bottom-to-top), and a series of comments that define additional
//! properties for each commit.
//!
//! Valid properties are:
//!
//! * Set a known changeset id for an already-existing commit
//!     # exists: COMMIT id
//!
//! * Set a bookmark on a commit
//!     # bookmark: COMMIT name
//!
//! * Set the content of a file.
//!     # modify: COMMIT path/to/file "content"
//!
//! * Mark a file as deleted.
//!     # delete: COMMIT path/to/file
//!
//! * Forget file that was about to be added (useful for getting rid of files
//!   that are added by default):
//!     # forget: COMMIT path/to/file
//!
//! Paths can be surrounded by quotes if they contain special characters.

use std::collections::BTreeMap;
use std::collections::HashMap;
use std::fmt::Display;
use std::io::Write;

use anyhow::anyhow;
use anyhow::Context;
use anyhow::Error;
use anyhow::Result;
use basename_suffix_skeleton_manifest::RootBasenameSuffixSkeletonManifest;
use blame::RootBlameV2;
use blobrepo::BlobRepo;
use bonsai_hg_mapping::BonsaiHgMappingRef;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkUpdateReason;
use bookmarks::BookmarksRef;
use changeset_info::ChangesetInfo;
use clap::Parser;
use context::CoreContext;
use deleted_manifest::RootDeletedManifestV2Id;
use derived_data_manager::BatchDeriveOptions;
use derived_data_manager::BonsaiDerivable;
use fastlog::RootFastlog;
use filenodes_derivation::FilenodesOnlyPublic;
use fsnodes::RootFsnodeId;
use futures::try_join;
use mercurial_derivation::MappedHgChangesetId;
use mononoke_app::args::RepoArgs;
use mononoke_app::MononokeApp;
use mononoke_types::ChangesetId;
use mononoke_types::DateTime;
use repo_derived_data::RepoDerivedDataRef;
use skeleton_manifest::RootSkeletonManifestId;
use tests_utils::drawdag::extend_from_dag_with_changes;
use tests_utils::drawdag::ChangeFn;
use tests_utils::CommitIdentifier;
use tests_utils::CreateCommitContext;
use tokio::io::AsyncReadExt;
use topo_sort::sort_topological;
use unodes::RootUnodeManifestId;

/// Create commits from a drawn DAG.
#[derive(Parser)]
pub struct CommandArgs {
    #[clap(flatten)]
    repo_args: RepoArgs,

    /// Disable creation of default files in each commit
    #[clap(long)]
    no_default_files: bool,

    /// Derive all derived data types for all commits
    #[clap(long)]
    derive_all: bool,

    /// Print hashes in HG format instead of bonsai
    #[clap(long)]
    print_hg_hashes: bool,
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum Action {
    Exists { name: String, id: ChangesetId },
    Bookmark { name: String, bookmark: BookmarkKey },
    Change { name: String, change: ChangeAction },
}

#[derive(Clone, Debug, Eq, PartialEq)]
enum ChangeAction {
    Modify {
        path: Vec<u8>,
        content: Vec<u8>,
    },
    Delete {
        path: Vec<u8>,
    },
    Forget {
        path: Vec<u8>,
    },
    Extra {
        key: String,
        value: Vec<u8>,
    },
    Message {
        message: String,
    },
    Author {
        author: String,
    },
    AuthorDate {
        author_date: DateTime,
    },
    Committer {
        committer: String,
    },
    CommitterDate {
        committer_date: DateTime,
    },
    Copy {
        path: Vec<u8>,
        content: Vec<u8>,
        parent: String,
        parent_path: Vec<u8>,
    },
}

impl Action {
    fn new(spec: &str) -> Result<Self> {
        if let Some((key, args)) = spec.trim().split_once(':') {
            let args = ActionArg::parse_args(args)
                .with_context(|| format!("Failed to parse args for '{}'", key))?;
            match (key, args.as_slice()) {
                ("exists", [name, id]) => {
                    let name = name.to_string()?;
                    let id = id.to_string()?.parse()?;
                    Ok(Action::Exists { name, id })
                }
                ("bookmark", [name, bookmark]) => {
                    let name = name.to_string()?;
                    let bookmark = bookmark.to_string()?.parse()?;
                    Ok(Action::Bookmark { name, bookmark })
                }
                ("message", [name, message]) => {
                    let name = name.to_string()?;
                    let message = message.to_string()?;
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Message { message },
                    })
                }
                ("author", [name, author]) => {
                    let name = name.to_string()?;
                    let author = author.to_string()?;
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Author { author },
                    })
                }
                ("author_date", [name, author_date]) => {
                    let name = name.to_string()?;
                    let author_date = DateTime::from_rfc3339(&author_date.to_string()?)?;
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::AuthorDate { author_date },
                    })
                }
                ("committer", [name, committer]) => {
                    let name = name.to_string()?;
                    let committer = committer.to_string()?;
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Committer { committer },
                    })
                }
                ("committer_date", [name, committer_date]) => {
                    let name = name.to_string()?;
                    let committer_date = DateTime::from_rfc3339(&committer_date.to_string()?)?;
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::CommitterDate { committer_date },
                    })
                }
                ("modify", [name, path, content]) => {
                    let name = name.to_string()?;
                    let path = path.to_bytes();
                    let content = content.to_bytes();
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Modify { path, content },
                    })
                }
                ("delete", [name, path]) => {
                    let name = name.to_string()?;
                    let path = path.to_bytes();
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Delete { path },
                    })
                }
                ("forget", [name, path]) => {
                    let name = name.to_string()?;
                    let path = path.to_bytes();
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Forget { path },
                    })
                }
                ("extra", [name, key, value]) => {
                    let name = name.to_string()?;
                    let key = key.to_string()?;
                    let value = value.to_bytes();
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Extra { key, value },
                    })
                }
                ("copy", [name, path, content, parent, parent_path]) => {
                    let name = name.to_string()?;
                    let path = path.to_bytes();
                    let content = content.to_bytes();
                    let parent = parent.to_string()?;
                    let parent_path = parent_path.to_bytes();
                    Ok(Action::Change {
                        name,
                        change: ChangeAction::Copy {
                            path,
                            content,
                            parent,
                            parent_path,
                        },
                    })
                }
                _ => Err(anyhow!("Invalid spec for key: {}", key)),
            }
        } else {
            Err(anyhow!("Invalid spec: {}", spec))
        }
    }
}

struct ActionArg(Vec<u8>);

impl ActionArg {
    fn new() -> Self {
        ActionArg(Vec::new())
    }

    fn to_bytes(&self) -> Vec<u8> {
        self.0.clone()
    }

    fn to_string(&self) -> Result<String> {
        let s = std::str::from_utf8(&self.0)
            .context("Expected UTF-8 string for drawdag action argument")?;
        Ok(s.to_string())
    }

    fn is_empty(&self) -> bool {
        self.0.is_empty()
    }

    fn push(&mut self, ch: char) {
        let mut buf = [0; 4];
        self.0
            .extend_from_slice(ch.encode_utf8(&mut buf).as_bytes());
    }

    fn push_byte(&mut self, byte: u8) {
        self.0.push(byte)
    }

    fn push_hex(&mut self, mut iter: impl Iterator<Item = char>) -> Result<()> {
        if let (Some(top_hex), Some(bottom_hex)) = (iter.next(), iter.next()) {
            if let (Some(top_digit), Some(bottom_digit)) =
                (top_hex.to_digit(16), bottom_hex.to_digit(16))
            {
                self.push_byte((top_digit * 0x10 + bottom_digit) as u8);
                return Ok(());
            }
        }
        Err(anyhow!("Expected two hex digits"))
    }

    fn parse_args(args: &str) -> Result<Vec<Self>> {
        let mut iter = args.trim().chars();
        let mut args = Vec::new();
        let mut arg = ActionArg::new();
        let mut in_quotes = false;
        while let Some(ch) = iter.next() {
            if in_quotes {
                match ch {
                    '"' => in_quotes = false,
                    '\\' => match iter
                        .next()
                        .ok_or_else(|| anyhow!("Unexpected end-of-line after '\\'"))?
                    {
                        '\\' => arg.push('\\'),
                        'r' => arg.push('\r'),
                        'n' => arg.push('\n'),
                        't' => arg.push('\t'),
                        'f' => arg.push('\u{0C}'),
                        'b' => arg.push('\u{08}'),
                        '"' => arg.push('"'),
                        'x' => arg.push_hex(&mut iter)?,
                        esc => return Err(anyhow!("Unexpected escape sequence: '\\{}'", esc)),
                    },
                    ch => arg.push(ch),
                }
            } else {
                match ch {
                    '"' => in_quotes = true,
                    ch if ch.is_whitespace() => {
                        if !arg.is_empty() {
                            args.push(arg);
                            arg = ActionArg::new();
                        }
                    }
                    ch if ch.is_alphanumeric() || "_-./".contains(ch) => {
                        arg.push(ch);
                    }
                    ch => return Err(anyhow!("Unexpected character: '{}'", ch)),
                }
            }
        }
        if in_quotes {
            return Err(anyhow!("Unterminated string literal"));
        }
        if !arg.is_empty() {
            args.push(arg);
        }
        Ok(args)
    }
}

fn print_name_hash_pairs(pairs: impl IntoIterator<Item = (String, impl Display)>) -> Result<()> {
    for (name, id) in pairs.into_iter() {
        writeln!(std::io::stdout(), "{}={}", name, id)?;
    }
    Ok(())
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<()> {
    let ctx = app.new_basic_context();

    let repo: BlobRepo = app
        .open_repo(&args.repo_args)
        .await
        .context("Failed to open repo")?;

    // Read DAG from stdin
    let mut input = String::new();
    tokio::io::stdin().read_to_string(&mut input).await?;

    let mut dag_buffer = String::new();
    let mut actions = Vec::new();
    for line in input.lines() {
        if let Some((dag_line, comment)) = line.split_once('#') {
            dag_buffer.push_str(dag_line);
            dag_buffer.push('\n');
            actions.push(Action::new(comment)?);
        } else {
            dag_buffer.push_str(line);
            dag_buffer.push('\n');
        }
    }

    let mut existing: BTreeMap<String, ChangesetId> = BTreeMap::new();
    let mut commit_changes: BTreeMap<String, Vec<ChangeAction>> = BTreeMap::new();
    let mut bookmarks: BTreeMap<BookmarkKey, String> = BTreeMap::new();

    for action in actions {
        match action {
            Action::Exists { name, id } => {
                existing.insert(name, id);
            }
            Action::Bookmark { name, bookmark } => {
                bookmarks.insert(bookmark, name);
            }
            Action::Change { name, change } => {
                commit_changes
                    .entry(name)
                    .or_insert_with(Vec::new)
                    .push(change);
            }
        }
    }

    let mut change_fns = BTreeMap::new();
    for (name, changes) in commit_changes {
        let apply: Box<ChangeFn<BlobRepo>> = Box::new(
            move |c: CreateCommitContext<BlobRepo>,
                  committed: &'_ BTreeMap<String, ChangesetId>| {
                apply_changes(c, committed, changes)
            },
        );
        change_fns.insert(name, apply);
    }

    let (commits, dag) = extend_from_dag_with_changes(
        &ctx,
        &repo,
        &dag_buffer,
        change_fns,
        existing,
        !args.no_default_files,
    )
    .await?;

    if !bookmarks.is_empty() {
        let mut txn = repo.bookmarks().create_transaction(ctx.clone());
        for (bookmark, name) in bookmarks {
            let target = commits
                .get(&name)
                .ok_or_else(|| anyhow!("No commit {} for bookmark {}", name, bookmark))?;
            let old_value = repo
                .bookmarks()
                .get(ctx.clone(), &bookmark)
                .await
                .with_context(|| format!("Failed to resolve bookmark '{}'", bookmark))?;
            // It's better to update/create rather than force_set which doesn't
            // save the old cid to the bookmark update log. (So it looks like
            // creation but it's update)
            match old_value {
                Some(old_value) => txn.update(
                    &bookmark,
                    *target,
                    old_value,
                    BookmarkUpdateReason::TestMove,
                ),
                None => txn.create(&bookmark, *target, BookmarkUpdateReason::TestMove),
            }?;
        }
        txn.commit().await?;
    }

    let any_derivation_needed = args.derive_all | args.print_hg_hashes;
    if any_derivation_needed {
        let dag: HashMap<_, _> = dag
            .into_iter()
            .map(|(k, v)| (k, v.into_iter().collect()))
            .collect();
        let sorted = sort_topological(&dag).ok_or_else(|| anyhow!("Graph has a cycle"))?;
        let csids = sorted
            .into_iter()
            .map(|name| {
                commits
                    .get(&name)
                    .cloned()
                    .ok_or_else(|| anyhow!("No commit found for {}", name))
            })
            .collect::<Result<Vec<_>>>()?;

        if args.derive_all {
            derive_all(&ctx, &repo, &csids).await?;
        } else {
            derive::<MappedHgChangesetId>(&ctx, &repo, &csids).await?;
        }
    }

    if args.print_hg_hashes {
        let mapping: HashMap<_, _> = repo
            .bonsai_hg_mapping()
            .get(&ctx, commits.values().copied().collect::<Vec<_>>().into())
            .await?
            .into_iter()
            .map(|entry| (entry.bcs_id, entry.hg_cs_id))
            .collect();
        let commits = commits
            .into_iter()
            .map(|(name, id)| {
                mapping
                    .get(&id)
                    .ok_or_else(|| anyhow!("Couldn't translate {}={} to hg", name, id))
                    .map(|hg_id| (name, hg_id))
            })
            .collect::<Result<Vec<_>>>()?;
        print_name_hash_pairs(commits)?;
    } else {
        print_name_hash_pairs(commits)?;
    }

    Ok(())
}

fn apply_changes<'a>(
    mut c: CreateCommitContext<'a, BlobRepo>,
    committed: &'_ BTreeMap<String, ChangesetId>,
    changes: Vec<ChangeAction>,
) -> CreateCommitContext<'a, BlobRepo> {
    for change in changes {
        match change {
            ChangeAction::Modify { path, content, .. } => c = c.add_file(path.as_slice(), content),
            ChangeAction::Delete { path, .. } => c = c.delete_file(path.as_slice()),
            ChangeAction::Forget { path, .. } => c = c.forget_file(path.as_slice()),
            ChangeAction::Extra { key, value, .. } => c = c.add_extra(key, value),
            ChangeAction::Message { message } => c = c.set_message(message),
            ChangeAction::Author { author } => c = c.set_author(author),
            ChangeAction::AuthorDate { author_date } => c = c.set_author_date(author_date),
            ChangeAction::Committer { committer } => c = c.set_committer(committer),
            ChangeAction::CommitterDate { committer_date } => {
                c = c.set_committer_date(committer_date)
            }
            ChangeAction::Copy {
                path,
                content,
                parent,
                parent_path,
                ..
            } => {
                let parent: CommitIdentifier =
                    committed.get(&parent).map_or(parent.into(), |&c| c.into());
                c = c.add_file_with_copy_info(
                    path.as_slice(),
                    content,
                    (parent, parent_path.as_slice()),
                )
            }
        }
    }
    c
}

async fn derive<D: BonsaiDerivable>(
    ctx: &CoreContext,
    repo: &BlobRepo,
    csids: &[ChangesetId],
) -> Result<()> {
    let mgr = repo.repo_derived_data().manager();
    mgr.derive_exactly_batch::<D>(
        ctx,
        csids.to_vec(),
        BatchDeriveOptions::Parallel { gap_size: None },
        None,
    )
    .await
    .with_context(|| format!("Failed to derive {}", D::NAME))?;
    Ok(())
}

async fn derive_all(ctx: &CoreContext, repo: &BlobRepo, csids: &[ChangesetId]) -> Result<()> {
    let mercurial = async {
        derive::<MappedHgChangesetId>(ctx, repo, csids).await?;
        derive::<FilenodesOnlyPublic>(ctx, repo, csids).await?;
        Ok::<_, Error>(())
    };
    let unodes = async {
        derive::<RootUnodeManifestId>(ctx, repo, csids).await?;
        try_join!(
            derive::<RootBlameV2>(ctx, repo, csids),
            derive::<RootDeletedManifestV2Id>(ctx, repo, csids),
            derive::<RootFastlog>(ctx, repo, csids),
        )?;
        Ok::<_, Error>(())
    };
    try_join!(
        mercurial,
        unodes,
        derive::<RootFsnodeId>(ctx, repo, csids),
        derive::<RootSkeletonManifestId>(ctx, repo, csids),
        derive::<ChangesetInfo>(ctx, repo, csids),
        derive::<RootBasenameSuffixSkeletonManifest>(ctx, repo, csids),
    )?;
    Ok(())
}

#[cfg(test)]
mod test {
    use super::*;

    #[test]
    fn test_action_specs() -> Result<()> {
        assert_eq!(
            Action::new(
                "exists: A aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa"
            )?,
            Action::Exists {
                name: "A".to_string(),
                id: "aaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaaa".parse()?,
            }
        );
        assert_eq!(
            Action::new("bookmark: \"A-bookmark\" \"main\"/\"bookmark\"")?,
            Action::Bookmark {
                name: "A-bookmark".to_string(),
                bookmark: "main/bookmark".parse()?,
            }
        );
        assert_eq!(
            Action::new(
                "modify: _1 path/to/file \"this has \\xaa content\\n\\ton \\x02 lines with \\\"quotes\\\"\""
            )?,
            Action::Change {
                name: "_1".to_string(),
                change: ChangeAction::Modify {
                    path: b"path/to/file".to_vec(),
                    content: b"this has \xaa content\n\ton \x02 lines with \"quotes\"".to_vec(),
                }
            }
        );
        assert_eq!(
            Action::new("delete: x path/\"to a deleted file\"")?,
            Action::Change {
                name: "x".to_string(),
                change: ChangeAction::Delete {
                    path: b"path/to a deleted file".to_vec(),
                }
            }
        );
        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::bail;
use anyhow::Context;
use anyhow::Result;
use bookmarks::BookmarkKey;
use bookmarks::BookmarkUpdateReason;
use bookmarks::Bookmarks;
use bookmarks::BookmarksRef;
use clap::Parser;
use clap::Subcommand;
use mononoke_app::args::RepoArgs;
use mononoke_app::MononokeApp;
use mononoke_types::ChangesetId;

/// Modify a bookmark
///
/// Like `mononoke-admin bookmarks`, except performs fewer checks and allows
/// overriding of the bookmark update reason, so only suitable for use in
/// tests.
#[derive(Parser)]
pub struct CommandArgs {
    #[clap(flatten)]
    repo_args: RepoArgs,

    /// Bookmark modification to perform.
    #[clap(subcommand)]
    op: BookmarkOperation,
}

#[derive(Subcommand)]
pub enum BookmarkOperation {
    Create {
        /// Bookmark to create
        bookmark: BookmarkKey,

        /// Bookmark update reason
        #[clap(long, value_enum, default_value = "test-move")]
        reason: BookmarkUpdateReason,

        /// Changeset bookmark should be created at
        #[clap(long, short = 't')]
        to: ChangesetId,
    },
    Update {
        /// Bookmark to update
        bookmark: BookmarkKey,

        /// Bookmark update reason
        #[clap(long, value_enum, default_value = "test-move")]
        reason: BookmarkUpdateReason,

        /// Changeset bookmark is being moved from
        ///
        /// If specified, the bookmark update will only succeed if the
        /// bookmark currently points at this changeset.  If omitted, the
        /// bookmark is force-moved to the new location, no matter where it
        /// currently points.
        #[clap(long, short = 'f')]
        from: Option<ChangesetId>,

        /// Changeset bookmark should be moved to
        #[clap(long, short = 't')]
        to: ChangesetId,
    },
    Delete {
        /// Bookmark to delete
        bookmark: BookmarkKey,

        /// Bookmark update reason
        #[clap(long, value_enum, default_value = "test-move")]
        reason: BookmarkUpdateReason,

        /// Changeset bookmark is being deleted from
        ///
        /// If specified, the bookmark delete will only succeed if the
        /// bookmarks currently points at this changeset.  If omitted, the
        /// bookmark is force-deleted no matter where it points to.
        #[clap(long, short = 'f')]
        from: Option<ChangesetId>,
    },
}

#[facet::container]
pub struct Repo {
    #[facet]
    bookmarks: dyn Bookmarks,
}

pub async fn run(app: MononokeApp, args: CommandArgs) -> Result<()> {
    let ctx = app.new_basic_context();

    let repo: Repo = app
        .open_repo(&args.repo_args)
        .await
        .context("Failed to open repo")?;

    let mut txn = repo.bookmarks().create_transaction(ctx.clone());

    match args.op {
        BookmarkOperation::Create {
            bookmark,
            reason,
            to,
        } => {
            txn.create(&bookmark, to, reason)?;
        }
        BookmarkOperation::Update {
            bookmark,
            reason,
            from: Some(from),
            to,
        } => {
            txn.update(&bookmark, to, from, reason)?;
        }
        BookmarkOperation::Update {
            bookmark,
            reason,
            from: None,
            to,
        } => {
            txn.force_set(&bookmark, to, reason)?;
        }
        BookmarkOperation::Delete {
            bookmark,
            reason,
            from: Some(from),
        } => {
            txn.delete(&bookmark, from, reason)?;
        }
        BookmarkOperation::Delete {
            bookmark,
            reason,
            from: None,
        } => {
            txn.force_delete(&bookmark, reason)?;
        }
    }

    let success = txn
        .commit()
        .await
        .context("Failed to commit bookmark transaction")?;

    if !success {
        bail!("Bookmark transaction failed");
    }

    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::cell::RefCell;
use std::collections::HashMap;
use std::ops::Deref;
use std::sync::Arc;
use std::thread_local;
use std::time::Duration;

use anyhow::Result;
use arc_swap::ArcSwap;
use arc_swap::ArcSwapOption;
use cached_config::ConfigHandle;
use futures::future::poll_fn;
use futures::Future;
use futures::FutureExt;
use once_cell::sync::OnceCell;
use slog::debug;
use slog::error;
use slog::warn;
use slog::Logger;
use stats::prelude::*;
use tokio::runtime::Handle;
use tunables_derive::Tunables;
use tunables_structs::Tunables as TunablesStruct;

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

static TUNABLES: OnceCell<MononokeTunables> = OnceCell::new();
static TUNABLES_WORKER_STATE: OnceCell<TunablesWorkerState> = OnceCell::new();

thread_local! {
    static TUNABLES_OVERRIDE: RefCell<Option<Arc<MononokeTunables>>> = RefCell::new(None);
}

pub enum TunablesReference {
    Override(Arc<MononokeTunables>),
    Static(&'static MononokeTunables),
}

impl Deref for TunablesReference {
    type Target = MononokeTunables;

    fn deref(&self) -> &MononokeTunables {
        match self {
            Self::Override(r) => r.as_ref(),
            Self::Static(r) => r,
        }
    }
}

pub fn tunables() -> TunablesReference {
    TUNABLES_OVERRIDE.with(|tunables_override| match *tunables_override.borrow() {
        Some(ref arc) => TunablesReference::Override(arc.clone()),
        None => TunablesReference::Static(TUNABLES.get_or_init(MononokeTunables::default)),
    })
}

// These types exist to simplify code generation in tunables-derive
pub type TunableBool = ArcSwapOption<bool>;
pub type TunableI64 = ArcSwapOption<i64>;

pub type TunableString = ArcSwapOption<String>;
pub type TunableVecOfStrings = ArcSwapOption<Vec<String>>;

pub type TunableBoolByRepo = ArcSwap<HashMap<String, bool>>;
pub type TunableStringByRepo = ArcSwap<HashMap<String, String>>;
pub type TunableVecOfStringsByRepo = ArcSwap<HashMap<String, Vec<String>>>;
pub type TunableI64ByRepo = ArcSwap<HashMap<String, i64>>;

#[derive(Tunables, Default, Debug)]
pub struct MononokeTunables {
    mutation_advertise_for_infinitepush: TunableBool,
    mutation_accept_for_infinitepush: TunableBool,
    mutation_generate_for_draft: TunableBool,
    warm_bookmark_cache_poll_interval_ms: TunableI64,
    /// Don't read from the BookmarksSubscription when updating the WBC, and instead poll for the
    /// entire list of bookmarks on every iteration.
    warm_bookmark_cache_disable_subscription: TunableBool,
    /// Maximum age of bookmarks subscriptions.
    bookmark_subscription_max_age_ms: TunableI64,
    bookmark_subscription_protect_master: TunableBool,
    max_scuba_msg_length: TunableI64,
    wishlist_read_qps: TunableI64,
    wishlist_write_qps: TunableI64,
    edenapi_large_tree_metadata_limit: TunableI64,
    edenapi_req_dumper_sample_ratio: TunableI64,
    command_monitor_interval: TunableI64,
    command_monitor_remote_logging: TunableI64,
    edenapi_request_monitor_interval: TunableI64,
    // Log all getfiles/gettreepack requests for paths that start with prefix
    // in a particular repo
    undesired_path_repo_name_to_log: TunableString,
    undesired_path_prefix_to_log: TunableString,
    undesired_path_regex_to_log: TunableString,
    pushrebase_disable_rebased_commit_validation: TunableBool,
    filenodes_disabled: TunableBool,
    filenodes_master_fallback_ratio: TunableI64,

    deduplicated_put_sampling_rate: TunableI64,
    disable_repo_client_warm_bookmarks_cache: TunableBool,
    remotefilelog_file_history_limit: TunableI64,
    disable_hooks_on_plain_push: TunableBool,
    run_hooks_on_additional_changesets: TunableBool,
    hooks_additional_changesets_limit: TunableI64,
    // SCS scuba sampling knobs
    scs_popular_methods_sampling_rate: TunableI64,
    scs_other_methods_sampling_rate: TunableI64,
    // When false error logs are never sampled
    scs_error_log_sampling: TunableBool,
    redacted_logging_sampling_rate: TunableI64,
    redaction_config_from_xdb: TunableBool,
    repo_client_bookmarks_timeout_secs: TunableI64,
    repo_client_clone_timeout_secs: TunableI64,
    repo_client_default_timeout_secs: TunableI64,
    repo_client_getbundle_timeout_secs: TunableI64,
    repo_client_getpack_timeout_secs: TunableI64,
    repo_client_concurrent_blob_uploads: TunableI64,
    repo_client_max_nodes_in_known_method: TunableI64,
    // How many trees is getting prepared at once
    repo_client_gettreepack_buffer_size: TunableI64,
    derived_data_slow_derivation_threshold_secs: TunableI64,
    disable_running_hooks_in_pushredirected_repo: TunableBool,
    scs_request_read_qps: TunableI64,
    scs_request_write_qps: TunableI64,
    enable_logging_commit_rewrite_data: TunableBool,
    // All blobstore read request with size bigger than
    // this threshold will be logged to scuba
    blobstore_read_size_logging_threshold: TunableI64,
    hash_validation_percentage: TunableI64,
    // Filter out commits that we already have in infinitepush. Shouldn't be needed if we have a
    // client exchanging commits with us, but when processing bundled uploads (i.e. commit cloud
    // filling), it might help a lot.
    filter_pre_existing_commits_on_infinitepush: TunableBool,
    backfill_read_qps: TunableI64,
    backfill_write_qps: TunableI64,
    disable_commit_scribe_logging_scs: TunableBool,
    xrepo_sync_disable_all_syncs: TunableBool,
    xrepo_disable_commit_sync_lease: TunableBool,

    // Use Background session class while deriving data. This makes derived data not write
    // data to blobstore sync queue if a write was successful to the main blobstore.
    derived_data_use_background_session_class: TunableBoolByRepo,
    commit_cloud_use_background_session_class: TunableBool,
    multiplex_blobstore_background_session_timeout_ms: TunableI64,

    allow_change_xrepo_mapping_extra: TunableBool,

    // Disable EdenAPI in http_service.
    disable_http_service_edenapi: TunableBool,

    // Disable putting hydrating manifests in .hg
    disable_hydrating_manifests_in_dot_hg: TunableBool,

    // Rendez vous configuration.
    rendezvous_dispatch_delay_ms: TunableI64,
    rendezvous_dispatch_max_threshold: TunableI64,

    unbundle_limit_num_of_commits_in_push: TunableI64,

    // Maximium negative caching age of a blob, in milliseconds
    // Negative means to not use weak consistency at all
    manifold_weak_consistency_max_age_ms: TunableI64,

    // -1: No override, use manifold server side config and rollout checks
    // if set to > -1, set the client side option to override (see manifoldblob code)
    manifold_request_priority_override: TunableI64,

    // Frequency at which to collect SQL connection pool stats
    sql_connection_pool_stats_collection_interval_ms: TunableI64,

    bookmarks_cache_ttl_ms: TunableI64,

    // Disable running SaveMappingPushrebaseHook on every Pushrebase
    disable_save_mapping_pushrebase_hook: TunableBool,

    // Set to 0 to disable compression
    zstd_compression_level: TunableI64,

    // Commits that aren't related (i.e. that are not ancestors of each other)
    // can be derived in parallel, and that's what derived data does.
    // derived_data_parallel_derivation_buffer limits
    // how many commits will be derived at once.
    derived_data_parallel_derivation_buffer: TunableI64,

    // Tunables to disable derived data derivation either for the full repo
    // or for specific derived data types inside a repo
    all_derived_data_disabled: TunableBoolByRepo,
    derived_data_types_disabled: TunableVecOfStringsByRepo,
    // How often to check if derived data is disabled or not
    derived_data_disabled_watcher_delay_secs: TunableI64,

    // Stops deriving on derivation workers. Will not drain Derivation Queue
    derived_data_disable_derivation_workers: TunableBool,

    // How long to wait before worker retries in case of an error
    // or empty Derivation queue.
    derivation_worker_sleep_duration: TunableI64,

    // How long client should wait between polls of Derived data service
    derivation_request_retry_delay: TunableI64,

    // Sets the size of the batch for derivaiton.
    derivation_batch_size: TunableI64,

    // Maximum time to wait for remote derivation request to finish in secs
    // before falling back to local derivation
    remote_derivation_fallback_timeout_secs: TunableI64,

    // Timeout for derivation request on service.
    dds_request_timeout: TunableI64,

    // Disable the parallel derivation for DM and default to serial
    deleted_manifest_disable_new_parallel_derivation: TunableBool,

    // multiplexed blobstore is_present/get new semantics rollout
    multiplex_blobstore_get_do_queue_lookup: TunableBool,
    multiplex_blobstore_is_present_do_queue_lookup: TunableBool,

    // Not in use.
    // TODO(mitrandir): clean it up
    fastlog_use_mutable_renames: TunableBoolByRepo,
    // Disable mutable renames for fastlog in case they cause problems.
    fastlog_disable_mutable_renames: TunableBoolByRepo,
    megarepo_api_dont_set_file_mutable_renames: TunableBool,
    megarepo_api_dont_set_directory_mutable_renames: TunableBool,

    force_unode_v2: TunableBool,
    fastlog_use_gen_num_traversal: TunableBool,

    // Changing the value of this tunable forces all mononoke instances
    // to reload segmented changelog. One can also specify jitter (or use default)
    segmented_changelog_force_reload: TunableI64ByRepo,
    segmented_changelog_force_reload_jitter_secs: TunableI64,

    // Override default progress logging sampling rate for segmented changelog parts
    segmented_changelog_idmap_log_sampling_rate: TunableI64,
    segmented_changelog_tailer_log_sampling_rate: TunableI64,

    // How many commits to walk back from the client heads before failing to rebuild SC
    segmented_changelog_client_max_commits_to_traverse: TunableI64,

    // Timeout for is_present call for multiplexed blobstore
    is_present_timeout_ms: TunableI64,

    // What timeout to use when doing filenode lookup.
    // Usually filenode lookup is used while generating hg changesets
    filenode_lookup_timeout_ms: TunableI64,

    // Sampling ratio percentage for warm boomark cache.
    warm_bookmark_cache_logging_sampling_pct: TunableI64,

    // Setting this tunable to a new non-zero value and restarting
    // mononoke hosts will invalidate the cache
    blobstore_memcache_sitever: TunableI64,
    // Setting this tunable to a new non-zero value and restarting
    // mononoke hosts will invalidate the cache
    sql_memcache_sitever: TunableI64,

    // Setting this tunable to a new non-zero value and restarting
    // mononoke hosts will invalidate bonsai_hg_mapping cache
    bonsai_hg_mapping_sitever: TunableI64,

    // Setting this tunable to a new non-zero value will update the
    // TTL for the mutation store cache
    hg_mutation_store_caching_ttl_secs: TunableI64,

    // Setting this tunable to a new non-zero value and restarting
    // mononoke hosts will invalidate hg mutation store cache
    hg_mutation_store_sitever: TunableI64,

    // EdenAPI requests that take long than this get logged unsampled
    edenapi_unsampled_duration_threshold_ms: TunableI64,

    // EdenAPI requests that take long than this get logged unsampled by request dumper
    edenapi_req_dumper_unsampled_duration_threshold_ms: TunableI64,

    // EdenAPI high load threshold (max number of concurrent requests to pass health check)
    edenapi_high_load_threshold: TunableI64,

    // Setting this tunable to a new non-zero value and restarting
    // mononoke hosts will invalidate mutable renames cache
    mutable_renames_sitever: TunableI64,

    // Setting this will enable hooks on pushrebase bookmark moves that were
    // initiated by a service (for user-initiated moves we always run hooks).
    // Most of them will run in no-op mode but some of them (namely,
    // verify_integrity) will need to do some work like logging and this is
    // why it may be useful to run them.
    enable_hooks_on_service_pushrebase: TunableBool,

    // Control whether the BYPASS_READONLY pushvar is restricted by an ACL
    enforce_bypass_readonly_acl: TunableBool,

    // Boolean to batch requests sent to Land Service
    batching_to_land_service: TunableBool,

    // Which region writes should be done to, in order to minimise latency.
    // This should align with underlying storage (SQL/Manifold) write regions.
    // Notice writes still work from any region, and this field is not necessarily
    // enforced.
    preferred_write_region: TunableString,

    // The replication_status call is problematic for SQL so we're experimenting
    // with removing it, but this tunable can be used as a quick killswitch to
    // enable them again.
    sql_lag_monitoring_blocklist: TunableVecOfStrings,

    // If set, the hook won't be created at all
    disable_check_write_permissions_hook: TunableBool,
    // If set, the check result will be discarded for user identities
    log_only_for_users_in_cwp_hook: TunableBool,
    // If set, the check result will be discarded for service identities
    log_only_for_services_in_cwp_hook: TunableBool,

    // If set, the wireproto implementation will only log the repo write ACL
    // check result.
    log_only_wireproto_write_acl: TunableBool,

    // If set the `draft` ACL action will be checked and logged on draft access
    // Unless `enforce_draft_acl` is set `read` action will still be used for
    // granting access.
    log_draft_acl_failures: TunableBool,

    // If set the `draft` ACL action will be used for `draft` access.
    enforce_draft_acl: TunableBool,

    // Force local pushrebase instead of talking to SCS or Land Service
    force_local_pushrebase: TunableBool,

    // Disable SCS pushredirecting commits landed on a small repo to a big one
    disable_scs_pushredirect: TunableBool,

    // Enable usage of basename_suffix_skeleton_manifest in commit_find_files
    disable_basename_suffix_skeleton_manifest: TunableBool,
    // Enable using BSSM for suffix queries. Might be inneficient for broad suffixes (like .php)
    enable_bssm_suffix_query: TunableBool,

    // List of targets in AOSP megarepo to apply squashing config overrides
    megarepo_squashing_config_override_targets: TunableVecOfStringsByRepo,
    // Override squashing limit for listed targets
    megarepo_override_squashing_limit: TunableI64ByRepo,
    // Override author check during squashing
    megarepo_override_author_check: TunableBoolByRepo,

    // Disable SQL queries being retried after admission control errors
    disable_sql_auto_retries: TunableBool,
    // Disable SQL queries being cached using `cacheable` keyword
    disable_sql_auto_cache: TunableBool,
    // Disable using rendezvous for batching WAL deletes.
    // TODO: delete once it using WAL here shows to be stable
    wal_disable_rendezvous_on_deletes: TunableBool,
    // Enable derivation on service per repo
    enable_remote_derivation: TunableBoolByRepo,

    // Enable using the new commit graph for is_ancestor queries
    enable_new_commit_graph_is_ancestor: TunableBoolByRepo,
    // Enable using new commit graph in common_base_with
    enable_new_commit_graph_common_base_with: TunableBoolByRepo,
    // Enable using new commit graph in finding commit history
    enable_new_commit_graph_commit_history: TunableBoolByRepo,
    // Enable using new commit graph in finding commit path history
    enable_new_commit_graph_commit_path_history: TunableBoolByRepo,

    // Enable streaming commit graph EdenAPI endpoint.
    enable_streaming_commit_graph_edenapi_endpoint: TunableBool,

    // Disable all prefetching in the commit graph
    disable_commit_graph_prefetch: TunableBool,
    // Disable memcache for commit graph prefetching
    disable_commit_graph_memcache_for_prefetch: TunableBool,
    // Max number of steps to make when prefetching
    commit_graph_prefetch_step_limit: TunableI64,

    // Interval between preloaded commit graph reloads in secs.
    preloaded_commit_graph_reloading_interval_secs: TunableI64,

    // Disable the fix to use isolation level read committed
    disable_wal_read_committed: TunableBool,

    // Disable sharing of large reads
    disable_large_blob_read_deduplication: TunableBool,
    // Enable double writing of Content Metadata.
    enable_content_metadata_double_writing: TunableBool,

    // Skip backsyncing for empty commits (except mapping changes via extras and merges)
    cross_repo_skip_backsyncing_ordinary_empty_commits: TunableBoolByRepo,
}

fn log_tunables(tunables: &TunablesStruct) -> String {
    serde_json::to_string(tunables)
        .unwrap_or_else(|e| format!("failed to serialize tunables: {}", e))
}

pub fn init_tunables_worker(
    logger: Logger,
    config_handle: ConfigHandle<TunablesStruct>,
    runtime_handle: Handle,
) -> Result<()> {
    init_tunables(&logger, &config_handle)?;
    if TUNABLES_WORKER_STATE
        .set(TunablesWorkerState {
            config_handle,
            logger,
        })
        .is_err()
    {
        panic!("Two or more tunables update threads exist at the same time");
    }
    runtime_handle.spawn(wait_and_update());

    Ok(())
}

pub fn init_tunables(logger: &Logger, config_handle: &ConfigHandle<TunablesStruct>) -> Result<()> {
    let tunables = config_handle.get();
    debug!(logger, "Initializing tunables: {}", log_tunables(&tunables));
    update_tunables(tunables)
}

/// Tunables are updated when the underlying config source notifies of a change.
/// Call this to force update them to the latest value provided by the config source.
/// Meant to be used in tests.
/// NOTE: if tunables are fetched from Configerator, you need to force update it as well.
pub fn force_update_tunables() {
    let state = TUNABLES_WORKER_STATE
        .get()
        .expect("Tunables worker state uninitialised");
    wait_and_update_iteration(state.config_handle.get(), &state.logger);
}

struct TunablesWorkerState {
    config_handle: ConfigHandle<TunablesStruct>,
    logger: Logger,
}

async fn wait_and_update() {
    let state = TUNABLES_WORKER_STATE
        .get()
        .expect("Tunables worker state uninitialised");
    let mut config_watcher = state
        .config_handle
        .watcher()
        .expect("Tunable backed by static config source");
    loop {
        match config_watcher.wait_for_next().await {
            Ok(new_tunables) => wait_and_update_iteration(new_tunables, &state.logger),
            Err(e) => {
                error!(
                    state.logger,
                    "Error in fetching latest config for tunable: {}.\n Exiting tunable updater", e
                );
                // Set the refresh failure count counter so that the oncall can be alerted
                // based on this metric
                STATS::refresh_failure_count.add_value(1);
                return;
            }
        }
    }
}

fn wait_and_update_iteration(new_tunables: Arc<TunablesStruct>, logger: &Logger) {
    debug!(
        logger,
        "Updating tunables to new: {}",
        log_tunables(&new_tunables),
    );
    if let Err(e) = update_tunables(new_tunables) {
        warn!(logger, "Failed to refresh tunables: {}", e);
        // Set the refresh failure count counter so that the oncall can be alerted
        // based on this metric
        STATS::refresh_failure_count.add_value(1);
    } else {
        // Add a value of 0 so that the counter won't get dead even if there
        // are no errors
        STATS::refresh_failure_count.add_value(0);
    }
}

fn update_tunables(new_tunables: Arc<TunablesStruct>) -> Result<()> {
    let tunables = tunables();
    tunables.update_bools(&new_tunables.killswitches);
    tunables.update_ints(&new_tunables.ints);
    tunables.update_strings(&new_tunables.strings);
    tunables.update_vec_of_strings(&new_tunables.vec_of_strings);

    if let Some(killswitches_by_repo) = &new_tunables.killswitches_by_repo {
        tunables.update_by_repo_bools(killswitches_by_repo);
    }

    if let Some(ints_by_repo) = &new_tunables.ints_by_repo {
        tunables.update_by_repo_ints(ints_by_repo);
    }

    if let Some(vec_of_strings_by_repo) = &new_tunables.vec_of_strings_by_repo {
        tunables.update_by_repo_vec_of_strings(vec_of_strings_by_repo);
    }
    Ok(())
}

/// A helper function to override tunables during a closure's execution.
/// This is useful for unit tests.
pub fn with_tunables<T>(new_tunables: MononokeTunables, f: impl FnOnce() -> T) -> T {
    TUNABLES_OVERRIDE.with(|t| *t.borrow_mut() = Some(Arc::new(new_tunables)));

    let res = f();

    TUNABLES_OVERRIDE.with(|tunables| *tunables.borrow_mut() = None);

    res
}

pub fn with_tunables_async<Out, Fut: Future<Output = Out> + Unpin>(
    new_tunables: MononokeTunables,
    fut: Fut,
) -> impl Future<Output = Out> {
    with_tunables_async_arc(Arc::new(new_tunables), fut)
}

pub fn with_tunables_async_arc<Out, Fut: Future<Output = Out> + Unpin>(
    new_tunables: Arc<MononokeTunables>,
    mut fut: Fut,
) -> impl Future<Output = Out> {
    poll_fn(move |cx| {
        TUNABLES_OVERRIDE.with(|t| *t.borrow_mut() = Some(new_tunables.clone()));

        let res = fut.poll_unpin(cx);

        TUNABLES_OVERRIDE.with(|tunables| *tunables.borrow_mut() = None);

        res
    })
}

pub fn override_tunables(new_tunables: Option<Arc<MononokeTunables>>) {
    TUNABLES_OVERRIDE.with(|t| *t.borrow_mut() = new_tunables);
}

/// Get a duration from a tunable or a provided default
/// value if the tunable is not set.
pub fn get_duration_from_tunable_or(
    get_tunable: impl Fn(&MononokeTunables) -> Option<i64>,
    into_duration: impl Fn(u64) -> Duration,
    default: u64,
) -> Duration {
    into_duration(get_tunable(tunables().deref()).map_or(default, |i| i as u64))
}

#[cfg(test)]
mod test {
    use std::collections::HashMap;

    use maplit::hashmap;

    use super::*;

    #[derive(Tunables, Default)]
    struct TestTunables {
        boolean: TunableBool,
        num: TunableI64,
        string: TunableString,
        vecofstrings: TunableVecOfStrings,

        repobool: TunableBoolByRepo,
        repobool2: TunableBoolByRepo,

        repoint: TunableI64ByRepo,
        repoint2: TunableI64ByRepo,

        repostr: TunableStringByRepo,
        repostr2: TunableStringByRepo,

        repovecofstrings: TunableVecOfStringsByRepo,
    }

    #[derive(Tunables, Default)]
    struct EmptyTunables {}

    fn s(a: &str) -> String {
        a.to_string()
    }

    #[test]
    fn test_override_tunables() {
        assert!(tunables().wishlist_write_qps().is_none());

        let res = with_tunables(
            MononokeTunables {
                wishlist_write_qps: ArcSwapOption::from(Some(Arc::new(2))),
                ..MononokeTunables::default()
            },
            || tunables().wishlist_write_qps().unwrap_or_default(),
        );

        assert_eq!(res, 2);
        assert!(tunables().wishlist_write_qps().is_none());
    }

    #[test]
    fn test_empty_tunables() {
        let bools = HashMap::new();
        let ints = HashMap::new();
        let empty = EmptyTunables::default();

        empty.update_bools(&bools);
        empty.update_ints(&ints);
        empty.update_strings(&HashMap::new());
        empty.update_vec_of_strings(&HashMap::new());
    }

    #[test]
    fn test_update_bool() {
        let mut d = HashMap::new();
        d.insert(s("boolean"), true);

        let test = TestTunables::default();
        assert!(test.boolean().is_none());
        test.update_bools(&d);
        assert_eq!(test.boolean(), Some(true));
    }

    #[test]
    fn test_revert_update() {
        // booleans
        let mut d = hashmap! {};
        d.insert(s("boolean"), true);

        let test = TestTunables::default();
        assert!(test.boolean().is_none());
        test.update_bools(&d);
        assert_eq!(test.boolean(), Some(true));

        test.update_bools(&hashmap! {});
        assert!(test.boolean().is_none());

        // ints
        let test = TestTunables::default();
        test.update_ints(&hashmap! { s("num") => 1});
        assert_eq!(test.num(), Some(1));

        test.update_ints(&hashmap! {});
        assert_eq!(test.num(), None);

        // strings
        let test = TestTunables::default();
        test.update_strings(&hashmap! { s("string") => s("string")});
        assert_eq!(test.string(), Some(Arc::new(s("string"))));

        test.update_strings(&hashmap! {});
        assert!(test.string().is_none());

        // by repo bools
        assert_eq!(test.by_repo_repobool("repo"), None);
        assert_eq!(test.by_repo_repobool("repo2"), None);

        test.update_by_repo_bools(&hashmap! {
            s("repo") => hashmap! {
                s("repobool") => true,
            },
            s("repo2") => hashmap! {
                s("repobool") => false,
            }
        });
        assert_eq!(test.by_repo_repobool("repo"), Some(true));
        assert_eq!(test.by_repo_repobool("repo2"), Some(false));

        test.update_by_repo_bools(&hashmap! {});
        assert_eq!(test.by_repo_repobool("repo"), None);
        assert_eq!(test.by_repo_repobool("repo2"), None);
    }

    #[test]
    fn test_update_int() {
        let mut d = HashMap::new();
        d.insert(s("num"), 10);

        let test = TestTunables::default();
        assert!(test.num().is_none());
        test.update_ints(&d);
        assert_eq!(test.num(), Some(10));
    }

    #[test]
    fn test_missing_int() {
        let mut d = HashMap::new();
        d.insert(s("missing"), 10);

        let test = TestTunables::default();
        assert!(test.num().is_none());
        test.update_ints(&d);
        assert!(test.num().is_none());
    }

    #[test]
    fn update_string() {
        let mut d = HashMap::new();
        d.insert(s("string"), s("value"));

        let test = TestTunables::default();
        assert!(test.string().is_none());
        test.update_strings(&d);
        assert_eq!(test.string().unwrap().as_str(), "value");
    }

    #[test]
    fn update_vec_of_strings() {
        let mut d = HashMap::new();
        d.insert(s("vecofstrings"), vec![s("value"), s("value2")]);

        let test = TestTunables::default();
        assert!(&test.vecofstrings().is_none());
        test.update_vec_of_strings(&d);
        assert_eq!(
            &test.vecofstrings().unwrap().as_slice(),
            &[s("value"), s("value2")]
        );
    }

    #[test]
    fn update_by_repo_bool() {
        let test = TestTunables::default();

        assert_eq!(test.by_repo_repobool("repo"), None);
        assert_eq!(test.by_repo_repobool("repo2"), None);

        test.update_by_repo_bools(&hashmap! {
            s("repo") => hashmap! {
                s("repobool") => true,
            },
            s("repo2") => hashmap! {
                s("repobool") => true,
            }
        });
        assert_eq!(test.by_repo_repobool("repo"), Some(true));
        assert_eq!(test.by_repo_repobool("repo2"), Some(true));

        test.update_by_repo_bools(&hashmap! {
            s("repo") => hashmap! {
                s("repobool") => true,
            }
        });
        assert_eq!(test.by_repo_repobool("repo2"), None);

        test.update_by_repo_bools(&hashmap! {
            s("repo") => hashmap! {
                s("repobool") => false,
            }
        });
        assert_eq!(test.by_repo_repobool("repo"), Some(false));

        test.update_by_repo_bools(&hashmap! {
            s(":default:") => hashmap! {
                s("repobool") => true,
            },
            s("repo") => hashmap! {
                s("repobool") => false,
            }
        });

        assert_eq!(test.by_repo_repobool("repo"), Some(false));
        assert_eq!(test.by_repo_repobool("repo2"), Some(true));

        test.update_by_repo_bools(&hashmap! {
            s(":override:") => hashmap! {
                s("repobool") => true,
            },
            s("repo") => hashmap! {
                s("repobool") => false,
            }
        });

        assert_eq!(test.by_repo_repobool("repo"), Some(true));
        assert_eq!(test.by_repo_repobool("repo2"), Some(true));
    }

    #[test]
    fn update_by_repo_two_bools() {
        let test = TestTunables::default();
        assert_eq!(test.by_repo_repobool("repo"), None);
        assert_eq!(test.by_repo_repobool2("repo"), None);

        test.update_by_repo_bools(&hashmap! {
            s("repo") => hashmap! {
                s("repobool") => true,
                s("repobool2") => true,
            }
        });

        assert_eq!(test.by_repo_repobool("repo"), Some(true));
        assert_eq!(test.by_repo_repobool2("repo"), Some(true));

        test.update_by_repo_bools(&hashmap! {
            s("repo") => hashmap! {
                s("repobool") => true,
                s("repobool2") => false,
            }
        });

        assert_eq!(test.by_repo_repobool("repo"), Some(true));
        assert_eq!(test.by_repo_repobool2("repo"), Some(false));
    }

    #[test]
    fn update_by_repo_str() {
        let test = TestTunables::default();

        assert_eq!(test.by_repo_repostr("repo"), None);
        assert_eq!(test.by_repo_repostr("repo2"), None);

        test.update_by_repo_strings(&hashmap! {
            s("repo") => hashmap! {
                s("repostr") => s("hello"),
            },
            s("repo2") => hashmap! {
                s("repostr") => s("world"),
            },
        });
        assert_eq!(test.by_repo_repostr("repo"), Some(s("hello")));
        assert_eq!(test.by_repo_repostr("repo2"), Some(s("world")));

        test.update_by_repo_strings(&hashmap! {
            s("repo") => hashmap! {
                s("repostr") => s("hello2"),
            },
        });
        assert_eq!(test.by_repo_repostr("repo"), Some(s("hello2")));
        assert_eq!(test.by_repo_repostr("repo2"), None);

        test.update_by_repo_strings(&hashmap! {
            s(":default:") => hashmap! {
                s("repostr") => s("hello3")
            },
            s("repo") => hashmap! {
                s("repostr") => s("hello"),
            },
        });

        assert_eq!(test.by_repo_repostr("repo"), Some(s("hello")));
        assert_eq!(test.by_repo_repostr("repo2"), Some(s("hello3")));

        test.update_by_repo_strings(&hashmap! {
            s(":override:") => hashmap! {
                s("repostr") => s("hello3")
            },
            s("repo") => hashmap! {
                s("repostr") => s("hello"),
            },
        });

        assert_eq!(test.by_repo_repostr("repo"), Some(s("hello3")));
        assert_eq!(test.by_repo_repostr("repo2"), Some(s("hello3")));
    }

    #[test]
    fn update_by_repo_two_strs() {
        let test = TestTunables::default();
        assert_eq!(test.by_repo_repostr("repo"), None);
        assert_eq!(test.by_repo_repostr2("repo"), None);

        test.update_by_repo_strings(&hashmap! {
            s("repo") => hashmap! {
                s("repostr") => s("hello"),
                s("repostr2") => s("world"),
            }
        });

        assert_eq!(test.by_repo_repostr("repo"), Some(s("hello")));
        assert_eq!(test.by_repo_repostr2("repo"), Some(s("world")));

        test.update_by_repo_strings(&hashmap! {
            s("repo") => hashmap! {
                s("repostr") => s("hello2"),
            }
        });

        assert_eq!(test.by_repo_repostr("repo"), Some(s("hello2")));
        assert_eq!(test.by_repo_repostr2("repo"), None);
    }

    #[test]
    fn update_by_repo_int() {
        let test = TestTunables::default();

        assert_eq!(test.by_repo_repoint("repo"), None);
        assert_eq!(test.by_repo_repoint("repo2"), None);

        test.update_by_repo_ints(&hashmap! {
            s("repo") => hashmap! {
                s("repoint") => 1,
            },
            s("repo2") => hashmap! {
                s("repoint") => 2,
            },
        });
        assert_eq!(test.by_repo_repoint("repo"), Some(1));
        assert_eq!(test.by_repo_repoint("repo2"), Some(2));

        test.update_by_repo_ints(&hashmap! {
            s("repo") => hashmap! {
                s("repoint") => 3,
            },
        });
        assert_eq!(test.by_repo_repoint("repo"), Some(3));
        assert_eq!(test.by_repo_repoint("repo2"), None);

        test.update_by_repo_ints(&hashmap! {
            s(":default:") => hashmap! {
                s("repoint") => 4
            },
            s("repo") => hashmap! {
                s("repoint") => 1,
            },
        });

        assert_eq!(test.by_repo_repoint("repo"), Some(1));
        assert_eq!(test.by_repo_repoint("repo2"), Some(4));

        test.update_by_repo_ints(&hashmap! {
            s(":override:") => hashmap! {
                s("repoint") => 4
            },
            s("repo") => hashmap! {
                s("repoint") => 1,
            },
        });

        assert_eq!(test.by_repo_repoint("repo"), Some(4));
        assert_eq!(test.by_repo_repoint("repo2"), Some(4));
    }

    #[test]
    fn update_by_repo_two_ints() {
        let test = TestTunables::default();
        assert_eq!(test.by_repo_repoint("repo"), None);
        assert_eq!(test.by_repo_repoint2("repo"), None);

        test.update_by_repo_ints(&hashmap! {
            s("repo") => hashmap! {
                s("repoint") => 1,
                s("repoint2") => 2,
            }
        });

        assert_eq!(test.by_repo_repoint("repo"), Some(1));
        assert_eq!(test.by_repo_repoint2("repo"), Some(2));

        test.update_by_repo_ints(&hashmap! {
            s("repo") => hashmap! {
                s("repoint") => 3
            }
        });

        assert_eq!(test.by_repo_repoint("repo"), Some(3));
        assert_eq!(test.by_repo_repoint2("repo"), None);
    }

    #[test]
    fn update_by_repo_vec_of_strings() {
        let test = TestTunables::default();
        assert_eq!(test.by_repo_repovecofstrings("repo"), None);

        test.update_by_repo_vec_of_strings(&hashmap! {
            s("repo") => hashmap! {
                s("unrelated") => vec![s("val1"), s("val2")],
            }
        });
        assert_eq!(test.by_repo_repovecofstrings("repo"), None);

        test.update_by_repo_vec_of_strings(&hashmap! {
            s("repo") => hashmap! {
                s("repovecofstrings") => vec![s("val1"), s("val2")],
            }
        });

        assert_eq!(
            test.by_repo_repovecofstrings("repo"),
            Some(vec![s("val1"), s("val2")])
        );

        test.update_by_repo_vec_of_strings(&hashmap! {
            s(":default:") => hashmap! {
                s("repovecofstrings") => vec![s("val3"), s("val4")],
            },
            s("repo") => hashmap! {
                s("repovecofstrings") => vec![s("val1"), s("val2")],
            },
        });

        assert_eq!(
            test.by_repo_repovecofstrings("repo"),
            Some(vec![s("val1"), s("val2")])
        );
        assert_eq!(
            test.by_repo_repovecofstrings("repo2"),
            Some(vec![s("val3"), s("val4")])
        );

        test.update_by_repo_vec_of_strings(&hashmap! {
            s(":override:") => hashmap! {
                s("repovecofstrings") => vec![s("val3"), s("val4")],
            },
            s("repo") => hashmap! {
                s("repovecofstrings") => vec![s("val1"), s("val2")],
            },
        });

        assert_eq!(
            test.by_repo_repovecofstrings("repo"),
            Some(vec![s("val3"), s("val4")])
        );
        assert_eq!(
            test.by_repo_repovecofstrings("repo2"),
            Some(vec![s("val3"), s("val4")])
        );
    }

    #[fbinit::test]
    async fn test_with_tunables_async(_fb: fbinit::FacebookInit) {
        let res = with_tunables_async(
            MononokeTunables {
                wishlist_write_qps: ArcSwapOption::from(Some(Arc::new(2))),
                ..MononokeTunables::default()
            },
            async { tunables().wishlist_write_qps().unwrap_or_default() }.boxed(),
        )
        .await;

        assert_eq!(res, 2);
    }
}


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

extern crate proc_macro;

use proc_macro2::TokenStream;
use quote::quote;
use syn::parse_macro_input;
use syn::Data;
use syn::DeriveInput;
use syn::Fields;
use syn::Ident;
use syn::Type;

const UNIMPLEMENTED_MSG: &str = "Only AtomicBool and AtomicI64 are supported";
const STRUCT_FIELD_MSG: &str = "Only implemented for named fields of a struct";

#[derive(Clone, PartialEq)]
enum TunableType {
    Bool,
    I64,
    String,
    VecOfStrings,
    ByRepoBool,
    ByRepoString,
    ByRepoI64,
    ByRepoVecOfStrings,
}

#[proc_macro_derive(Tunables)]
// This proc macro accepts a struct and provides methods that get the atomic
// values stored inside of it. It does this by generating methods
// named get_<field>(). The macro also generates methods that update the
// atomic values inside of the struct, using a provided HashMap.
pub fn derive_tunables(input: proc_macro::TokenStream) -> proc_macro::TokenStream {
    let parsed_input = parse_macro_input!(input as DeriveInput);

    let struct_name = parsed_input.ident;
    let names_and_types = parse_names_and_types(parsed_input.data).into_iter();

    let getter_methods = generate_getter_methods(names_and_types.clone());
    let updater_methods = generate_updater_methods(names_and_types);

    let expanded = quote! {
        impl #struct_name {
            #updater_methods
            #getter_methods
        }
    };

    expanded.into()
}

impl TunableType {
    fn external_type(&self) -> TokenStream {
        match self {
            Self::Bool => quote! { bool },
            Self::I64 => quote! { i64 },
            Self::String => quote! { Arc<String> },
            Self::VecOfStrings => quote! { Arc<Vec<String>> },
            Self::ByRepoBool => quote! { bool },
            Self::ByRepoString => quote! { String },
            Self::ByRepoI64 => quote! { i64 },
            Self::ByRepoVecOfStrings => quote! { Vec<String> },
        }
    }

    fn by_repo_value_type(&self) -> TokenStream {
        match self {
            Self::Bool | Self::I64 | Self::String | Self::VecOfStrings => {
                panic!("Expected ByRepo flavor of tunable")
            }
            Self::ByRepoBool => quote! { bool },
            Self::ByRepoI64 => quote! { i64 },
            Self::ByRepoString => quote! { String },
            Self::ByRepoVecOfStrings => quote! { Vec<String> },
        }
    }

    fn update_container_type(&self) -> TokenStream {
        match self {
            Self::Bool => quote! { HashMap<String, bool> },
            Self::I64 => quote! { HashMap<String, i64> },
            Self::String => quote! { HashMap<String, String> },
            Self::VecOfStrings => quote! { HashMap<String, Vec<String>> },
            Self::ByRepoBool => quote! { HashMap<String, HashMap<String, bool>> },
            Self::ByRepoString => quote! { HashMap<String, HashMap<String, String>> },
            Self::ByRepoI64 => quote! { HashMap<String, HashMap<String, i64>> },
            Self::ByRepoVecOfStrings => quote! { HashMap<String, HashMap<String, Vec<String>>> },
        }
    }

    fn generate_getter_method(&self, name: Ident) -> TokenStream {
        let external_type = self.external_type();

        match &self {
            Self::Bool | Self::I64 => {
                quote! {
                    pub fn #name(&self) -> Option<#external_type> {
                        self.#name.load_full().map(|arc_v| *arc_v)
                    }
                }
            }
            Self::String | Self::VecOfStrings => {
                quote! {
                    pub fn #name(&self) -> Option<#external_type> {
                        self.#name.load_full()
                    }
                }
            }
            Self::ByRepoBool | Self::ByRepoI64 | Self::ByRepoString | Self::ByRepoVecOfStrings => {
                let by_repo_method = quote::format_ident!("by_repo_{}", name);
                quote! {
                    pub fn #by_repo_method(&self, repo: &str) -> Option<#external_type> {
                        // If :override: has a value set for this tunable then use it, otherwise
                        // lookup the repo specific value, if not found lookup the value for :default:.
                        let values = self.#name.load_full();
                        values.get(":override:").or_else(|| values.get(repo)).or_else(|| values.get(":default:")).cloned()
                    }
                }
            }
        }
    }
}

fn generate_getter_methods<I>(names_and_types: I) -> TokenStream
where
    I: Iterator<Item = (Ident, TunableType)> + std::clone::Clone,
{
    let mut methods = TokenStream::new();

    for (name, ty) in names_and_types {
        methods.extend(ty.generate_getter_method(name));
    }

    methods
}

fn generate_updater_methods<I>(names_and_types: I) -> TokenStream
where
    I: Iterator<Item = (Ident, TunableType)> + std::clone::Clone,
{
    let mut methods = TokenStream::new();

    methods.extend(generate_updater_method(
        names_and_types.clone(),
        TunableType::Bool,
        quote::format_ident!("update_bools"),
    ));

    methods.extend(generate_updater_method(
        names_and_types.clone(),
        TunableType::I64,
        quote::format_ident!("update_ints"),
    ));

    methods.extend(generate_updater_method(
        names_and_types.clone(),
        TunableType::String,
        quote::format_ident!("update_strings"),
    ));

    methods.extend(generate_updater_method(
        names_and_types.clone(),
        TunableType::VecOfStrings,
        quote::format_ident!("update_vec_of_strings"),
    ));

    methods.extend(generate_updater_method(
        names_and_types.clone(),
        TunableType::ByRepoBool,
        quote::format_ident!("update_by_repo_bools"),
    ));

    methods.extend(generate_updater_method(
        names_and_types.clone(),
        TunableType::ByRepoString,
        quote::format_ident!("update_by_repo_strings"),
    ));

    methods.extend(generate_updater_method(
        names_and_types.clone(),
        TunableType::ByRepoI64,
        quote::format_ident!("update_by_repo_ints"),
    ));

    methods.extend(generate_updater_method(
        names_and_types,
        TunableType::ByRepoVecOfStrings,
        quote::format_ident!("update_by_repo_vec_of_strings"),
    ));

    methods
}

fn generate_updater_method<I>(
    names_and_types: I,
    ty: TunableType,
    method_name: Ident,
) -> TokenStream
where
    I: Iterator<Item = (Ident, TunableType)> + std::clone::Clone,
{
    let names = names_and_types.filter(|(_, t)| *t == ty).map(|(n, _)| n);

    let mut names = names.peekable();
    let mut body = TokenStream::new();

    if names.peek().is_some() {
        match ty {
            TunableType::I64 | TunableType::Bool => {
                body.extend(quote! {
                    #(
                        self.#names.swap(tunables.get(stringify!(#names)).map(|v| Arc::new(*v)));
                    )*
                });
            }
            TunableType::String | TunableType::VecOfStrings => {
                body.extend(quote! {
                    #(
                        self.#names.swap(tunables.get(stringify!(#names)).map(|v| Arc::new(v.clone()))
                    );)*
                });
            }
            TunableType::ByRepoBool
            | TunableType::ByRepoString
            | TunableType::ByRepoI64
            | TunableType::ByRepoVecOfStrings => {
                let by_repo_value_type = ty.by_repo_value_type();
                body.extend(quote! {
                    #(
                        let mut new_values_by_repo: HashMap<String, #by_repo_value_type> = HashMap::new();
                        for (repo, val_by_tunable) in tunables {
                                for (tunable, val) in val_by_tunable {
                                    match tunable.as_ref() {
                                        stringify!(#names) => {
                                            new_values_by_repo.insert((*repo).clone(), (*val).clone());
                                        }
                                        _ => {}
                                    }
                                }
                        }
                        self.#names.swap(Arc::new(new_values_by_repo));
                    )*
                });
            }
        }
    }

    let update_container_type = ty.update_container_type();
    quote! {
        pub fn #method_name(&self, tunables: &#update_container_type) {
            #body
        }
    }
}

fn parse_names_and_types(data: Data) -> Vec<(Ident, TunableType)> {
    match data {
        Data::Struct(data) => match data.fields {
            Fields::Named(fields) => fields
                .named
                .into_iter()
                .filter_map(|f| f.clone().ident.map(|i| (i, resolve_type(f.ty))))
                .collect::<Vec<_>>(),
            _ => unimplemented!("{}", STRUCT_FIELD_MSG),
        },
        _ => unimplemented!("{}", STRUCT_FIELD_MSG),
    }
}

fn resolve_type(ty: Type) -> TunableType {
    // TODO: Handle full paths to the types, such as
    // std::sync::atomic::AtomicBool, rather than just the type name.
    if let Type::Path(p) = ty {
        if let Some(ident) = p.path.get_ident() {
            match &ident.to_string()[..] {
                "TunableBool" => return TunableType::Bool,
                "TunableI64" => return TunableType::I64,
                "TunableVecOfStrings" => return TunableType::VecOfStrings,
                // TunableString is a type alias of ArcSwap<String>.
                // p.path.get_ident() returns None for ArcSwap<String>
                // and it makes it harder to parse it.
                // We use TunableString as a workaround
                "TunableString" => return TunableType::String,
                "TunableBoolByRepo" => return TunableType::ByRepoBool,
                "TunableI64ByRepo" => return TunableType::ByRepoI64,
                "TunableStringByRepo" => return TunableType::ByRepoString,
                "TunableVecOfStringsByRepo" => return TunableType::ByRepoVecOfStrings,
                _ => unimplemented!("{}, found: {}", UNIMPLEMENTED_MSG, &ident.to_string()[..]),
            }
        }
    }

    unimplemented!("{}", UNIMPLEMENTED_MSG);
}


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

use std::collections::HashSet;

use crate::detail::graph::EdgeType;
use crate::detail::graph::NodeType;
use crate::detail::progress::ProgressStateCountByType;
use crate::detail::progress::ProgressStateMutex;
use crate::detail::progress::ProgressSummary;
use crate::detail::state::StepStats;
use crate::detail::tail::TailParams;
use crate::detail::walk::RepoWalkParams;

pub const SCRUB: &str = "scrub";
pub const COMPRESSION_BENEFIT: &str = "compression_benefit";
pub const VALIDATE: &str = "validate";
pub const CORPUS: &str = "corpus";

// Per repo things we don't pass into the walk
#[derive(Clone)]
pub struct RepoSubcommandParams {
    pub progress_state: ProgressStateMutex<ProgressStateCountByType<StepStats, ProgressSummary>>,
    pub tail_params: TailParams,
    pub lfs_threshold: Option<u64>,
}

// These don't vary per repo
#[derive(Clone)]
pub struct JobWalkParams {
    pub enable_derive: bool,
    pub quiet: bool,
    pub error_as_data_node_types: HashSet<NodeType>,
    pub error_as_data_edge_types: HashSet<EdgeType>,
    pub repo_count: usize,
}

#[derive(Clone)]
pub struct JobParams {
    pub walk_params: JobWalkParams,
    pub per_repo: Vec<(RepoSubcommandParams, RepoWalkParams)>,
}

mononoke_app::subcommands! {
    mod compression_benefit;
    mod corpus;
    mod scrub;
    mod validate;
}


/*
 * 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(async_closure)]

mod args;
mod commands;
mod detail;
mod setup;

use std::num::NonZeroU32;

use anyhow::Error;
use blobstore_factory::BlobstoreArgDefaults;
use blobstore_factory::ReadOnlyStorage;
use clap::ArgGroup;
use clap::Parser;
use cmdlib_caching::CacheMode;
use cmdlib_caching::CachelibSettings;
use cmdlib_scrubbing::ScrubAppExtension;
use fbinit::FacebookInit;
use metaconfig_types::WalkerJobType;
use mononoke_app::args::MultiRepoArgs;
use mononoke_app::fb303::Fb303AppExtension;
use mononoke_app::fb303::ReadyFlagService;
use mononoke_app::MononokeApp;
use mononoke_app::MononokeAppBuilder;
use multiplexedblob::SrubWriteOnly;

#[derive(Parser)]
#[clap(group(
    ArgGroup::new("walkerargs")
        .required(true)
        .multiple(true)
        .args(&["repo_id", "repo_name", "sharded_service_name", "walker_type"]),
))]
struct WalkerArgs {
    /// List of Repo IDs or Repo Names used when sharded-service-name
    /// is absent.
    #[clap(flatten)]
    pub repos: MultiRepoArgs,

    /// The name of ShardManager service to be used when the walker
    /// functionality is desired to be executed in a sharded setting.
    #[clap(long, conflicts_with = "multirepos", requires = "walker_type")]
    pub sharded_service_name: Option<String>,

    /// The type of the walker job that needs to run for the current
    /// repo.
    #[clap(long, value_enum, conflicts_with = "multirepos")]
    pub walker_type: Option<WalkerJobType>,
}

#[fbinit::main]
fn main(fb: FacebookInit) -> Result<(), Error> {
    // FIXME: Investigate why some SQL queries kicked off by the walker take 30s or more.
    newfilenodes::disable_sql_timeouts();

    let service = ReadyFlagService::new();

    let cachelib_defaults = CachelibSettings {
        cache_size: 2 * 1024 * 1024 * 1024,
        ..Default::default()
    };

    let blobstore_defaults = BlobstoreArgDefaults {
        read_qps: NonZeroU32::new(20000),
        cachelib_attempt_zstd: Some(false),
        put_behaviour: Some(blobstore::PutBehaviour::OverwriteAndLog),
        ..Default::default()
    };

    let scrub_extension = ScrubAppExtension {
        write_only_missing: Some(SrubWriteOnly::SkipMissing),
        ..Default::default()
    };

    let read_only_storage = ReadOnlyStorage(true);

    let subcommands = commands::subcommands();
    let app = MononokeAppBuilder::new(fb)
        .with_app_extension(scrub_extension)
        .with_default_cachelib_settings(cachelib_defaults)
        .with_arg_defaults(CacheMode::LocalOnly)
        .with_arg_defaults(blobstore_defaults)
        .with_arg_defaults(read_only_storage)
        .with_app_extension(Fb303AppExtension {})
        .build_with_subcommands::<WalkerArgs>(subcommands)?;

    // TODO: we may want to set_ready after the repo setup is done
    service.set_ready();

    app.run_with_monitoring_and_logging(async_main, "walker", service)
}

async fn async_main(app: MononokeApp) -> Result<(), Error> {
    commands::dispatch(app).await
}